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_path = "/home/dataset-assist-0/usr/lh/ysh/dw/RL/AR/data/DyMesh_50000v_16f_0000_part_00"
output_debug_image = "mesh_motion_diff_check.png"

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:
            for member in tar:
                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:
                            vertices = data['vertices']
                            faces = data['faces']
                            
                            if vertices.shape[0] == 16:
                                # ======= 开始核心差值计算 =======
                                fig = plt.figure(figsize=(32, 3))
                                base_frame = vertices[0] # 以第一帧为基准
                                
                                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
                                
                                print(f"正在放大渲染物体 [{member.name}] 的每帧运动轨迹...")
                                
                                for frame_idx in range(16):
                                    ax = fig.add_subplot(1, 16, frame_idx + 1, projection='3d')
                                    v = vertices[frame_idx]
                                    
                                    if frame_idx == 0:
                                        # 第一帧作为基准，显示为冷色调
                                        ax.plot_trisurf(v[:, 0], v[:, 1], v[:, 2], triangles=faces, 
                                                        color='cyan', edgecolor='none', alpha=0.6)
                                        ax.set_title("Base Frame 1", fontsize=10, color='blue')
                                    else:
                                        # 计算当前帧每个顶点相对于第一帧的欧氏距离（位移量）
                                        displacements = np.linalg.norm(v - base_frame, axis=1)
                                        max_disp = displacements.max()
                                        
                                        # 如果有位移，把位移映射为色彩（动的越多越红，没动的部位是蓝色）
                                        # 这样即使只有 0.001 的微小位移，也会在视觉上变红！
                                        if max_disp > 0:
                                            colors = plt.cm.jet(displacements / max_disp)
                                        else:
                                            colors = 'blue'
                                            
                                        # 渲染带有“运动热力图”的 Mesh
                                        surf = ax.plot_trisurf(v[:, 0], v[:, 1], v[:, 2], triangles=faces, 
                                                                edgecolor='none', alpha=0.8)
                                        surf.set_facecolors(colors)
                                        ax.set_title(f"F{frame_idx+1} (Max:{max_disp:.4f})", fontsize=9)
                                    
                                    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')
                                
                                plt.tight_layout()
                                plt.savefig(output_debug_image, dpi=150, bbox_inches='tight')
                                plt.close()
                                print(f"📊 运动热力图已生成: {output_debug_image}，快去看看哪里变红了！")
                                break