"""Retarget Dyn-HaMR to the supplied L20 with JSON coupling and moving wrist."""
import os
os.environ.setdefault('MUJOCO_GL','osmesa')
os.environ.setdefault('OMP_NUM_THREADS','4')
import argparse,json,time
from pathlib import Path
import xml.etree.ElementTree as ET
import numpy as np
import torch,mujoco
from scipy.spatial.transform import Rotation
from scipy.ndimage import gaussian_filter1d
from dex_retargeting.retargeting_config import RetargetingConfig
from l20_calibrated import ROOT,Calibration,CalibratedAdaptor,build_assets,palm_basis,CHAINS,TIPS

OUT=ROOT/'output/l20_2047635068'

def main():
    global OUT
    parser=argparse.ArgumentParser();parser.add_argument('--frames',type=int,default=0)
    parser.add_argument('--output-dir',type=Path,default=OUT)
    parser.add_argument('--consistent-objective',action=argparse.BooleanOptionalAction,default=True)
    args=parser.parse_args();OUT=args.output_dir.resolve()
    torch.set_num_threads(4)
    fixed,urdf=build_assets(OUT)
    cal=Calibration()
    config=dict(type='position',urdf_path=str(urdf),target_joint_names=cal.active,
                target_link_names=[f'landmark_{i:02d}' for i in range(21)],
                target_link_human_indices=list(range(21)),low_pass_alpha=-1,normal_delta=.002)
    retarget=RetargetingConfig.from_dict(config.copy()).build()
    if args.consistent_objective:
        from l20_consistent_optimizer import ConsistentPositionOptimizer
        retarget.optimizer.__class__=ConsistentPositionOptimizer
    robot=retarget.optimizer.robot
    adaptor=CalibratedAdaptor(robot,cal);retarget.optimizer.set_kinematic_adaptor(adaptor)
    bounds=np.array([cal.bounds[n] for n in cal.active]);retarget.joint_limits=bounds
    retarget.optimizer.set_joint_limit(bounds,epsilon=0)
    neutral_q=np.zeros(robot.dof)
    for i,n in enumerate(cal.active):neutral_q[robot.get_joint_index(n)]=np.clip(cal.values[n][cal.raw['baseline_command_u8'][cal.tables[n]['motor_index']]],*bounds[i])
    neutral_q=adaptor.forward_qpos(neutral_q)
    retarget.set_qpos(neutral_q)
    sites=[robot.get_link_index(f'landmark_{i:02d}') for i in range(21)]
    def fk(q):
        robot.compute_forward_kinematics(q)
        return np.array([robot.get_link_pose(k)[:3,3] for k in sites])
    neutral=fk(neutral_q);basis_robot=palm_basis(neutral)
    human=np.load(ROOT/'output/dex_2047635068/human_joints.npz')
    joints=human['joints'][:args.frames or None]
    targets=[];wrist_rot=[]
    for t,p in enumerate(joints):
        basis=palm_basis(p)
        align=basis@basis_robot.T
        aligned=(p-p[0])@align
        target=neutral.copy()
        for chain in CHAINS:
            for a,b in zip(chain[:-1],chain[1:]):
                direction=aligned[b]-aligned[a]
                target[b]=target[a]+direction/np.linalg.norm(direction)*np.linalg.norm(neutral[b]-neutral[a])
        targets.append(target)
        wrist_rot.append(Rotation.from_rotvec(human['root_orient'][t]).as_matrix()@align)
    targets=np.array(targets);wrist_rot=np.array(wrist_rot)
    qraw=[];codes=[];start=time.monotonic()
    for t,target in enumerate(targets):
        q=retarget.retarget(target)
        qraw.append(q.copy());codes.append(retarget.optimizer.opt.last_optimize_result())
        if t%200==0: print(f'L20 dex {t}/{len(targets)}',flush=True)
    qraw=np.array(qraw)
    active=qraw[:,adaptor.idx_pin2target]
    active=np.clip(gaussian_filter1d(active,1,axis=0),bounds[:,0],bounds[:,1])
    qpos=[]
    for a in active:
        q=np.zeros(robot.dof);q[adaptor.idx_pin2target]=a;qpos.append(adaptor.forward_qpos(q))
    qpos=np.array(qpos);actual=np.array([fk(q) for q in qpos])
    commands=np.array([cal.command(dict(zip(robot.dof_joint_names,q))) for q in qpos])
    decoded=np.array([cal.decode(c,robot.dof_joint_names) for c in commands])
    wrist_world=human['wrist_world'][:len(qpos)]
    # A single rigid transform makes the initial palm upright for visualization.
    scene_rot=wrist_rot[0].T
    wrist_pos=(wrist_world-wrist_world[0])@scene_rot.T+np.array([0,0,.2])
    wrist_scene=np.einsum('ij,tjk->tik',scene_rot,wrist_rot)
    quat=Rotation.from_matrix(wrist_scene).as_quat()[:,[3,0,1,2]]
    for i in range(1,len(quat)):
        if quat[i]@quat[i-1]<0:quat[i]*=-1
    np.savez_compressed(OUT/'motion.npz',qpos=qpos,qpos_raw=qraw,joint_names=robot.dof_joint_names,
        active_qpos=active,active_joint_names=cal.active,command_u8=commands,decoded_qpos=decoded,
        wrist_pos=wrist_pos,wrist_quat_wxyz=quat,wrist_world=wrist_world,wrist_world_R=wrist_rot,
        scene_rotation=scene_rot,scene_translation=np.array([0,0,.2])-scene_rot@wrist_world[0],
        targets=targets,actual=actual,fps=30,time=np.arange(len(qpos))/30)
    np.savetxt(OUT/'trajectory.csv',np.c_[np.arange(len(qpos))/30,wrist_pos,quat,qpos],delimiter=',',
       header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz']+robot.dof_joint_names),comments='')
    np.savetxt(OUT/'commands_u8.csv',np.c_[np.arange(len(qpos)),commands],delimiter=',',fmt='%d',
       header=','.join(['frame']+[f'slot_{i}' for i in range(20)]),comments='')
    limits=np.array([cal.urdf_limits[n] for n in robot.dof_joint_names])
    violation=float(np.maximum(limits[:,0]-qpos,0).max()+np.maximum(qpos-limits[:,1],0).max())
    assert violation<1e-6 and np.isfinite(qpos).all()
    err=np.linalg.norm(actual[:,TIPS]-targets[:,TIPS],axis=-1)*1000
    # Compare independent FK libraries, with all 21 nonlinear-coupled angles supplied.
    model=mujoco.MjModel.from_xml_path(str(fixed));data=mujoco.MjData(model)
    addresses=[model.jnt_qposadr[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_JOINT,n)] for n in robot.dof_joint_names]
    ids=[mujoco.mj_name2id(model,mujoco.mjtObj.mjOBJ_SITE,f'landmark_{i:02d}') for i in range(21)]
    fk_error=0
    for t in np.linspace(0,len(qpos)-1,min(20,len(qpos)),dtype=int):
        data.qpos[addresses]=qpos[t];mujoco.mj_forward(model,data)
        fk_error=max(fk_error,float(np.abs(data.site_xpos[ids]-actual[t]).max()))
    assert fk_error<1e-5
    report=dict(frames=len(qpos),fps=30,independent_joints=16,urdf_joints=21,
       calibration_model=cal.raw['model'],source_directory_name='L20',
       calibration_interpretation='q_urdf=angle_rad[command]+urdf_zero_offset_rad; passive joints share motor command; not hardware verified',
       reserved_command_slots=[11,12,13,14],tip_error_mean_mm=float(err.mean()),tip_error_p95_mm=float(np.percentile(err,95)),
       tip_error_finger_mm=dict(zip(['thumb','index','middle','ring','pinky'],err.mean(0).tolist())),
       wrist_translation_range_m=np.ptp(wrist_pos,axis=0).tolist(),
       wrist_rotation_max_from_initial_rad=float(np.linalg.norm(Rotation.from_matrix(wrist_scene).as_rotvec(),axis=1).max()),
       urdf_limit_violation_rad=violation,quantization_max_error_rad=float(np.abs(decoded-qpos).max()),
       pinocchio_mujoco_fk_error_m=fk_error,seconds=time.monotonic()-start,
       optimizer_status_counts={str(k):int(v) for k,v in zip(*np.unique(codes,return_counts=True))},
       consistent_objective=args.consistent_objective,
       object_contact_stage='not applied in this hand-only retargeting run')
    (OUT/'retarget_validation.json').write_text(json.dumps(report,indent=2))
    (OUT/'dex_config.json').write_text(json.dumps(config,indent=2))
    # Moving wrist replay model; contact scene construction is a separate stage.
    tree=ET.parse(fixed);root=tree.getroot()
    body=root.find("worldbody/body[@name='hand_base_link']")
    ET.SubElement(body,'freejoint',name='wrist_free')
    root.find('option').set('gravity','0 0 0')
    tree.write(OUT/'l20_moving.xml')
    print(json.dumps(report,indent=2),flush=True)

if __name__=='__main__':main()