# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import annotations

import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
import trimesh

MESH_EXTS = (".stl", ".obj", ".ply", ".glb", ".gltf")
POINT_CLOUD_EXTS = (".pcd", ".xyz", ".pts", ".npy", ".npz")
GEOMETRY_EXTS = MESH_EXTS + POINT_CLOUD_EXTS
DEFAULT_NUM_VERTS = 2048


def resolve_num_verts_for_mesh(
    mesh_path: Union[str, Path],
    num_verts: int,
    num_latents: int | None,
) -> int:
    """GLB proj_cond value: file verts if < num_latents, else num_verts (not for RoPE)."""
    path = Path(mesh_path)
    if path.suffix.lower() != ".glb" or Mesh.is_point_cloud_file(str(path)):
        return int(num_verts)
    if num_latents is None:
        return int(num_verts)
    mesh = Mesh.load_mesh(str(path), normalize=False, preprocess=True)
    file_num_verts = int(mesh.verts.shape[0])
    if file_num_verts < int(num_latents):
        return file_num_verts
    return int(num_verts)


def _normalize_points(
    points: torch.Tensor,
    normalize_by: str = "bsphere",
    size: float = 2.0,
) -> torch.Tensor:
    if normalize_by == "bbox":
        bbox_min, bbox_max = points.min(0).values, points.max(0).values
        return (points - (bbox_min + bbox_max) / 2) * (size / (bbox_max - bbox_min).max())
    if normalize_by == "bsphere":
        center = (points.min(0).values + points.max(0).values) / 2
        radius = torch.linalg.norm(points - center, dim=-1).max() * 2
        return (points - center) * (size / radius)
    raise NotImplementedError(f"Invalid normalize_by: {normalize_by}")


def _resample_points(points: torch.Tensor, num_points: int) -> torch.Tensor:
    if points.shape[0] == 0:
        raise ValueError("Point cloud is empty")
    if points.shape[0] >= num_points:
        idx = torch.randperm(points.shape[0], device=points.device)[:num_points]
        return points[idx]
    idx = torch.randint(0, points.shape[0], (num_points,), device=points.device)
    return points[idx]


def _resample_point_cloud(points: torch.Tensor, num_points: int) -> torch.Tensor:
    if points.shape[0] == 0:
        raise ValueError("Point cloud is empty")
    if points.shape[0] < num_points:
        raise ValueError(
            f"Point cloud must have at least {num_points} points, got {int(points.shape[0])}"
        )
    idx = torch.randperm(points.shape[0], device=points.device)[:num_points]
    return points[idx]


def _read_point_cloud_file(filename: str) -> torch.Tensor:
    ext = os.path.splitext(filename)[1].lower()
    if ext == ".npy":
        arr = np.load(filename)
    elif ext == ".npz":
        data = np.load(filename)
        for key in ("points", "pts", "xyz", "vertices", "data"):
            if key in data:
                arr = data[key]
                break
        else:
            raise ValueError(f"No point array found in {filename}")
    elif ext in (".xyz", ".pts"):
        arr = np.loadtxt(filename, dtype=np.float32)
    elif ext == ".pcd":
        with open(filename, "r", encoding="utf-8", errors="ignore") as f:
            lines = [line.strip() for line in f if line.strip() and not line.startswith("#")]
        start = 0
        for i, line in enumerate(lines):
            header = line.upper()
            if header.startswith("POINTS") or header.startswith("WIDTH"):
                start = i + 1
                break
        arr = np.loadtxt(lines[start:], dtype=np.float32)
    else:
        raise NotImplementedError(f"Unsupported point cloud format: {ext}")

    arr = np.asarray(arr, dtype=np.float32)
    if arr.ndim == 1:
        if arr.shape[0] % 3 != 0:
            raise ValueError(f"Invalid point array shape in {filename}")
        arr = arr.reshape(-1, 3)
    if arr.shape[-1] < 3:
        raise ValueError(f"Expected at least 3 coordinates per point in {filename}")
    return torch.from_numpy(arr[..., :3].copy())


def _try_read_point_cloud_ply(filename: str) -> Optional[torch.Tensor]:
    loaded = trimesh.load(filename, process=False)
    if isinstance(loaded, trimesh.PointCloud):
        return torch.from_numpy(np.asarray(loaded.vertices, dtype=np.float32))
    if isinstance(loaded, trimesh.Trimesh) and len(loaded.faces) == 0:
        return torch.from_numpy(np.asarray(loaded.vertices, dtype=np.float32))
    if isinstance(loaded, trimesh.Scene):
        points = []
        for geometry in loaded.geometry.values():
            if isinstance(geometry, trimesh.PointCloud):
                points.append(geometry.vertices)
            elif isinstance(geometry, trimesh.Trimesh) and len(geometry.faces) == 0:
                points.append(geometry.vertices)
        if points:
            return torch.from_numpy(np.concatenate(points, axis=0).astype(np.float32))
    return None


def detect_and_triangulate_holes(
    triangles: List[Tuple[int, int, int]],
    max_cycle_size: int = 20,
) -> List[Tuple[int, int, int]]:
    """Find boundary loops and fan-triangulate holes in a triangle soup."""
    edge_to_triangles: Dict[Tuple[int, int], List[Tuple[int, int, int]]] = {}
    for tri in triangles:
        edges_in_tri = [
            tuple(sorted([tri[0], tri[1]])),
            tuple(sorted([tri[1], tri[2]])),
            tuple(sorted([tri[0], tri[2]])),
        ]
        for edge in edges_in_tri:
            edge_to_triangles.setdefault(edge, []).append(tri)

    boundary_edges = {edge for edge, tris in edge_to_triangles.items() if len(tris) == 1}
    if not boundary_edges:
        return []

    boundary_vertices: set[int] = set()
    boundary_adjacency: Dict[int, List[int]] = {}
    for v1, v2 in boundary_edges:
        boundary_vertices.update((v1, v2))
        boundary_adjacency.setdefault(v1, []).append(v2)
        boundary_adjacency.setdefault(v2, []).append(v1)

    def find_cycles_dfs(start_vertex: int) -> List[List[int]]:
        cycles: List[List[int]] = []
        visited_in_path: set[int] = set()
        path: List[int] = []

        def dfs(current: int) -> None:
            visited_in_path.add(current)
            path.append(current)
            has_unvisited_neighbor = False
            for neighbor in boundary_adjacency.get(current, []):
                if neighbor == start_vertex and len(path) >= 3:
                    if 4 <= len(path) <= max_cycle_size:
                        cycles.append(path[:])
                elif neighbor not in visited_in_path and len(path) < max_cycle_size:
                    has_unvisited_neighbor = True
                    dfs(neighbor)
            if not has_unvisited_neighbor and len(path) >= 3:
                if 4 <= len(path) + 1 <= max_cycle_size:
                    if start_vertex not in boundary_adjacency.get(current, []):
                        missing_edge = tuple(sorted([current, start_vertex]))
                        if missing_edge in edge_to_triangles:
                            cycles.append(path[:])
            path.pop()
            visited_in_path.remove(current)

        dfs(start_vertex)
        return cycles

    all_cycles: List[List[int]] = []
    processed_cycles: set[Tuple[int, ...]] = set()
    for start_v in boundary_vertices:
        for cycle in find_cycles_dfs(start_v):
            min_idx = cycle.index(min(cycle))
            normalized = tuple(cycle[min_idx:] + cycle[:min_idx])
            reversed_cycle = tuple(reversed(normalized))
            if normalized not in processed_cycles and reversed_cycle not in processed_cycles:
                processed_cycles.add(normalized)
                all_cycles.append(cycle)

    new_triangles: List[Tuple[int, int, int]] = []
    for loop in all_cycles:
        v0 = loop[0]
        for i in range(1, len(loop) - 1):
            new_triangles.append(tuple(sorted([v0, loop[i], loop[i + 1]])))
    return new_triangles


def extract_mesh_from_verts_normals_edges(
    vertices: torch.Tensor,
    vertices_normal: torch.Tensor,
    edges: torch.Tensor,
    fill_holes: bool = False,
    max_cycle_size: int = 10,
    device: torch.device | None = None,
) -> Mesh:
    """Build a Mesh from predicted vertices, normals, and edge pairs."""
    if device is None:
        device = vertices.device

    vertices = vertices.to(device, dtype=torch.float32)
    vertices_normal = vertices_normal.to(device, dtype=torch.float32)
    edges = edges.to(device, dtype=torch.long)

    num_vertices = vertices.shape[0]
    if edges.numel() == 0 or num_vertices < 3:
        raise ValueError(
            f"Cannot extract mesh: invalid input (edges={edges.numel()}, vertices={num_vertices})"
        )

    edges = torch.clamp(edges, 0, num_vertices - 1)
    edges = edges[edges[:, 0] != edges[:, 1]]
    if edges.numel() == 0:
        raise ValueError("Cannot extract mesh: no edges after removing self-loops")

    adjacency_list: List[List[int]] = [[] for _ in range(num_vertices)]
    edge_set: set[Tuple[int, int]] = set()
    for e in edges.cpu().numpy():
        v1, v2 = int(e[0]), int(e[1])
        if (v1, v2) not in edge_set and (v2, v1) not in edge_set:
            adjacency_list[v1].append(v2)
            adjacency_list[v2].append(v1)
            edge_set.add((v1, v2))

    max_neighbors = max(len(adj) for adj in adjacency_list) if adjacency_list else 0
    if max_neighbors == 0:
        raise ValueError("Cannot extract mesh: no vertex neighbors in edge graph")

    adj_tensor = torch.full((num_vertices, max_neighbors), -1, dtype=torch.long, device=device)
    for i, neighbors in enumerate(adjacency_list):
        if neighbors:
            adj_tensor[i, : len(neighbors)] = torch.tensor(neighbors, dtype=torch.long, device=device)

    triangles: List[Tuple[int, int, int]] = []
    for v1, v2 in edge_set:
        n1_mask = adj_tensor[v1] >= 0
        n2_mask = adj_tensor[v2] >= 0
        if not (n1_mask.any() and n2_mask.any()):
            continue
        neighbors_v1 = adj_tensor[v1][n1_mask]
        neighbors_v2 = adj_tensor[v2][n2_mask]
        matches = neighbors_v1.unsqueeze(1) == neighbors_v2.unsqueeze(0)
        if not matches.any():
            continue
        v1_indices, _ = torch.where(matches)
        for v3 in neighbors_v1[v1_indices].unique():
            if v3 != v1 and v3 != v2:
                triangles.append(tuple(sorted([v1, v2, int(v3.item())])))

    triangles = list(set(triangles))
    if not triangles:
        raise ValueError("Cannot extract mesh: no triangles formed from edges")

    if fill_holes:
        extra = detect_and_triangulate_holes(triangles, max_cycle_size=max_cycle_size)
        if extra:
            triangles = list(set(triangles + extra))

    triangle_tensor = torch.tensor(triangles, dtype=torch.long, device=device)
    triangle_vertices = vertices[triangle_tensor]
    edge1 = triangle_vertices[:, 1] - triangle_vertices[:, 0]
    edge2 = triangle_vertices[:, 2] - triangle_vertices[:, 0]
    face_normals = torch.linalg.cross(edge1, edge2, dim=1)
    face_normals_norm = torch.linalg.norm(face_normals, dim=1, keepdim=True)
    fallback = torch.tensor([0.0, 0.0, 1.0], device=device, dtype=torch.float32)
    face_normals = torch.where(
        face_normals_norm > 1e-8,
        face_normals / face_normals_norm,
        fallback.expand_as(face_normals),
    )

    triangle_normals = vertices_normal[triangle_tensor]
    avg_vertex_normals = triangle_normals.mean(dim=1)
    avg_normals_norm = torch.linalg.norm(avg_vertex_normals, dim=1, keepdim=True)
    avg_vertex_normals = torch.where(
        avg_normals_norm > 1e-8,
        avg_vertex_normals / avg_normals_norm,
        fallback.expand_as(avg_vertex_normals),
    )

    faces = triangle_tensor.clone()
    flip_mask = (face_normals * avg_vertex_normals).sum(dim=1) < 0
    faces[flip_mask, 1], faces[flip_mask, 2] = faces[flip_mask, 2], faces[flip_mask, 1]
    faces = torch.clamp(faces, 0, num_vertices - 1)

    mesh = Mesh(verts=vertices, faces=faces, device=device)
    mesh._v_nrm = vertices_normal
    return mesh


class Mesh:
    """Mesh class for loading, augmentation, and topology feature extraction."""

    def __init__(self, verts: torch.Tensor, faces: torch.Tensor, device: torch.device | None = None) -> None:
        """Args:
            verts: (N, 3) vertex positions.
            faces: (F, 3) triangle indices.
            device: target device; defaults to verts.device.
        """
        self.device = device if device is not None else verts.device
        self.verts = verts.to(self.device, dtype=torch.float32)
        self.faces = faces.to(self.device, dtype=torch.int64)
        self._v_nrm: Optional[torch.Tensor] = None
        self._edges: Optional[torch.Tensor] = None

    @classmethod
    def is_point_cloud_file(cls, filename: str) -> bool:
        """True for point-cloud formats, including faceless `.ply` point clouds."""
        ext = os.path.splitext(filename)[1].lower()
        if ext in POINT_CLOUD_EXTS:
            return True
        if ext == ".ply":
            return _try_read_point_cloud_ply(filename) is not None
        return False

    @classmethod
    def load_mesh(
        cls,
        filename: str,
        normalize: bool = False,
        normalize_by: str = "bbox",
        obj_scale: float = 2.0,
        preprocess: bool = True,
        device: torch.device | None = None,
    ) -> Mesh:
        """Load mesh from disk (.stl/.obj/.ply/.glb/.gltf).

        STL files are rotated -90 degrees about X during load.
        Scene graphs are flattened into a single mesh.

        Args:
            normalize: whether to call normalize() after load.
            normalize_by: "bbox" or "bsphere".
            obj_scale: target extent after normalization (passed as size).
            preprocess: whether to merge vertices and drop degenerate faces.
        """
        _, ext = os.path.splitext(filename)
        ext = ext.lower()
        if ext not in (".stl", ".obj", ".ply", ".glb", ".gltf"):
            raise NotImplementedError(f"Unsupported mesh format: {ext}")
        if cls.is_point_cloud_file(filename):
            raise ValueError(
                f"{filename} is a point cloud; use load_rope_points() or pass the path to "
                "MeshFlowPipeline.run() instead of load_mesh()."
            )

        mesh = trimesh.load(filename, force="mesh", process=False, skip_materials=True)
        if ext == ".stl":
            mesh.apply_transform(trimesh.transformations.rotation_matrix(np.radians(-90), [1, 0, 0]))
        if isinstance(mesh, trimesh.Scene):
            parts = []
            for _, node in mesh.graph.to_flattened().items():
                name = node["geometry"]
                if name in mesh.geometry and isinstance(mesh.geometry[name], trimesh.Trimesh):
                    parts.append(mesh.geometry[name].apply_transform(node["transform"]))
            mesh = trimesh.util.concatenate(parts)

        m = cls(
            torch.tensor(mesh.vertices, dtype=torch.float32),
            torch.tensor(mesh.faces, dtype=torch.int64),
            device=device,
        )
        if preprocess:
            m.preprocess()
        if normalize:
            m.normalize(normalize_by=normalize_by, size=obj_scale)
        return m

    @classmethod
    def load_rope_points(
        cls,
        filename: str,
        n_verts: int,
        normalize_by: str = "bsphere",
        obj_scale: float = 2.0,
        preprocess: bool = True,
        device: torch.device | None = None,
    ) -> torch.Tensor:
        """Load a mesh or point cloud and return fixed-size points for RoPE conditioning."""
        device = device or torch.device("cpu")
        ext = os.path.splitext(filename)[1].lower()

        if ext in POINT_CLOUD_EXTS:
            points = _read_point_cloud_file(filename).to(device)
        elif ext == ".ply":
            points = _try_read_point_cloud_ply(filename)
            if points is None:
                mesh = cls.load_mesh(
                    filename,
                    normalize=True,
                    normalize_by=normalize_by,
                    obj_scale=obj_scale,
                    preprocess=preprocess,
                    device=device,
                )
                surface_pts, _ = mesh.sample_surface(n_verts)
                return surface_pts
            points = points.to(device)
        else:
            mesh = cls.load_mesh(
                filename,
                normalize=True,
                normalize_by=normalize_by,
                obj_scale=obj_scale,
                preprocess=preprocess,
                device=device,
            )
            surface_pts, _ = mesh.sample_surface(n_verts)
            return surface_pts

        points = _normalize_points(points, normalize_by=normalize_by, size=obj_scale)
        return _resample_point_cloud(points, n_verts)

    def to_trimesh(self) -> trimesh.Trimesh:
        """Export to a host-side mesh object (detached from autograd)."""
        return trimesh.Trimesh(
            vertices=self.verts.detach().cpu().numpy(),
            faces=self.faces.detach().cpu().numpy(),
            process=False,
        )

    @property
    def v_nrm(self) -> torch.Tensor:
        """Per-vertex normals (area-weighted face normal sum), shape (N, 3)."""
        if self._v_nrm is None:
            v0, v1, v2 = self.verts[self.faces[:, 0]], self.verts[self.faces[:, 1]], self.verts[self.faces[:, 2]]
            fn = torch.linalg.cross(v1 - v0, v2 - v0)
            v_nrm = torch.zeros_like(self.verts)
            v_nrm.scatter_add_(0, self.faces[:, 0:1].expand(-1, 3), fn)
            v_nrm.scatter_add_(0, self.faces[:, 1:2].expand(-1, 3), fn)
            v_nrm.scatter_add_(0, self.faces[:, 2:3].expand(-1, 3), fn)
            v_nrm = torch.where(
                (v_nrm * v_nrm).sum(-1, keepdim=True) > 1e-20,
                v_nrm,
                torch.tensor([0.0, 0.0, 1.0], device=self.device),
            )
            self._v_nrm = F.normalize(v_nrm, dim=1)
        return self._v_nrm

    @property
    def edges(self) -> torch.Tensor:
        """Unique undirected edges, shape (E, 2)."""
        if self._edges is None:
            e = torch.cat([self.faces[:, [0, 1]], self.faces[:, [1, 2]], self.faces[:, [2, 0]]], dim=0)
            self._edges = torch.unique(torch.sort(e, dim=1)[0], dim=0)
        return self._edges

    def preprocess(self) -> Mesh:
        """Clean mesh in-place: merge vertices, remove degenerate/unreferenced geometry."""
        tm = self.to_trimesh()
        tm.process()
        tm.merge_vertices(merge_tex=True, merge_norm=True)
        tm.update_faces(tm.nondegenerate_faces())
        tm.update_faces(tm.unique_faces())
        tm.remove_unreferenced_vertices()
        self.verts = torch.from_numpy(tm.vertices).to(self.verts).float()
        self.faces = torch.from_numpy(tm.faces).to(self.faces).long()
        self._v_nrm = self._edges = None
        return self

    def normalize(self, normalize_by: str = "bsphere", size: float = 2.0) -> Mesh:
        """Center and scale vertices in-place to fit a canonical volume.

        Args:
            normalize_by: "bbox" (axis-aligned box) or "bsphere" (bounding sphere).
            size: target extent along the longest axis / diameter.
        """
        if normalize_by == "bbox":
            bbox_min, bbox_max = self.verts.min(0).values, self.verts.max(0).values
            self.verts = (self.verts - (bbox_min + bbox_max) / 2) * (size / (bbox_max - bbox_min).max())
        elif normalize_by == "bsphere":
            center = (self.verts.min(0).values + self.verts.max(0).values) / 2
            self.verts = (self.verts - center) * (
                size / (torch.linalg.norm(self.verts - center, dim=-1).max() * 2)
            )
        else:
            raise NotImplementedError(f"Invalid normalize_by: {normalize_by}")
        self._v_nrm = self._edges = None
        return self

    def rotate(self, yaw: float = 0.0, pitch: float = 0.0, roll: float = 0.0) -> Mesh:
        """Apply ZYX Euler rotation (degrees) and return a new Mesh."""
        y, p, r = np.radians(yaw), np.radians(pitch), np.radians(roll)
        cy, sy, cp, sp, cr, sr = np.cos(y), np.sin(y), np.cos(p), np.sin(p), np.cos(r), np.sin(r)
        R = torch.tensor(
            [
                [cy * cr + sy * sp * sr, -cy * sr + sy * sp * cr, sy * cp],
                [cp * sr, cp * cr, -sp],
                [-sy * cr + cy * sp * sr, sy * sr + cy * sp * cr, cy * cp],
            ],
            dtype=torch.float32,
            device=self.device,
        )
        return Mesh(self.verts @ R.T, self.faces, device=self.device)

    @torch.no_grad()
    def sample_surface(self, num_points: int, return_face_idx: bool = False):
        """Uniformly sample points on the mesh surface.

        Returns:
            points, face_normals; with return_face_idx=True also per-sample face indices.
        """
        tm = self.to_trimesh()
        points, face_idx = tm.sample(num_points, return_index=True)
        points = torch.from_numpy(points.astype(np.float32)).to(self.device)
        normals = torch.from_numpy(tm.face_normals[face_idx].astype(np.float32)).to(self.device)
        if return_face_idx:
            return points, normals, torch.from_numpy(face_idx).to(self.device, dtype=torch.long)
        return points, normals

    @torch.no_grad()
    def build_topology(
        self,
        topology_feats_type: Optional[List[str]] = None,
        n_verts: int = 8192,
        max_degree: int = 50,
    ) -> Dict[str, Any]:
        """Build padded topology tensors for MeshFlow VAE / inference.

        Supported features: verts_normal, neighbor_points, degree, adjacency_matrix.

        Returns:
            Dict with padded_verts, verts_mask, verts, faces, num_verts, and requested features.
        """
        if topology_feats_type is None:
            topology_feats_type = ["verts_normal", "neighbor_points", "degree", "adjacency_matrix"]

        m = self.preprocess().normalize(size=2.0)
        num_verts = int(m.verts.shape[0])
        if num_verts > n_verts:
            raise ValueError(f"Mesh has {num_verts} verts > n_verts={n_verts}")
        if num_verts == 0:
            raise ValueError("Mesh has no vertices")

        device = m.device
        edges, v_nrm = m.edges, m.v_nrm
        n_negative_verts = n_verts - num_verts

        negative_verts = torch.zeros(0, 3, dtype=m.verts.dtype, device=device)
        negative_normals = torch.zeros(0, 3, dtype=v_nrm.dtype, device=device)
        nearest_verts_idx = None
        if n_negative_verts > 0:
            negative_verts, negative_normals, negative_face_idx = m.sample_surface(
                n_negative_verts, return_face_idx=True
            )
            face_vertices = m.faces[negative_face_idx]
            dist2 = (face_vertices - negative_verts[:, None, :]).pow(2).sum(dim=-1)
            nearest_choice = dist2.argmin(dim=1)
            nearest_verts_idx = face_vertices[
                torch.arange(n_negative_verts, device=device),
                nearest_choice,
            ]

        padded = torch.zeros(n_verts, 3, dtype=m.verts.dtype, device=device)
        padded[:num_verts] = m.verts
        if n_negative_verts > 0:
            padded[num_verts:] = negative_verts.to(dtype=m.verts.dtype, device=device)

        verts_mask = (torch.arange(n_verts, device=device) < num_verts).float().unsqueeze(-1)
        ret: Dict[str, Any] = {
            "padded_verts": padded,
            "verts_mask": verts_mask,
            "num_verts": num_verts,
            "num_faces": int(m.faces.shape[0]),
            "num_edges": int(edges.shape[0]),
            "verts": m.verts,
            "faces": m.faces,
        }

        for feat in topology_feats_type:
            if feat == "verts_normal":
                out = torch.zeros(n_verts, 3, dtype=v_nrm.dtype, device=device)
                out[:num_verts] = v_nrm
                if n_negative_verts > 0:
                    out[num_verts:] = negative_normals.to(dtype=v_nrm.dtype, device=device)
            elif feat == "adjacency_matrix":
                out = torch.zeros(n_verts, n_verts, dtype=torch.bool, device=device)
                out[edges[:, 0], edges[:, 1]] = True
                out[edges[:, 1], edges[:, 0]] = True
                if n_negative_verts > 0:
                    neg_slice = slice(num_verts, n_verts)
                    out[neg_slice, :] = out[nearest_verts_idx, :]
                    out[:, neg_slice] = out[:, nearest_verts_idx]
            elif feat == "degree":
                degrees = torch.bincount(edges.flatten(), minlength=num_verts)
                if (degrees > max_degree).any():
                    raise ValueError(f"Vertex degree exceeds max_degree ({max_degree}).")
                out = torch.zeros(n_verts, dtype=torch.long, device=device)
                out[:num_verts] = degrees
                if n_negative_verts > 0:
                    out[num_verts:] = out[nearest_verts_idx]
                out = out.unsqueeze(-1)
            elif feat == "neighbor_points":
                all_e = torch.cat([edges, edges.flip(1)], dim=0)
                src, dst = all_e[:, 0], all_e[:, 1]
                order = torch.argsort(src)
                src, dst = src[order], dst[order]
                if torch.bincount(src, minlength=num_verts).max() > max_degree:
                    raise ValueError(f"Vertex degree exceeds max_degree ({max_degree}).")
                first = torch.searchsorted(src, src)
                out = torch.zeros(n_verts, max_degree, 3, dtype=m.verts.dtype, device=device)
                out[src, torch.arange(len(src), device=device) - first] = m.verts[dst]
                if n_negative_verts > 0:
                    out[num_verts:] = out[nearest_verts_idx]
                out = out.view(n_verts, -1)
            else:
                raise ValueError(f"Unknown topology_feats_type: {feat}")
            ret[feat] = out
        return ret