"""Read-only trajectory diagnosis; frame differences are not ground-truth errors."""
import json
import sys
from pathlib import Path
import numpy as np
import torch
from scipy.spatial.transform import Rotation
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT/'third_party/Dyn-HaMR/dyn-hamr'))
from body_model import MANO, run_mano
from vis.tools import smooth_results
torch.set_num_threads(4)
BASE = ROOT/'output/20260915_171525_dynhamr'
OUT = BASE/'diagnosis'
OUT.mkdir(exist_ok=True)
def stats(x):
    x=np.asarray(x)
    return dict(rms=float(np.sqrt(np.mean(x*x))), p95=float(np.percentile(x,95)), maximum=float(np.max(x)))
def positions(x):
    d=np.linalg.norm(np.diff(x,axis=0),axis=-1)*1000
    dd=np.linalg.norm(np.diff(x,n=2,axis=0),axis=-1)*1000
    return dict(step_mm=stats(d), second_difference_mm=stats(dd), top_step_frames=(np.argsort(d)[-8:][::-1]+1).tolist())
def rotations(x):
    r=Rotation.from_rotvec(x.reshape(-1,3)).as_matrix().reshape(x.shape[:-1]+(3,3))
    inc=r[1:]@np.swapaxes(r[:-1],-1,-2)
    a=Rotation.from_matrix(inc.reshape(-1,3,3)).magnitude()*180/np.pi
    return stats(a)
model=MANO(model_path=str(ROOT/'third_party/Dyn-HaMR/_DATA/data/mano'),batch_size=704,pose2rot=True)
report={}
saved={}
for stage in ['smooth_fit','prior']:
    path=sorted((BASE/'optimization'/stage).glob('*world_results.npz'))[-1]
    arr=np.load(path)
    for filtered in [False,True]:
        d={k:torch.from_numpy(arr[k].copy()).float() for k in ['trans','root_orient','pose_body','is_right','betas']}
        if filtered:
            d['root_orient'],d['pose_body'],d['betas'],d['trans']=smooth_results(d['root_orient'],d['pose_body'],d['betas'],d['is_right'],d['trans'])
        with torch.no_grad():
            j=run_mano(model,d['trans'],d['root_orient'],d['pose_body'],d['is_right'],d['betas'])['joints'].numpy()
        key=stage+('_render_filtered' if filtered else '_raw')
        saved[key]=j
        report[key]={}
        for b in range(2):
            side='right' if arr['is_right'][b,0]>.5 else 'left'
            report[key][side]=dict(wrist_world=positions(j[b,:,0]),root_rotation_step_deg=rotations(d['root_orient'][b].numpy()),finger_rotation_step_deg=rotations(d['pose_body'][b].numpy().reshape(352,15,3)))
        report[key]['relative_wrist']=positions(j[1,:,0]-j[0,:,0])
        if stage=='prior':
            R=arr['cam_R'][0]; t=arr['cam_t'][0]
            for b in range(2):
                side='right' if arr['is_right'][b,0]>.5 else 'left'
                report[key][side]['wrist_camera']=positions(np.einsum('tij,tj->ti',R,j[b,:,0])+t)
cam=np.load(ROOT/'output/20260915_171525/rgbd_camera.npz')
report['camera_position']=positions(cam['c2w'][:,:3,3])
report['camera_rotation_step_deg']=rotations(Rotation.from_matrix(cam['c2w'][:,:3,:3]).as_rotvec())
report['caveat']='Frame differences include true motion; not a causal decomposition or accuracy metric. Render filtering replicated for static views; source camera render excludes translation filtering.'
np.savez_compressed(OUT/'joints_diagnostic.npz',**saved,time=cam['time'])
(OUT/'metrics.json').write_text(json.dumps(report,indent=2))
print(json.dumps(report,indent=2))