"""Export Dyn-HaMR MANO joints, preserving the exact reconstruction FK convention."""
import os
os.environ.setdefault('OMP_NUM_THREADS', '4')
import sys
from pathlib import Path
import numpy as np
import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'third_party/Dyn-HaMR/dyn-hamr'))
from body_model.mano_wrapper import MANO

def main():
    torch.set_num_threads(4)
    source = ROOT / 'output/results_2047635068/world_results.npz'
    out = ROOT / 'output/dex_2047635068'
    out.mkdir(exist_ok=True)
    data = np.load(source)
    candidates = np.flatnonzero(np.all(data['is_right'] > .5, axis=1))
    assert len(candidates) == 1, 'Expected exactly one consistent right-hand track'
    track = int(candidates[0])
    pose = data['pose_body'][track].reshape(-1, 45)
    model = MANO(model_path=str(ROOT / 'third_party/Dyn-HaMR/_DATA/data/mano'),
                 batch_size=128, pose2rot=True)
    joints = []
    wrists = []
    with torch.no_grad():
        for start in range(0, len(pose), 128):
            p = torch.from_numpy(pose[start:start+128]).float()
            n = len(p)
            result = model(hand_pose=p, global_orient=torch.zeros(n, 3),
                           transl=torch.zeros(n, 3),
                           betas=torch.from_numpy(data['betas'][track]).float().expand(n, -1))
            j = result.joints.numpy()
            joints.append(j - j[:, :1])
            world = model(hand_pose=p,
                          global_orient=torch.from_numpy(data['root_orient'][track,start:start+n]).float(),
                          transl=torch.from_numpy(data['trans'][track,start:start+n]).float(),
                          betas=torch.from_numpy(data['betas'][track]).float().expand(n,-1))
            wrists.append(world.joints[:,0].numpy())
    joints = np.concatenate(joints)
    assert joints.shape == (len(pose), 21, 3) and np.isfinite(joints).all()
    np.savez_compressed(out / 'human_joints.npz', joints=joints, fps=30,
                        source_track=track, source=str(source.resolve()),
                        wrist_world=np.concatenate(wrists),
                        root_orient=data['root_orient'][track], trans=data['trans'][track])
    print(f'Exported {joints.shape}, right track {track}, to {out}', flush=True)

if __name__ == '__main__':
    main()