import torch
import numpy as np
from kornia.geometry.conversions import (
    quaternion_to_axis_angle,
    axis_angle_to_quaternion,
)
import pymeshlab as pml
from plyfile import PlyData, PlyElement

from typing import Optional

pytorch3d_capable = True
try:
    import pytorch3d
    from pytorch3d.ops import knn_points
except ImportError:
    pytorch3d_capable = False

from .mesh import PointCloud
from shared_utils.sh_utils import SH2RGB, RGB2SH

def _base_face_areas(face_vertices_0, face_vertices_1, face_vertices_2):
    """Base function to compute the face areas."""
    x1, x2, x3 = torch.split(face_vertices_0 - face_vertices_1, 1, dim=-1)
    y1, y2, y3 = torch.split(face_vertices_1 - face_vertices_2, 1, dim=-1)

    a = (x2 * y3 - x3 * y2) ** 2
    b = (x3 * y1 - x1 * y3) ** 2
    c = (x1 * y2 - x2 * y1) ** 2
    areas = torch.sqrt(a + b + c) * 0.5

    return areas


def _base_sample_points_selected_faces(face_vertices, face_features=None):
    """Base function to sample points over selected faces, sample one point per face.
       The coordinates of the face vertices are interpolated to generate new samples.

    Args:
        face_vertices (tuple of torch.Tensor):
            Coordinates of vertices, corresponding to selected faces to sample from.
            A tuple of 3 entries corresponding to each of the face vertices.
            Each entry is a torch.Tensor of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`.
        face_features (tuple of torch.Tensor, Optional):
            Features of face vertices, corresponding to selected faces to sample from.
            A tuple of 3 entries corresponding to each of the face vertices.
            Each entry is a torch.Tensor of shape
            :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`.

    Returns:
        (torch.Tensor, torch.Tensor):
            Sampled point coordinates of shape :math:`(\\text{batch_size}, \\text{num_samples}, 3)`.
            Sampled points interpolated features of shape
            :math:`(\\text{batch_size}, \\text{num_samples}, \\text{feature_dim})`.
            If `face_vertices_features` arg is not specified, the returned interpolated features are None.
    """

    face_vertices0, face_vertices1, face_vertices2 = face_vertices

    sampling_shape = tuple(int(d) for d in face_vertices0.shape[:-1]) + (1,)
    # u is proximity to middle point between v1 and v2 against v0.
    # v is proximity to v2 against v1.
    #
    # The probability density for u should be f_U(u) = 2u.
    # However, torch.rand use a uniform (f_X(x) = x) distribution,
    # so using torch.sqrt we make a change of variable to have the desired density
    # f_Y(y) = f_X(y ^ 2) * |d(y ^ 2) / dy| = 2y
    u = torch.sqrt(torch.rand(sampling_shape,
                              device=face_vertices0.device,
                              dtype=face_vertices0.dtype))

    v = torch.rand(sampling_shape,
                   device=face_vertices0.device,
                   dtype=face_vertices0.dtype)
    w0 = 1 - u
    w1 = u * (1 - v)
    w2 = u * v

    points = w0 * face_vertices0 + w1 * face_vertices1 + w2 * face_vertices2

    features = None
    if face_features is not None:
        face_features0, face_features1, face_features2 = face_features
        features = w0 * face_features0 + w1 * face_features1 + \
            w2 * face_features2

    return points, features

# Modified from https://github.com/NVIDIAGameWorks/kaolin/blob/master/kaolin/ops/mesh/trianglemesh.py#L159
def sample_points(vertices, faces, num_samples, areas=None, face_features=None):
    r"""Uniformly sample points over the surface of triangle meshes.

    First, face on which the point is sampled is randomly selected,
    with the probability of selection being proportional to the area of the face.
    then the coordinate on the face is uniformly sampled.

    If ``face_features`` is defined for the mesh faces,
    the sampled points will be returned with interpolated features as well,
    otherwise, no feature interpolation will occur.

    Args:
        vertices (torch.Tensor):
            The vertices of the meshes, of shape
            :math:`(\text{batch_size}, \text{num_vertices}, 3)`.
        faces (torch.LongTensor):
            The faces of the mesh, of shape :math:`(\text{num_faces}, 3)`.
        num_samples (int):
            The number of point sampled per mesh.
            Also the number of faces sampled per mesh, and then sample a single point per face.
        areas (torch.Tensor, optional):
            The areas of each face, of shape :math:`(\text{batch_size}, \text{num_faces})`,
            can be preprocessed, for fast on-the-fly sampling,
            will be computed if None (default).
        face_features (torch.Tensor, optional):
            Per-vertex-per-face features, matching ``faces`` order,
            of shape :math:`(\text{batch_size}, \text{num_faces}, 3, \text{feature_dim})`.
            For example:

                1. Texture uv coordinates would be of shape
                   :math:`(\text{batch_size}, \text{num_faces}, 3, 2)`.
                2. RGB color values would be of shape
                   :math:`(\text{batch_size}, \text{num_faces}, 3, 3)`.

            When specified, it is used to interpolate the features for new sampled points.

    See also:
        :func:`~kaolin.ops.mesh.index_vertices_by_faces` for conversion of features defined per vertex
        and need to be converted to per-vertex-per-face shape of :math:`(\text{num_faces}, 3)`.

    Returns:
        (torch.Tensor, torch.LongTensor, (optional) torch.Tensor):
            the pointclouds of shape :math:`(\text{batch_size}, \text{num_samples}, 3)`,
            and the indexes of the faces selected,
            of shape :math:`(\text{batch_size}, \text{num_samples})`.

            If ``face_features`` arg is specified, then the interpolated features of sampled points of shape
            :math:`(\text{batch_size}, \text{num_samples}, \text{feature_dim})` are also returned.
    """
    if faces.shape[-1] != 3:
        raise NotImplementedError("sample_points is only implemented for triangle meshes")
    faces_0, faces_1, faces_2 = torch.split(faces, 1, dim=1)        # (num_faces, 3) -> tuple of (num_faces,)
    face_v_0 = torch.index_select(vertices, 1, faces_0.reshape(-1))  # (batch_size, num_faces, 3)
    face_v_1 = torch.index_select(vertices, 1, faces_1.reshape(-1))  # (batch_size, num_faces, 3)
    face_v_2 = torch.index_select(vertices, 1, faces_2.reshape(-1))  # (batch_size, num_faces, 3)

    if areas is None:
        areas = _base_face_areas(face_v_0, face_v_1, face_v_2).squeeze(-1)
    face_dist = torch.distributions.Categorical(areas)
    face_choices = face_dist.sample([num_samples]).transpose(0, 1)
    _face_choices = face_choices.unsqueeze(-1).repeat(1, 1, 3)
    v0 = torch.gather(face_v_0, 1, _face_choices)  # (batch_size, num_samples, 3)
    v1 = torch.gather(face_v_1, 1, _face_choices)  # (batch_size, num_samples, 3)
    v2 = torch.gather(face_v_2, 1, _face_choices)  # (batch_size, num_samples, 3)
    face_vertices_choices = (v0, v1, v2)

    # UV coordinates are available, make sure to calculate them for sampled points as well
    face_features_choices = None
    if face_features is not None:
        feat_dim = face_features.shape[-1]
        # (num_faces, 3) -> tuple of (num_faces,)
        _face_choices = face_choices[..., None, None].repeat(1, 1, 3, feat_dim)
        face_features_choices = torch.gather(face_features, 1, _face_choices)
        face_features_choices = tuple(
            tmp_feat.squeeze(2) for tmp_feat in torch.split(face_features_choices, 1, dim=2))

    points, point_features = _base_sample_points_selected_faces(
        face_vertices_choices, face_features_choices)

    if point_features is not None:
        return points, face_choices, point_features
    else:
        return points, face_choices

def poisson_mesh_reconstruction(points, normals=None):
    # points/normals: [N, 3] np.ndarray

    import open3d as o3d

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)

    # outlier removal
    pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=10)

    # normals
    if normals is None:
        pcd.estimate_normals()
    else:
        pcd.normals = o3d.utility.Vector3dVector(normals[ind])

    # visualize
    o3d.visualization.draw_geometries([pcd], point_show_normal=False)

    mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
        pcd, depth=9
    )
    vertices_to_remove = densities < np.quantile(densities, 0.1)
    mesh.remove_vertices_by_mask(vertices_to_remove)

    # visualize
    o3d.visualization.draw_geometries([mesh])

    vertices = np.asarray(mesh.vertices)
    triangles = np.asarray(mesh.triangles)

    print(
        f"[INFO] poisson mesh reconstruction: {points.shape} --> {vertices.shape} / {triangles.shape}"
    )

    return vertices, triangles


def decimate_mesh(
    verts, faces, target=5e4, backend="pymeshlab", remesh=False, optimalplacement=True, verbose=True
):
    """ perform mesh decimation.

    Args:
        verts (np.ndarray): mesh vertices, float [N, 3]
        faces (np.ndarray): mesh faces, int [M, 3]
        target (int): targeted number of faces
        backend (str, optional): algorithm backend, can be "pymeshlab" or "pyfqmr". Defaults to "pymeshlab".
        remesh (bool, optional): whether to remesh after decimation. Defaults to False.
        optimalplacement (bool, optional): For flat mesh, use False to prevent spikes. Defaults to True.
        verbose (bool, optional): whether to print the decimation process. Defaults to True.

    Returns:
        Tuple[np.ndarray]: vertices and faces after decimation.
    """

    _ori_vert_shape = verts.shape
    _ori_face_shape = faces.shape

    if backend == "pyfqmr":
        import pyfqmr

        solver = pyfqmr.Simplify()
        solver.setMesh(verts, faces)
        solver.simplify_mesh(target_count=target, preserve_border=False, verbose=False)
        verts, faces, normals = solver.getMesh()
    else:
        m = pml.Mesh(verts, faces)
        ms = pml.MeshSet()
        ms.add_mesh(m, "mesh")  # will copy!
        
        ms.meshing_merge_close_vertices()

        # filters
        # ms.meshing_decimation_clustering(threshold=pml.PercentageValue(1))
        ms.meshing_decimation_quadric_edge_collapse(
            targetfacenum=int(target), optimalplacement=optimalplacement
        )

        if remesh:
            # ms.apply_coord_taubin_smoothing()
            ms.meshing_isotropic_explicit_remeshing(
                iterations=3, targetlen=pml.PercentageValue(1)
            )

        # extract mesh
        m = ms.current_mesh()
        m.compact()
        verts = m.vertex_matrix()
        faces = m.face_matrix()

    if verbose:
        print(f"[INFO] mesh decimation: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}")

    return verts, faces


def clean_mesh(
    verts,
    faces,
    v_pct=1,
    min_f=64,
    min_d=20,
    repair=True,
    remesh=True,
    remesh_size=0.01,
):
    # verts: [N, 3]
    # faces: [N, 3]

    _ori_vert_shape = verts.shape
    _ori_face_shape = faces.shape

    m = pml.Mesh(verts, faces)
    ms = pml.MeshSet()
    ms.add_mesh(m, "mesh")  # will copy!

    # filters
    ms.meshing_remove_unreferenced_vertices()  # verts not refed by any faces

    if v_pct > 0:
        ms.meshing_merge_close_vertices(
            threshold=pml.Percentage(v_pct)
        )  # 1/10000 of bounding box diagonal

    ms.meshing_remove_duplicate_faces()  # faces defined by the same verts
    ms.meshing_remove_null_faces()  # faces with area == 0

    if min_d > 0:
        ms.meshing_remove_connected_component_by_diameter(
            mincomponentdiag=pml.Percentage(min_d)
        )

    if min_f > 0:
        ms.meshing_remove_connected_component_by_face_number(mincomponentsize=min_f)

    if repair:
        # ms.meshing_remove_t_vertices(method=0, threshold=40, repeat=True)
        ms.meshing_repair_non_manifold_edges(method=0)
        ms.meshing_repair_non_manifold_vertices(vertdispratio=0)

    if remesh:
        # ms.apply_coord_taubin_smoothing()
        ms.meshing_isotropic_explicit_remeshing(
            iterations=3, targetlen=pml.AbsoluteValue(remesh_size)
        )

    # extract mesh
    m = ms.current_mesh()
    verts = m.vertex_matrix()
    faces = m.face_matrix()

    print(
        f"[INFO] mesh cleaning: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}"
    )

    return verts, faces

def construct_list_of_gs_attributes(features_dc, features_rest, scaling, rotation):
    l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
    for i in range(features_dc.shape[1]*features_dc.shape[2]):
        l.append('f_dc_{}'.format(i))
    for i in range(features_rest.shape[1]*features_rest.shape[2]):
        l.append('f_rest_{}'.format(i))
    l.append('opacity')
    for i in range(scaling.shape[1]):
        l.append('scale_{}'.format(i))
    for i in range(rotation.shape[1]):
        l.append('rot_{}'.format(i))
    return l

def calculate_max_sh_degree_from_gs_ply(plydata):
    extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
    #assert len(extra_f_names)!=3*(max_sh_degree + 1) ** 2 - 3:
    max_sh_degree = int(((len(extra_f_names) + 3) / 3) ** 0.5 - 1)
    return max_sh_degree, extra_f_names

def write_gs_ply(xyz, normals, f_dc, f_rest, opacities, scale, rotation, list_of_attributes):
    dtype_full = [(attribute, 'f4') for attribute in list_of_attributes]

    elements = np.empty(xyz.shape[0], dtype=dtype_full)
    attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
    elements[:] = list(map(tuple, attributes))
    el = PlyElement.describe(elements, 'vertex')
    return PlyData([el])

def read_gs_ply(plydata):
    xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
                    np.asarray(plydata.elements[0]["y"]),
                    np.asarray(plydata.elements[0]["z"])),  axis=1)
    opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]

    features_dc = np.zeros((xyz.shape[0], 3, 1))
    features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
    features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
    features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])

    max_sh_degree, extra_f_names = calculate_max_sh_degree_from_gs_ply(plydata)

    features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
    for idx, attr_name in enumerate(extra_f_names):
        features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
    # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
    features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1))

    scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
    scales = np.zeros((xyz.shape[0], len(scale_names)))
    for idx, attr_name in enumerate(scale_names):
        scales[:, idx] = np.asarray(plydata.elements[0][attr_name])

    rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
    rots = np.zeros((xyz.shape[0], len(rot_names)))
    for idx, attr_name in enumerate(rot_names):
        rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
        
    return xyz, features_dc, features_extra, opacities, scales, rots

def ply_to_points_cloud(plydata):
    xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata)
    
    features_dc = np.transpose(features_dc, (0, 2, 1))  # equivalent of torch.transpose(features_dc, 1, 2)
    features_extra = np.transpose(features_extra, (0, 2, 1))
    shs = np.concatenate((features_dc, features_extra), axis=1)
    normals = np.zeros_like(xyz)
    pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals)
    return pcd


def get_target_axis_and_scale(axis_string, scale_value=1.0):
    """
    Coordinate system inverts when:
    1. Any of the axis inverts
    2. Two of the axises switch
    
    If coordinate system inverts twice in a row then it will not be inverted
    """
    axis_names = ["x", "y", "z"]
    
    target_axis, target_scale, coordinate_invert_count = [], [], 0
    axis_switch_count = 0
    for i in range(len(axis_names)):
        s = axis_string[i]
        if s[0] == "-":
            target_scale.append(-scale_value)
            coordinate_invert_count += 1
        else:
            target_scale.append(scale_value)
            
        new_axis_i = axis_names.index(s[1])
        if new_axis_i != i:
            axis_switch_count += 1
        target_axis.append(new_axis_i)
        
    if axis_switch_count == 2:
        coordinate_invert_count += 1
        
    return target_axis, target_scale, coordinate_invert_count

def switch_vector_axis(vector3s, target_axis):
    """
    Example:
        vector3s = torch.tensor([[1, 2, 3], [3, 2, 1], [2, 3, 1]])  # shape (N, 3)

        target_axis = (2, 0, 1) # or [2, 0, 1]
        vector3s[:, [0, 1, 2]] = vector3s[:, target_axis]
        
        # Result: tensor([[3, 1, 2], [1, 3, 2], [1, 2, 3]])
    """
    vector3s[:, [0, 1, 2]] = vector3s[:, target_axis]
    return vector3s

def switch_ply_axis_and_scale(plydata, target_axis, target_scale, coordinate_invert_count):
    """
    Args:
        target_axis (array): shape (3)
        target_scale (array): shape (3)
    """
    xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata)
    normals = np.zeros_like(xyz)
    features_dc_2d = features_dc.reshape(features_dc.shape[0], features_dc.shape[1]*features_dc.shape[2])
    features_extra_2d = features_extra.reshape(features_extra.shape[0], features_extra.shape[1]*features_extra.shape[2])
    
    target_scale = torch.tensor(target_scale).float().cuda()
    xyz = switch_vector_axis(torch.tensor(xyz).float().cuda() * target_scale, target_axis).detach().cpu().numpy()
    scales = switch_vector_axis(torch.tensor(scales).float().cuda(), target_axis).detach().cpu().numpy()
    
    # change rotation representation from quaternion (w, x, y, z) to axis angle vector (x, y, z) to make swich axis easier
    rots_axis_angle = quaternion_to_axis_angle(torch.tensor(rots).float().cuda())
    rots_axis_angle = switch_vector_axis(rots_axis_angle * target_scale, target_axis)
    """
    Since axis–angle vector is composed of axis (unit vector/direction) and clockwise radians angle (vector magnitude),
    so in order to invert the sign of angle when coordinate system inverts, we also need to invert the direction of axis–angle vector
    """
    if coordinate_invert_count % 2 != 0:
        rots_axis_angle = -rots_axis_angle
    rots = axis_angle_to_quaternion(rots_axis_angle).detach().cpu().numpy()
    
    return write_gs_ply(xyz, normals, features_dc_2d, features_extra_2d, opacities, scales, rots, construct_list_of_gs_attributes(features_dc, features_extra, scales, rots))
    
def switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal=False):
    """
    Args:
        target_axis (array): shape (3)
        target_scale (array): shape (3)
    """
    target_scale = torch.tensor(target_scale).float().cuda()
    mesh.v = switch_vector_axis(mesh.v * target_scale, target_axis)
    mesh.vn = switch_vector_axis(mesh.vn * target_scale, target_axis)
    if flip_normal:
        mesh.vn *= -1
    return mesh


def marching_cubes_density_to_mesh(get_density_func, grid_size=256, S=128, density_thresh=10, decimate_target=5e4):
    from mcubes import marching_cubes
    from kiui.mesh_utils import clean_mesh, decimate_mesh
    
    
    sigmas = np.zeros([grid_size, grid_size, grid_size], dtype=np.float32)

    X = torch.linspace(-1, 1, grid_size).split(S)
    Y = torch.linspace(-1, 1, grid_size).split(S)
    Z = torch.linspace(-1, 1, grid_size).split(S)

    for xi, xs in enumerate(X):
        for yi, ys in enumerate(Y):
            for zi, zs in enumerate(Z):
                xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing='ij')
                pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [S, 3]
                val = get_density_func(pts)
                sigmas[xi * S: xi * S + len(xs), yi * S: yi * S + len(ys), zi * S: zi * S + len(zs)] = val.reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [S, 1] --> [x, y, z]

    print(f'[INFO] marching cubes thresh: {density_thresh} ({sigmas.min()} ~ {sigmas.max()})')

    vertices, triangles = marching_cubes(sigmas, density_thresh)
    vertices = vertices / (grid_size - 1.0) * 2 - 1
    
    # clean
    vertices = vertices.astype(np.float32)
    triangles = triangles.astype(np.int32)
    vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01)
    if triangles.shape[0] > decimate_target:
        vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False)
    
    return vertices, triangles

def color_func_to_albedo(mesh, get_rgb_func, texture_resolution=1024, padding=2, batch_size=640000, device="cuda", force_cuda_rast=False, glctx=None):
    import nvdiffrast.torch as dr
    from kiui.op import uv_padding
    
    if glctx is None:
        if force_cuda_rast:
            glctx = dr.RasterizeCudaContext()
        else:
            glctx = dr.RasterizeGLContext()
    
    # render uv maps
    h = w = texture_resolution
    if mesh.vt is None:
        mesh.auto_uv()

    uv = mesh.vt * 2.0 - 1.0 # uvs to range [-1, 1]
    uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4]

    rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), mesh.ft, (h, w)) # [1, h, w, 4]
    xyzs, _ = dr.interpolate(mesh.v.unsqueeze(0), rast, mesh.f) # [1, h, w, 3]
    mask, _ = dr.interpolate(torch.ones_like(mesh.v[:, :1]).unsqueeze(0), rast, mesh.f) # [1, h, w, 1]

    # masked query 
    xyzs = xyzs.view(-1, 3)
    mask = (mask > 0).view(-1)
    
    albedo = torch.zeros(h * w, 3, device=device, dtype=torch.float32)

    if mask.any():
        print(f"[INFO] querying texture...")

        xyzs = xyzs[mask] # [M, 3]

        # batched inference to avoid OOM
        batch = []
        head = 0
        while head < xyzs.shape[0]:
            tail = min(head + batch_size, xyzs.shape[0])
            batch.append(get_rgb_func(xyzs[head:tail]).float())
            head += batch_size

        albedo[mask] = torch.cat(batch, dim=0)
    
    albedo = albedo.view(h, w, -1)
    mask = mask.view(h, w)
    albedo = uv_padding(albedo, mask, padding)
    
    return albedo

@torch.no_grad()
def K_nearest_neighbors_func(
    points: torch.Tensor,
    K: int,
    query: Optional[torch.Tensor] = None,
    return_dist=False,
    skip_first_index=True,
):
    if not pytorch3d_capable:
        raise ImportError("pytorch3d is not installed, which is required for KNN")
    
    # query/points: Tensor of shape (N, P1/P2, D) giving a batch of N point clouds, each containing up to P1/P2 points of dimension D
    if query is None:
        query = points
    dist, idx, nn = knn_points(query[None, ...], points[None, ...], K=K, return_nn=True)

    # idx: Tensor of shape (N, P1, K)
    # nn: Tensor of shape (N, P1, K, D)

    # take the index 1 since index 0 is the point itself
    if skip_first_index:
        if not return_dist:
            return nn[0, :, 1:, :], idx[0, :, 1:]
        else:
            return nn[0, :, 1:, :], idx[0, :, 1:], dist[0, :, 1:]
    else:
        if not return_dist:
            return nn[0, :, :, :], idx[0, :, :]
        else:
            return nn[0, :, :, :], idx[0, :, :], dist[0, :, :]


def reduce_mesh_faces_inplace(mesh_obj, max_faces: int = 40000, backend="pymeshlab", optimalplacement=True, verbose=True):
    """
    Reduce faces for Mesh/FastMesh objects using decimate_mesh (in-place modification)
    
    Args:
        mesh_obj: Mesh or FastMesh object (modified in-place)
        max_faces: maximum number of faces
        backend: algorithm backend ("pymeshlab" or "pyfqmr")
        optimalplacement: optimal placement for decimation
        verbose: print progress
        
    Returns:
        Same mesh object (modified in-place)
    """
    
    faces_before = mesh_obj.f.shape[0]
    
    # Skip if already under limit
    if faces_before <= max_faces:
        if verbose:
            print(f"Mesh already has {faces_before} faces (≤ {max_faces}) - skipping")
        return mesh_obj
    
    try:
        
        # Convert to numpy
        vertices_np = mesh_obj.v.detach().cpu().numpy().astype(np.float32)
        faces_np = mesh_obj.f.detach().cpu().numpy().astype(np.int32)
        
        # Apply decimation
        reduced_vertices, reduced_faces = decimate_mesh(
            vertices_np, faces_np, 
            target=max_faces,
            backend=backend,
            optimalplacement=optimalplacement,
            verbose=verbose
        )
        
        # Update existing mesh object (in-place)
        device = mesh_obj.v.device
        mesh_obj.v = torch.from_numpy(reduced_vertices).contiguous().float().to(device)
        mesh_obj.f = torch.from_numpy(reduced_faces).contiguous().int().to(device)
        
        # Recalculate normals
        mesh_obj.auto_normal()
        
        # Stats
        if verbose:
            faces_after = mesh_obj.f.shape[0]
            reduction_ratio = ((faces_before - faces_after) / faces_before) * 100
        
        return mesh_obj
        
    except Exception as e:
        if verbose:
            print(f"Face reduction error: {e}, keep original mesh")
        return mesh_obj


def interpolate_texture_map_attr(mesh, texture_size: int = 256, batch_size: int = 64, interpolate_color=True, interpolate_position=False):
    # Get UV coordinates and faces
    if mesh.vt is None:
        mesh.auto_uv()
        
    # Get Faces on UV
    texture_size_minus_one = texture_size - 1
    faces = mesh.ft
    verts_uvs = mesh.vt * texture_size_minus_one
    verts_uvs_idx = verts_uvs.to(torch.long)

    vmapping = mesh.get_default_vt_to_v_mapping()
    # Get vertex colors
    interpolate_color = interpolate_color and mesh.vc is not None
    texture_map = mesh.albedo
    if interpolate_color:
        verts_colors = mesh.vc[vmapping]
        # Create a blank texture map
        texture_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device)

    # Get vertex positions
    position_map = None
    if interpolate_position:
        verts_positions = mesh.v[vmapping]
        position_map = torch.zeros((texture_size, texture_size, 3), device=verts_colors.device)
    
    # Create a grid of UV coordinates
    grid_x, grid_y = torch.meshgrid(torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device), torch.linspace(0, texture_size_minus_one, texture_size, dtype=verts_uvs.dtype, device=verts_colors.device))
    grid = torch.stack([grid_x, grid_y], dim=-1)
    
    batch_idx_range = torch.arange(0, batch_size**2, 1, device=verts_colors.device, dtype=torch.long)
    
    all_face_verts = verts_uvs_idx[faces]
    for uv_i in range(0, texture_size_minus_one, batch_size):
        for uv_j in range(0, texture_size_minus_one, batch_size):
            # Divide faces into squre batchs
            faces_mask = (uv_i <= all_face_verts[:, :, 0]) & (all_face_verts[:, :, 0] <= uv_i + batch_size) & (uv_j <= all_face_verts[:, :, 1]) & (all_face_verts[:, :, 1] <= uv_j + batch_size)
            faces_mask, _ = torch.max(faces_mask, dim=1)
            faces_mask = faces_mask.nonzero().squeeze(-1)
            batch_faces = faces[faces_mask]
            
            uvs = verts_uvs[batch_faces]
            if interpolate_color:
                colors = verts_colors[batch_faces]
            if interpolate_position:
                positions = verts_positions[batch_faces]
                        
            v0, v1, v2 = uvs[:, 0], uvs[:, 1], uvs[:, 2]
            v0_0, v0_1 = v0[:, 0], v0[:, 1]
            v1_0, v1_1 = v1[:, 0], v1[:, 1]
            v2_0, v2_1 = v2[:, 0], v2[:, 1]
            
            sub_grid = grid[uv_i:uv_i + batch_size, uv_j:uv_j + batch_size].reshape(-1, 2).unsqueeze(1)
                        
            # Compute barycentric coordinates for the batch of faces
            denom = (v1_1 - v2_1) * (v0_0 - v2_0) + (v2_0 - v1_0) * (v0_1 - v2_1)
            a = ((v1_1 - v2_1) * (sub_grid[:, :, 0] - v2_0) + (v2_0 - v1_0) * (sub_grid[:, :, 1] - v2_1)) / denom
            b = ((v2_1 - v0_1) * (sub_grid[:, :, 0] - v2_0) + (v0_0 - v2_0) * (sub_grid[:, :, 1] - v2_1)) / denom
            c = 1 - a - b
            
            # Mask for points inside the triangle
            mask = (a >= 0) & (b >= 0) & (c >= 0)
            mask, mask_idx = torch.max(mask, dim=1)
            
            # Get texture coordinates
            sub_grid = sub_grid.squeeze(1)
            uv_coords = sub_grid[mask].to(torch.long)
            
            # Interpolate colors
            a, b, c = a[batch_idx_range, mask_idx][mask], b[batch_idx_range, mask_idx][mask], c[batch_idx_range, mask_idx][mask]
            
            if interpolate_color:
                colors = colors[mask_idx][mask]
                interpolated_colors = (a[:, None] * colors[:, 0] + (b[:, None] * colors[:, 1]) + (c[:, None] * colors[:, 2]))
                texture_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_colors
            if interpolate_position:
                positions = positions[mask_idx][mask]
                interpolated_positions = (a[:, None] * positions[:, 0] + (b[:, None] * positions[:, 1]) + (c[:, None] * positions[:, 2]))
                position_map[uv_coords[:, 1], uv_coords[:, 0]] = interpolated_positions
            
    return texture_map, position_map