import os
import json
import numpy as np
import torch
from scipy.spatial.transform import Rotation as R
from hmr4d.utils.smplx_utils import make_smplx

def test_kabsch():
    input_dir = "/mnt/c/Temp/SyntheticDataset"
    seqs = ["sequence_100_biboo_birthday_speech_explosion_1.jsonl", "sequence_101_biboo_birthday_speech_explosion_2.jsonl"]
    
    device = "cuda"
    model = make_smplx("supermotion").to(device).eval()

    fix_rot = R.from_euler("z", 180, degrees=True).as_matrix()
    C = np.diag([1.0, -1.0, 1.0])
    C4 = np.diag([1.0, -1.0, 1.0, 1.0])
    fix_mat = np.eye(4); fix_mat[:3, :3] = fix_rot

    for seq in seqs:
        p = os.path.join(input_dir, seq)
        with open(p, "r") as f: lines = f.readlines()
        
        row = json.loads(lines[1]) # frame 0
        
        # parse unity
        R_cam_w = R.from_quat(row["cam_rot_world"]).as_matrix()
        R_pel_w = R.from_quat(row["pelvis_rot_world"]).as_matrix()
        R_rel_unity = R_cam_w.T @ R_pel_w
        R_cv = C @ R_rel_unity @ C
        R_final = R_cv @ R.from_euler("z", 180, degrees=True).as_matrix()
        go_c = R.from_matrix(R_final).as_rotvec().astype(np.float32)

        pelvis_cam_cv = np.asarray(row["smpl_incam_transl"], dtype=np.float64).reshape(3)
        root_cam_cv = np.asarray(row.get("smpl_root_incam_transl", [0.0, 0.0, 0.0]), dtype=np.float64).reshape(3)
        pelvis_cam_cv = pelvis_cam_cv # + np.array([0.0, -0.02, 0.0], dtype=np.float64)
        target_cam_cv = pelvis_cam_cv

        R_cam_w_cv = fix_rot @ (C @ R_cam_w @ C)
        R_pelvis_w_cv = R_cam_w_cv @ R_final
        go_w = R.from_matrix(R_pelvis_w_cv).as_rotvec().astype(np.float32)
        
        pos_cv_raw = (C @ row['pelvis_pos_world'])
        pelvis_pos_w_cv = fix_rot @ pos_cv_raw

        pose = np.asarray(row["smplx_pose"], dtype=np.float32)
        body_pose = pose[3:66].astype(np.float32)
        betas = np.zeros(10, dtype=np.float32)
        
        # compute pel0
        with torch.no_grad():
            bt = torch.from_numpy(betas[None]).to(device)
            bp = torch.from_numpy(body_pose[None]).to(device)
            
            go_c_t = torch.from_numpy(go_c[None]).to(device)
            pel0_c = model(betas=bt, global_orient=go_c_t, body_pose=bp).joints[0, 0].cpu().numpy()
            
            go_w_t = torch.from_numpy(go_w[None]).to(device)
            pel0_w = model(betas=bt, global_orient=go_w_t, body_pose=bp).joints[0, 0].cpu().numpy()

        tr_c = target_cam_cv - pel0_c
        tr_w = pelvis_pos_w_cv - pel0_w
        
        # math check
        R_c_v = R.from_rotvec(go_c).as_matrix()
        R_w_v = R.from_rotvec(go_w).as_matrix()
        R_w2c_v = R_c_v @ R_w_v.T
        
        j0_c = pel0_c + tr_c
        j0_w = pel0_w + tr_w
        t_w2c_derived = j0_c - R_w2c_v @ j0_w

        cam_T_wc = np.eye(4)
        cam_T_wc[:3, :3] = R_cam_w
        cam_T_wc[:3, 3] = row["cam_pos_world"]
        cam_T_wc_cv = fix_mat @ (C4 @ cam_T_wc @ C4)
        T_w2c_exported = np.linalg.inv(cam_T_wc_cv)
        
        # Kabsch verify manually
        T_c2w_exported = np.linalg.inv(T_w2c_exported)
        
        with torch.no_grad():
            tr_c_t = torch.from_numpy(tr_c[None]).to(device)
            j_c = model(betas=bt, global_orient=go_c_t, body_pose=bp, transl=tr_c_t).joints[0, :22].cpu().numpy()
            
            tr_w_t = torch.from_numpy(tr_w[None]).to(device)
            j_w = model(betas=bt, global_orient=go_w_t, body_pose=bp, transl=tr_w_t).joints[0, :22].cpu().numpy()
        
        src_mean = j_c.mean(axis=0)
        dst_mean = j_w.mean(axis=0)
        X = j_c - src_mean
        Y = j_w - dst_mean
        H = X.T @ Y
        U, _, Vt = np.linalg.svd(H)
        R_align = Vt.T @ U.T
        if np.linalg.det(R_align) < 0:
            Vt[-1, :] *= -1
            R_align = Vt.T @ U.T
        t_align = dst_mean - src_mean @ R_align
        
        T_c2w_kabsch = np.eye(4)
        T_c2w_kabsch[:3, :3] = R_align
        T_c2w_kabsch[:3, 3] = t_align
        
        t_err = np.linalg.norm(T_c2w_kabsch[:3, 3] - T_c2w_exported[:3, 3])
        
        print(f"[{seq}]")
        print(f"  t_err: {t_err:.3f}m")
        print(f"  t_c2w_kabsch: {T_c2w_kabsch[:3, 3]}")
        print(f"  T_c2w_exported: {T_c2w_exported[:3, 3]}")
        print(f"  diff(kabsch - exported): {T_c2w_kabsch[:3, 3] - T_c2w_exported[:3, 3]}")
        print(f"  t_w2c_derived: {t_w2c_derived}")
        print(f"  T_w2c_exported: {T_w2c_exported[:3, 3]}")
        print(f"  delta_T_w2c (derived - exported): {t_w2c_derived - T_w2c_exported[:3, 3]}")

if __name__ == "__main__":
    test_kabsch()