#!/usr/bin/env python3
"""
Test rotation computation on actual Unity export data.
This will show what's actually causing the 168° roll error.
"""
import json
import numpy as np
from scipy.spatial.transform import Rotation as R
from pathlib import Path
from glob import glob

def test_actual_frame():
    """Load first frame from Unity export and compare rotation computations."""

    # Find Unity export
    jsonl_files = sorted(glob("./raw_dataset/sequence_*.jsonl"))
    if not jsonl_files:
        print("ERROR: No Unity export found in ./raw_dataset/")
        print("Please run this from GENMO/ directory")
        return

    with open(jsonl_files[0], "r") as f:
        lines = f.readlines()

    if len(lines) < 2:
        print("ERROR: No data in JSONL file")
        return

    # Parse frame 0 (skip header line)
    row = json.loads(lines[1])

    print("="*70)
    print("TESTING ACTUAL UNITY EXPORT DATA (Frame 0)")
    print("="*70)
    print()

    # Extract Unity quaternions
    cam_quat = np.array(row["cam_rot_world"], dtype=np.float64)
    pel_quat = np.array(row["pelvis_rot_world"], dtype=np.float64)

    print(f"Unity quaternions (XYZW):")
    print(f"  cam_rot_world:    {cam_quat}")
    print(f"  pelvis_rot_world: {pel_quat}")
    print()

    # Convert to matrices
    R_cam_w_unity = R.from_quat(cam_quat).as_matrix()
    R_pel_w_unity = R.from_quat(pel_quat).as_matrix()

    print("Unity rotations (Euler YXZ):")
    print(f"  Camera: {R.from_matrix(R_cam_w_unity).as_euler('YXZ', degrees=True)}")
    print(f"  Pelvis: {R.from_matrix(R_pel_w_unity).as_euler('YXZ', degrees=True)}")
    print()

    # Coordinate conversion
    C = np.diag([1.0, -1.0, 1.0])

    # === OLD METHOD ===
    print("--- OLD METHOD (Compute relative, then convert) ---")
    R_w2c_unity = R_cam_w_unity.T
    R_rel_unity = R_w2c_unity @ R_pel_w_unity
    R_cv_old = C @ R_rel_unity @ C

    euler_old = R.from_matrix(R_cv_old).as_euler('YXZ', degrees=True)
    print(f"global_orient_c (YXZ): yaw={euler_old[0]:7.2f}°, pitch={euler_old[1]:7.2f}°, roll={euler_old[2]:7.2f}°")
    print()

    # === NEW METHOD ===
    print("--- NEW METHOD (Convert, then compute relative) ---")
    R_cam_w_cv = C @ R_cam_w_unity @ C
    R_pel_w_cv = C @ R_pel_w_unity @ C
    R_w2c_cv = R_cam_w_cv.T
    R_pel_c_cv = R_w2c_cv @ R_pel_w_cv

    euler_new = R.from_matrix(R_pel_c_cv).as_euler('YXZ', degrees=True)
    print(f"global_orient_c (YXZ): yaw={euler_new[0]:7.2f}°, pitch={euler_new[1]:7.2f}°, roll={euler_new[2]:7.2f}°")
    print()

    # === COMPARE WITH UNITY'S smpl_incam_rot ===
    if "smpl_incam_rot" in row:
        unity_incam_quat = np.array(row["smpl_incam_rot"], dtype=np.float64)
        R_unity_incam = R.from_quat(unity_incam_quat).as_matrix()
        # Unity gives this in Unity convention, need to convert to CV
        R_unity_incam_cv = C @ R_unity_incam @ C
        euler_unity = R.from_matrix(R_unity_incam_cv).as_euler('YXZ', degrees=True)

        print("--- UNITY'S EXPORTED smpl_incam_rot (converted to CV) ---")
        print(f"global_orient_c (YXZ): yaw={euler_unity[0]:7.2f}°, pitch={euler_unity[1]:7.2f}°, roll={euler_unity[2]:7.2f}°")
        print()

        print("--- COMPARISON ---")
        diff_old = euler_old - euler_unity
        diff_new = euler_new - euler_unity
        print(f"OLD vs Unity: yaw={diff_old[0]:7.2f}°, pitch={diff_old[1]:7.2f}°, roll={diff_old[2]:7.2f}°")
        print(f"NEW vs Unity: yaw={diff_new[0]:7.2f}°, pitch={diff_new[1]:7.2f}°, roll={diff_new[2]:7.2f}°")

    print()
    print("="*70)
    print("DIAGNOSIS:")
    diff = euler_new - euler_old
    if np.allclose(diff, 0, atol=1.0):
        print("  Both methods give same result!")
        print("  The 168° roll error must come from somewhere else:")
        print("    - Check body_pose export")
        print("    - Check if Unity's smpl_incam_rot matches derived rotation")
        print("    - Check SMPL model forward pass")
    else:
        print(f"  Methods differ: yaw={diff[0]:.2f}°, pitch={diff[1]:.2f}°, roll={diff[2]:.2f}°")
        print("  This explains the rotation error!")

if __name__ == "__main__":
    test_actual_frame()