import json
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

# --- 配置 ---
import json
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

# --- 配置 ---
import json
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# --- 配置 ---
# ------------
file_path = 'data_inference/inference_reward_record_20260104_105636/episode_0001/data.jsonl'  # 你的文件名

def visualize_trajectory(filename):
    timestamps = []
    # 假设 states 有两个列表，分别对应 state 0 和 state 1
    states_0 = []
    states_1 = []
    # Actions 分别对应 left 和 right
    actions_left = []
    actions_right = []
    # Grippers
    gripper_left = []
    gripper_right = []
    LIMIT = 100
    print(f"正在读取文件: {filename} ...")
    try:
        count = 0
        with open(filename, 'r') as f:
            for line in f:
                d = json.loads(line)
                timestamps.append(d['timestamp'])
                
                # 提取 State 的前3个值 (XYZ)
                # 注意：根据您的数据结构，states是一个包含两个列表的列表
                # 这里假设 states[0] 和 states[1] 分别对应两个实体的状态
                states_0.append(d['states'][0][:3])
                states_1.append(d['states'][1][:3])
                
                # 提取 Action 的前3个值 (XYZ)
                actions_left.append(d['actions']['left_arm']['pose'][:3])
                actions_right.append(d['actions']['right_arm']['pose'][:3])
                
                # 提取 Gripper
                gripper_left.append(d['actions']['left_gripper'][0])
                gripper_right.append(d['actions']['right_gripper'][0])
                count += 1
                if count > LIMIT:
                    break
    except FileNotFoundError:
        print(f"错误: 找不到文件 {filename}。请确保文件在当前目录下或路径正确。")
        return

    # 转换为 numpy 数组方便处理
    timestamps = np.array(timestamps)
    # 将时间戳转换为相对时间（从0开始）
    rel_time = timestamps - timestamps[0]
    
    states_0 = np.array(states_0)
    states_1 = np.array(states_1)
    actions_left = np.array(actions_left)
    actions_right = np.array(actions_right)
    gripper_left = np.array(gripper_left)
    gripper_right = np.array(gripper_right)

    # 1. 绘制 XYZ 分量对比图 (State vs Action)
    fig1, axes = plt.subplots(3, 1, figsize=(12, 10), sharex=True)
    labels = ['X', 'Y', 'Z']
    
    for i in range(3):
        # 绘制 State
        axes[i].plot(rel_time, states_0[:, i], label='State[0]', linestyle='--', alpha=0.7)
        axes[i].plot(rel_time, states_1[:, i], label='State[1]', linestyle='--', alpha=0.7)
        # 绘制 Action
        axes[i].plot(rel_time, actions_left[:, i], label='Action Left', linewidth=1.5)
        axes[i].plot(rel_time, actions_right[:, i], label='Action Right', linewidth=1.5)
        
        axes[i].set_ylabel(f'{labels[i]} (m)')
        axes[i].grid(True, alpha=0.3)
        if i == 0:
            axes[i].legend(loc='upper right')
            axes[i].set_title('XYZ Position: State vs Action')

    axes[2].set_xlabel('Time (s)')
    plt.tight_layout()
    plt.savefig('1_xyz_comparison.png')
    print("已保存: 1_xyz_comparison.png")

    # 2. 绘制 3D 轨迹图
    fig2 = plt.figure(figsize=(10, 8))
    ax = fig2.add_subplot(111, projection='3d')

    # 为了图表清晰，这里只画出轨迹，不画时间点
    ax.plot(states_0[:, 0], states_0[:, 1], states_0[:, 2], label='State[0]', linestyle='--')
    ax.plot(states_1[:, 0], states_1[:, 1], states_1[:, 2], label='State[1]', linestyle='--')
    ax.plot(actions_left[:, 0], actions_left[:, 1], actions_left[:, 2], label='Action Left')
    ax.plot(actions_right[:, 0], actions_right[:, 1], actions_right[:, 2], label='Action Right')

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_title('3D Trajectories')
    ax.legend()
    plt.savefig('2_3d_trajectory.png')
    print("已保存: 2_3d_trajectory.png")

    # 3. 绘制 Gripper 状态
    plt.figure(figsize=(12, 4))
    plt.plot(rel_time, gripper_left, label='Left Gripper')
    plt.plot(rel_time, gripper_right, label='Right Gripper')
    plt.xlabel('Time (s)')
    plt.ylabel('Value')
    plt.title('Gripper State')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig('3_gripper.png')
    print("已保存: 3_gripper.png")

    # 4. 绘制 Timestamp 频率 (Delta Time)
    plt.figure(figsize=(12, 4))
    dt = np.diff(timestamps)
    # 移除异常大的跳跃点以便观察（可选）
    # dt = dt[dt < 0.5] 
    
    plt.plot(rel_time[1:], dt, marker='.', linestyle='-', linewidth=0.5, markersize=3)
    mean_dt = np.mean(dt)
    freq = 1.0 / mean_dt if mean_dt > 0 else 0
    
    plt.xlabel('Time (s)')
    plt.ylabel('Delta Time (s)')
    plt.title(f'Control Frequency Stability (Mean dt: {mean_dt:.4f}s, ~{freq:.1f} Hz)')
    plt.grid(True, alpha=0.3)
    plt.savefig('4_frequency.png')
    print("已保存: 4_frequency.png")

    plt.show()

if __name__ == "__main__":
    visualize_trajectory(file_path)