import io
import os
import tarfile
import pickle
import zstandard
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

# 数据集路径列表（支持跨包连读）
archive_paths = [
    f"/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/DyMesh_50000v_16f_0000_part_{str(i).zfill(2)}"
    for i in range(1, 4)
]

# 配置输出目标
output_dir = "mesh_vis_results"
caption_log_path = os.path.join(output_dir, "mesh_captions_map.txt")
os.makedirs(output_dir, exist_ok=True)

TARGET_OBJECT_COUNT = 100
collected_count = 0

print(f"📁 结果将保存在文件夹: {output_dir}/")
print(f"🚀 开始跨包搜寻并逐个渲染前 {TARGET_OBJECT_COUNT} 个有效的 16 帧物体...\n")

with open(caption_log_path, 'w', encoding='utf-8') as cap_file:
    cap_file.write(f"前 {TARGET_OBJECT_COUNT} 个物体的文本标注与图片映射表\n" + "="*80 + "\n")
    
    # 开始遍历压缩包
    for archive_path in archive_paths:
        if collected_count >= TARGET_OBJECT_COUNT:
            break
        print(f"🔍 正在扫描: {archive_path}")
        
        try:
            with open(archive_path, 'rb') as fh:
                dctx = zstandard.ZstdDecompressor()
                with dctx.stream_reader(fh) as reader:
                    with tarfile.open(fileobj=reader, mode='r|') as tar:
                        while True:
                            try:
                                member = tar.next()
                                if member is None:
                                    break
                            except tarfile.ReadError:
                                break  # 捕捉单包末尾截断，跳向下一包
                            
                            if member.isfile():
                                f = tar.extractfile(member)
                                if f is not None:
                                    try:
                                        mem_file = io.BytesIO(f.read())
                                        data = pickle.load(mem_file)
                                        
                                        # 验证有效性
                                        if isinstance(data, dict) and 'vertices' in data and 'faces' in data:
                                            vertices = data['vertices']
                                            faces = data['faces']
                                            
                                            if vertices.shape[0] == 16:
                                                collected_count += 1
                                                raw_name = os.path.basename(member.name)
                                                # 提取前缀哈希名，比如 048adc72...
                                                short_name = raw_name.split('.')[0] 
                                                
                                                # 定义当前物体的独立图片文件名
                                                img_name = f"Obj_{collected_count:03d}_{short_name}.png"
                                                img_path = os.path.join(output_dir, img_name)
                                                
                                                caption = data.get('caption', '无文本描述')
                                                
                                                # 1. 记录文本到日志
                                                cap_file.write(f"【物体 {collected_count:03d}】\n图片: {img_name}\n原始路径: {member.name}\n描述: {caption}\n")
                                                cap_file.write("-" * 80 + "\n")
                                                
                                                # 2. 为当前物体独立画一张 1行 16列 的长条图
                                                fig = plt.figure(figsize=(32, 2.2))
                                                
                                                # 计算该物体的全局归一化视角坐标，防止帧间缩放抖动
                                                all_v = vertices.reshape(-1, 3)
                                                max_range = (all_v.max(axis=0) - all_v.min(axis=0)).max() / 2.0
                                                mid_x, mid_y, mid_z = (all_v.max(axis=0) + all_v.min(axis=0)) / 2.0
                                                
                                                for frame_idx in range(16):
                                                    ax = fig.add_subplot(1, 16, frame_idx + 1, projection='3d')
                                                    v = vertices[frame_idx]
                                                    
                                                    # 渲染 3D 表面
                                                    ax.plot_trisurf(v[:, 0], v[:, 1], v[:, 2], triangles=faces, 
                                                                    cmap='coolwarm', edgecolor='none', alpha=0.8)
                                                    
                                                    # 锁定视角
                                                    ax.set_xlim(mid_x - max_range, mid_x + max_range)
                                                    ax.set_ylim(mid_y - max_range, mid_y + max_range)
                                                    ax.set_zlim(mid_z - max_range, mid_z + max_range)
                                                    ax.axis('off')
                                                    ax.set_title(f"F {frame_idx+1}", fontsize=9)
                                                
                                                # 给整张单图加一个总标题（可选，显示Obj编号）
                                                fig.suptitle(f"Object {collected_count:03d} - {short_name}", fontsize=12, x=0.05, y=0.5, ha='left', weight='bold', color='blue')
                                                
                                                # 保存当前物体的图片
                                                plt.tight_layout(pad=0.5)
                                                plt.savefig(img_path, dpi=100, bbox_inches='tight')
                                                plt.close()
                                                
                                                print(f"   [💾 已渲染并保存 {collected_count}/{TARGET_OBJECT_COUNT}] -> {img_name}")
                                                
                                                if collected_count >= TARGET_OBJECT_COUNT:
                                                    break
                                    except Exception:
                                        continue
        except FileNotFoundError:
            print(f" ❌ 未找到文件: {archive_path}，跳过。")
            continue

print("\n" + "="*50)
print(f"🎉 任务完美完成！")
print(f"1. 所有的 16帧 独立图片均保存在目录: {output_dir}/")
print(f"2. 文本标注与图片名称的映射对照表已写入: {caption_log_path}")
print("="*50)