import io
import os
import tarfile
import pickle
import zstandard
import numpy as np
import torch
import trimesh
from scipy.spatial import KDTree

# 引入 VAE 包装类和几何提取工具
from models.vae import TripoSGVaeWrapper
from triposg.inference_utils import hierarchical_extract_geometry

# 配置路径
archive_path = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/DyMesh_50000v_16f_0001_part_00"
triposg_weights = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/TripoSG/pretrained_weights/TripoSG"

def get_mesh_frames(archive_file, limit=100):
    """从数据集中提取前 N 个有效物体的第0帧"""
    extracted_data = []
    with open(archive_file, 'rb') as fh:
        dctx = zstandard.ZstdDecompressor()
        with dctx.stream_reader(fh) as reader:
            with tarfile.open(fileobj=reader, mode='r|') as tar:
                for member in tar:
                    if len(extracted_data) >= limit:
                        break
                    if member.isfile():
                        f = tar.extractfile(member)
                        if f is not None:
                            data = pickle.load(io.BytesIO(f.read()))
                            if isinstance(data, dict) and 'vertices' in data and 'faces' in data:
                                if data['vertices'].shape[0] >= 1:
                                    v_f0 = data['vertices'][0]
                                    faces = data['faces']
                                    short_name = os.path.basename(member.name).split('.')[0]
                                    extracted_data.append((v_f0, faces, short_name))
    return extracted_data

if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16

    print(f"📦 正在从数据包中读取前 100 个测试 Mesh...")
    mesh_list = get_mesh_frames(archive_path, limit=100)
    
    # 初始化 VAE
    print("🔮 正在初始化冻结的 TripoSG VAE 模型...")
    vae_wrapper = TripoSGVaeWrapper(weights_dir=triposg_weights, device=device, dtype=dtype)

    metrics_sum = {"mae_r2r": 0.0, "mae_raw2r": 0.0, "cd": 0.0, "fscore": 0.0}
    count = 0

    print(f"\n{'ID':<5} | {'Recon->Raw':<12} | {'Raw->Recon':<12} | {'CD':<10} | {'F-Score':<8}")
    print("-" * 65)

    for i, (raw_vertices, raw_faces, obj_name) in enumerate(mesh_list):
        # 预处理
        mesh_raw = trimesh.Trimesh(vertices=raw_vertices, faces=raw_faces, process=False)
        max_edge = np.max(mesh_raw.bounding_box.extents)
        mesh_raw.vertices = (mesh_raw.vertices - mesh_raw.center_mass) * (2.0 / max_edge)

        # 编码与解码
        v_tensor = torch.from_numpy(mesh_raw.vertices).to(device=device, dtype=dtype)
        f_tensor = torch.from_numpy(raw_faces).to(device=device, dtype=torch.int32)
        
        latent_code = vae_wrapper.encode_mesh(v_tensor, f_tensor, num_samples=204800)
        
        geometric_func = lambda x: vae_wrapper.decode_latent(latent_code, sampled_points=x)
        output_geometry = hierarchical_extract_geometry(
            geometric_func, device=device, 
            bounds=(-1.005, -1.005, -1.005, 1.005, 1.005, 1.005),
            dense_octree_depth=8, hierarchical_octree_depth=9
        )
        
        recon_v, recon_f = output_geometry[0]
        mesh_recon = trimesh.Trimesh(vertices=recon_v.astype(np.float32), faces=recon_f, process=False)

        # 评估
        pts_raw, _ = trimesh.sample.sample_surface(mesh_raw, 204800)
        pts_recon, _ = trimesh.sample.sample_surface(mesh_recon, 204800)
        
        tree_raw = KDTree(pts_raw)
        tree_recon = KDTree(pts_recon)
        
        d_r2r = np.mean(tree_raw.query(pts_recon, k=1)[0])
        d_raw2r = np.mean(tree_recon.query(pts_raw, k=1)[0])
        cd = d_r2r + d_raw2r
        
        tau = 0.01 * np.linalg.norm(mesh_raw.bounding_box.extents)
        precision = np.mean(tree_raw.query(pts_recon, k=1)[0] < tau)
        recall = np.mean(tree_recon.query(pts_raw, k=1)[0] < tau)
        f_score = (2 * precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0

        # 记录
        metrics_sum["mae_r2r"] += d_r2r
        metrics_sum["mae_raw2r"] += d_raw2r
        metrics_sum["cd"] += cd
        metrics_sum["fscore"] += f_score
        count += 1
        
        print(f"{i+1:<5} | {d_r2r:.4f}       | {d_raw2r:.4f}       | {cd:.4f}     | {f_score*100:.2f}%")

    # 输出平均值
    if count > 0:
        print("-" * 65)
        print(f"平均值: | {metrics_sum['mae_r2r']/count:.4f}       | {metrics_sum['mae_raw2r']/count:.4f}       | {metrics_sum['cd']/count:.4f}     | {(metrics_sum['fscore']/count)*100:.2f}%")