# source from AnimateAnyMesh iccv2025

import torch
import numpy as np
import time
import torch
from pytorch3d.io import save_obj

def mesh_preprocess(vertices, faces, max_length=4096):
    if vertices.ndim == 4:
        vertices, faces = vertices[0], faces[0]
    center = (vertices[0].max(dim=0)[0] + vertices[0].min(dim=0)[0]) / 2
    vertices = vertices - center
    v_max = vertices[0].abs().max()
    vertices = vertices / v_max
    valid_mask = torch.ones(vertices.shape[1], dtype=torch.bool)
    faces = torch.tensor(faces, dtype=torch.int64)
    faces_max_length = int(max_length * 2.5)
    vertices = torch.cat([vertices, torch.zeros(vertices.shape[0], max_length-vertices.shape[1], 3)], dim=1)
    vertices_color = torch.cat([vertices_color, -1.0 * torch.ones(max_length-vertices_color.shape[0], 3)], dim=0)
    faces = torch.cat([faces, -1 * torch.ones(faces_max_length-faces.shape[0], 3).to(torch.int64)], dim=0)
    valid_mask = torch.cat([valid_mask, torch.zeros(max_length-valid_mask.shape[0], dtype=torch.bool)])[None]
    valid_length = valid_mask.sum(dim=-1)
    adj_matrix = get_adjacency_matrix(vertices[0][None], faces[None], valid_length)
    return vertices, vertices_color, faces, valid_mask, valid_length, adj_matrix

def find_indices_in_merged(vertices_list, merged_vertices):
    indices_list = []
    for vertices in vertices_list:
        matches = (merged_vertices.unsqueeze(1) == vertices.unsqueeze(0))
        matches = matches.all(dim=2)  
        indices = matches.nonzero()[:, 0]
        indices = indices.reshape(vertices.shape[0])
        indices_list.append(indices)
    return indices_list

# We remove the topology merging to save computation cost
def merge_vertices_with_indices(vertices_list, faces_list):
    """
    Args:
        vertices_list: list of torch.Tensor[(Ni,3)]
        faces_list: list of torch.Tensor[(Fi,3)]
    Returns:
        merged_vertices: torch.Tensor 
        merged_faces: torch.Tensor 
        indices_list: list of torch.Tensor 
    """
    all_vertices = torch.cat(vertices_list, dim=0)
    all_faces = torch.cat(faces_list, dim=0)

    
    indices_list = []
    start_idx = 0
    inverse_indices = torch.arange(len(all_vertices))
    for vertices in vertices_list:
        indices_list.append(inverse_indices[start_idx:start_idx + len(vertices)])
        start_idx += len(vertices)
    return all_vertices, all_faces, indices_list
