import os
import sys

triposg_root = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/TripoSG"
if triposg_root not in sys.path:
    sys.path.append(triposg_root)

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

# 现在可以完美导入了
from triposg.models.autoencoders import TripoSGVAEModel


class TripoSGVaeWrapper(nn.Module):
    def __init__(self, weights_dir: str = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/TripoSG/pretrained_weights/TripoSG", device: str = "cuda", dtype: torch.dtype = torch.float16):
        """
        加载预训练且冻结的 TripoSG VAE 模型，集成内置的 Mesh 采样与归一化预处理。
        """
        super().__init__()
        self.device = device
        self.dtype = dtype
        
        # 1. 加载预训练 VAE
        assert os.path.exists(os.path.join(weights_dir, "vae")), f"Error: 找不到 VAE 子文件夹，请确认路径: {weights_dir}"
        self.vae = TripoSGVAEModel.from_pretrained(
            weights_dir,
            subfolder="vae",
        ).to(self.device, dtype=self.dtype)
        
        # 2. 彻底冻结所有 VAE 参数
        self.vae.eval()
        for param in self.vae.parameters():
            param.requires_grad = False

    def preprocess_mesh(self, vertices: torch.Tensor, faces: torch.Tensor, num_samples: int = 204800) -> torch.Tensor:
        verts_np = vertices.detach().cpu().numpy()
        faces_np = faces.detach().cpu().numpy()
        
        mesh = trimesh.Trimesh(vertices=verts_np, faces=faces_np, process=True)
        if mesh.is_empty: return None
        
        components = mesh.split(only_watertight=False)
        mesh = max(components, key=lambda m: len(m.vertices))
        
        distances = np.linalg.norm(mesh.vertices - mesh.center_mass, axis=1)
        if len(distances) > 0:
            mean_dist = np.mean(distances)
            std_dist = np.std(distances)
            mask = (distances < (mean_dist + 3.0 * (std_dist + 1e-6))).astype(bool)
            
            if not np.all(mask):
                # 1. 过滤顶点
                new_vertices = mesh.vertices[mask]
                
                # 2. 核心步骤：构建索引映射表 (Old Index -> New Index)
                # 创建一个数组，存储原索引到新索引的映射
                map_old_to_new = np.full(len(mesh.vertices), -1, dtype=int)
                map_old_to_new[mask] = np.arange(len(new_vertices))
                
                # 3. 过滤面片：只保留所有顶点都在 mask 里的面片
                valid_faces_mask = mask[mesh.faces].all(axis=1)
                new_faces = mesh.faces[valid_faces_mask]
                
                # 4. 更新面片的顶点索引：将旧编号替换为新编号
                new_faces = map_old_to_new[new_faces]
                
                # 现在重建的 mesh 是合法且连续的
                mesh = trimesh.Trimesh(vertices=new_vertices, faces=new_faces, process=True)
            
        if mesh.is_empty: return None
        
        # 后续归一化和采样逻辑保持不变...
        q_min, q_max = np.percentile(mesh.vertices, [2.5, 97.5], axis=0)
        max_extent = np.max(q_max - q_min)
        scale = 2.0 / (max_extent + 1e-6)
        normalized_vertices = (mesh.vertices - mesh.center_mass) * scale
        
        norm_mesh = trimesh.Trimesh(vertices=normalized_vertices, faces=mesh.faces, process=False)
        
        surface_points, face_indices = trimesh.sample.sample_surface(norm_mesh, count=num_samples)
        surface_normals = norm_mesh.face_normals[face_indices]
        
        surface_tensor = torch.cat([torch.from_numpy(surface_points), torch.from_numpy(surface_normals)], dim=-1).float()
        surface_tensor = torch.clamp(surface_tensor, -10.0, 10.0)
        
        return surface_tensor.unsqueeze(0).to(device=self.device, dtype=torch.float16)
        

    @torch.no_grad()
    def encode_mesh(self, vertices: np.ndarray, faces: np.ndarray, num_samples: int = 204800) -> torch.Tensor:
        # 1. 预处理数据
        surface_tensor = self.preprocess_mesh(vertices, faces, num_samples=num_samples)
        
        # --- 增加深度快照 ---
        if torch.isnan(surface_tensor).any():
            print("[CRITICAL] Input to VAE contains NaN!")
        
        # 检查输入数值范围 (这是最关键的一步)
        # 如果 min/max 超过了 (-10, 10)，VAE 极易溢出
        if surface_tensor.min() < -100 or surface_tensor.max() > 100:
            print(f"[CRITICAL] Input range extreme: min={surface_tensor.min()}, max={surface_tensor.max()}")
        
        # 2. VAE 编码
        encoder_output = self.vae.encode(surface_tensor)
        latent_dist = encoder_output.latent_dist
        
        # 检查分布参数
        if torch.isnan(latent_dist.mean).any() or torch.isnan(latent_dist.var).any():
            print("[CRITICAL] VAE produced NaN in latent_dist params!")
            
        latent_sample = latent_dist.sample()
        
        # 4. 编码后检查
        if torch.isnan(latent_sample).any():
            print("[Debug] VAE output NaN!")
            # 此时 surface_tensor 已经包含了原始点云信息
            # 建议保存 debug 数据以便后续排查
            torch.save(surface_tensor, "nan_surface_tensor.pt")
            raise ValueError("VAE produced NaN")
            
        return latent_sample

    @torch.no_grad()
    def decode_latent(self, latent_sample: torch.Tensor, sampled_points: torch.Tensor) -> torch.Tensor:
        """
        传入隐特征向量以及需要查询的 3D 空间坐标点，解码返回这些坐标点上的 SDF/Occupancy 预测值
        """
        return self.vae.decode(latent_sample, sampled_points=sampled_points).sample