"""Register yesterday's Dyn-HaMR hands to the RGB-D odometry world used by the object tracks.

Three corrections, all rigid/similarity, no articulation change:
1. World frame: Dyn-HaMR world differs from the RGB-D odometry world (object tracks, table plane)
   by a constant translation (~12.6 cm). Hands are re-expressed in the odometry world.
2. Scale: MANO silhouettes match the image at the wrong depth (right hand ~10% too close),
   so a per-hand constant scale about the camera centre fixes depth while keeping the silhouette.
3. Residual per-frame depth shift along the wrist ray (smoothed), from the held-out depth fit.
Outputs go to output/20260915_171525_dynhamr/registered/ ; nothing upstream is overwritten.
"""
import os, sys, json, shutil
os.environ.setdefault('OMP_NUM_THREADS', '4')
from pathlib import Path
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'third_party/Dyn-HaMR/dyn-hamr')); sys.path.insert(0, str(ROOT / 'third_party/FoundationPose')); sys.path.insert(0, str(ROOT / 'scripts'))
import numpy as np, cv2, torch, trimesh, open3d as o3d, nvdiffrast.torch as dr
from scipy.ndimage import gaussian_filter1d
from body_model import MANO, run_mano
from Utils import nvdiffrast_render
from red_box_distance import signed_distance

SRC = ROOT / 'docs/20260915_171525'; BASE = ROOT / 'output/20260915_171525_dynhamr'
PRIOR = BASE / 'corrected/prior/20260915_171525_000000_world_results.npz'
DIAG = ROOT / 'output/hand_depth_diagnosis_20260915/frame_metrics.json'
FP = ROOT / 'output/foundationpose_spider_20260915/foundationpose_objects.npz'
OUT = Path(os.environ.get('HF_REG_OUT', BASE / 'registered')); OUT.mkdir(exist_ok=True)
CAMERA = Path(os.environ.get('HF_CAMERA_FILE', BASE / 'rgbd_camera.npz')); ALIGN = os.environ.get('HF_ALIGN', 'wrist')   # 'wrist' | 'mcp'
N = 352; torch.set_num_threads(4)
meta = json.loads((SRC / 'intrinsics.json').read_text()); K = np.array([[meta['fx'], 0, meta['cx']], [0, meta['fy'], meta['cy']], [0, 0, 1.]])

p = dict(np.load(PRIOR)); cam = np.load(CAMERA); fp = np.load(FP); fm = json.load(open(DIAG))
model = MANO(model_path=str(ROOT / 'third_party/Dyn-HaMR/_DATA/data/mano'), batch_size=2 * N, pose2rot=True)
with torch.no_grad():
    o = run_mano(model, *[torch.tensor(p[k]).float() for k in ['trans', 'root_orient', 'pose_body', 'is_right', 'betas']])
J = o['joints'].numpy(); V = o['vertices'].numpy(); faces = np.asarray(model.faces if hasattr(model, 'faces') else model.bm.faces).astype(np.int32)
sides = ['right' if p['is_right'][b, 0] > .5 else 'left' for b in range(2)]
R = p['cam_R'][0]; tD = p['cam_t'][0]; c2w = cam['c2w']
w2cR = np.transpose(c2w[:, :3, :3], (0, 2, 1)); tO = -np.einsum('tij,tj->ti', w2cR, c2w[:, :3, 3])
rot_gap = float(np.abs(R - w2cR).max()); assert rot_gap < (1e-5 if ALIGN == 'wrist' and CAMERA.name == 'rgbd_camera.npz' else 0.2), 'camera rotation differs too much'
offset = np.einsum('tji,tj->ti', R, tD - tO); offset_mean = offset.mean(0)
if CAMERA.name == 'rgbd_camera.npz': assert offset.std(0).max() < .005
Jc = np.einsum('tij,btkj->btki', R, J) + tD[None, :, None, :]; Vc = np.einsum('tij,btvj->btvi', R, V) + tD[None, :, None, :]

ctx = dr.RasterizeCudaContext(); eye = torch.eye(4, device='cuda')[None]; ft = torch.as_tensor(faces, device='cuda', dtype=torch.int32)
def render(v):
    mt = {'pos': torch.as_tensor(v, device='cuda', dtype=torch.float32), 'faces': ft, 'vnormals': torch.zeros((len(v), 3), device='cuda'), 'vertex_color': torch.ones((len(v), 3), device='cuda')}
    with torch.inference_mode(): _, d, _ = nvdiffrast_render(K=K, H=480, W=848, ob_in_cams=eye, glctx=ctx, mesh_tensors=mt)
    return d[0].cpu().numpy()
scenes = {}
for name in ['upper', 'lower']:
    mesh = trimesh.load(ROOT / f'output/collision_fix_20260915/collision_v2/{name}_closed.ply', process=False); assert mesh.is_watertight
    sc = o3d.t.geometry.RaycastingScene(); sc.add_triangles(o3d.t.geometry.TriangleMesh.from_legacy(o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(mesh.vertices), o3d.utility.Vector3iVector(mesh.faces)))); scenes[name] = (sc, mesh)
def tip_gap(pts_world, name, t):
    T = fp[name + '_T_world'][t]; local = (pts_world - T[:3, 3]) @ T[:3, :3]
    return float(np.min(np.abs(signed_distance(scenes[name][0], scenes[name][1], local))) * 1000)

yy, xx = np.indices((480, 848)); train = ((xx // 8 + yy // 8) % 2) == 0
report = {'camera_file': str(CAMERA), 'l20_alignment': ALIGN, 'offset_dynhamr_to_odometry_world_m': offset_mean.tolist(), 'offset_norm_mm': float(np.linalg.norm(offset_mean) * 1000), 'offset_std_mm': (offset.std(0) * 1000).tolist(), 'hands': {}}
Jw_new = np.zeros_like(J); Vc_new = np.zeros_like(Vc)
for b, side in enumerate(sides):
    rows = {r['frame']: r for r in fm[side]}
    wz = Jc[b, :, 0, 2]; dz = np.array([rows[t]['fit_shift_z_mm'] / 1000 if rows[t]['depth_supported'] else np.nan for t in range(N)])
    sup = np.isfinite(dz); k = float(np.median(1 + dz[sup] / wz[sup]))
    resid = dz - (k - 1) * wz; resid = np.interp(np.arange(N), np.flatnonzero(sup), resid[sup]); resid = gaussian_filter1d(resid, 2)
    ray = Jc[b, :, 0] / wz[:, None]
    Jc_new = k * Jc[b] + (ray * resid[:, None])[:, None, :]; Vc_new[b] = k * Vc[b] + (ray * resid[:, None])[:, None, :]
    Jw_new[b] = np.einsum('tji,tkj->tki', w2cR, Jc_new - tO[:, None, :])   # world of the chosen camera file
    # independent check: odometry camera must reproduce the corrected camera coordinates
    back = np.einsum('tij,tkj->tki', w2cR, Jw_new[b]) + tO[:, None, :]; assert np.abs(back - Jc_new).max() < 1e-5
    report['hands'][side] = {'scale_about_camera': k, 'depth_supported_frames': int(sup.sum()), 'residual_shift_mm': {'median_abs': float(np.median(np.abs(resid)) * 1000), 'max_abs': float(np.abs(resid).max() * 1000)},
                             'wrist_move_mm': {'median': float(np.median(np.linalg.norm(Jw_new[b, :, 0] - J[b, :, 0], axis=1)) * 1000), 'max': float(np.max(np.linalg.norm(Jw_new[b, :, 0] - J[b, :, 0], axis=1)) * 1000)}}
# held-out depth residual after correction (same masks/method as the diagnosis), plus tip gaps
cap = cv2.VideoCapture(str(SRC / 'color.mp4')); after = {s: [] for s in sides}; gaps = {s: {'before': [], 'after': []} for s in sides}; keep = {60, 175, 240, 300}
pair = {'right': 'upper', 'left': 'lower'}
for t in range(N):
    ok, bgr = cap.read(); assert ok
    z = cv2.imread(str(SRC / 'depth' / f'{t:06d}.png'), -1).astype('float32') * meta['depth_scale_m']
    hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV); skin = cv2.inRange(cv2.cvtColor(bgr, cv2.COLOR_BGR2YCrCb), np.array([0, 133, 77]), np.array([255, 173, 127])) > 0
    skin &= ~(((hsv[:, :, 0] < 10) | (hsv[:, :, 0] > 170)) & (hsv[:, :, 1] > 115)); skin &= ~((hsv[:, :, 0] > 95) & (hsv[:, :, 0] < 130) & (hsv[:, :, 1] > 115))
    d1 = [render(Vc_new[b, t]) for b in range(2)]
    for b, side in enumerate(sides):
        other = d1[1 - b]; occ = (other > .1) & (other < d1[b]); interior = cv2.erode((d1[b] > .1).astype('uint8'), np.ones((5, 5), np.uint8)) > 0
        test = skin & interior & (z > .1) & (z < .85) & ~occ & ~train
        if test.sum() >= 100: after[side].append(float(np.median(np.abs(z[test] - d1[b][test])) * 1000))
        name = pair[side]; lim = (t >= 150) if side == 'right' else (t < 177)
        if fp[name + '_valid'][t] and lim:
            tips = J[b, t, [4, 8, 12, 16, 20]]; gaps[side]['before'].append(tip_gap(tips, name, t)); gaps[side]['after'].append(tip_gap(Jw_new[b, t, [4, 8, 12, 16, 20]], name, t))
    if t in keep:
        img = bgr.copy()
        for b, col in enumerate([(255, 128, 0), (0, 200, 255)]):
            cnts, _ = cv2.findContours((d1[b] > .1).astype('uint8'), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE); cv2.drawContours(img, cnts, -1, col, 2)
        cv2.imwrite(str(OUT / f'overlay_registered_{t:04d}.jpg'), img)
    if t % 100 == 0: print('frame', t, flush=True)
cap.release()
def rms2(x): return float(np.sqrt(np.mean(np.sum(np.diff(x, n=2, axis=0) ** 2, axis=-1))) * 1000)
for b, side in enumerate(sides):
    before = [r['abs_residual_before_mm'] for r in fm[side] if r['abs_residual_before_mm'] is not None]
    report['hands'][side].update({'heldout_abs_residual_mm': {'before_median': float(np.median(before)), 'after_median': float(np.median(after[side])), 'after_p95': float(np.percentile(after[side], 95))},
                                  'nearest_tip_to_box_mm': {'pipeline_before_median': float(np.median(gaps[side]['before'])), 'after_median': float(np.median(gaps[side]['after'])), 'frames': len(gaps[side]['after'])},
                                  'wrist_second_difference_rms_mm': {'before': rms2(J[b, :, 0]), 'after': rms2(Jw_new[b, :, 0])}})
# exports: human joints + L20 motion with wrist replaced; finger retargeting is direction-only so joints are unchanged
for b, side in enumerate(sides):
    old = dict(np.load(BASE / f'human_joints_{side}.npz', allow_pickle=True))
    old.update(joints=Jw_new[b] - Jw_new[b, :, :1], wrist_world=Jw_new[b, :, 0], source=str(PRIOR), depth_correction_applied=True,
               coordinate_note='RGB-D odometry world (first camera frame), same as FoundationPose object tracks. Constant world offset applied; per-hand scale about camera centre; smoothed residual depth shift along wrist ray. Wrist-relative joints scaled accordingly.',
               scale_about_camera=np.float64(report['hands'][side]['scale_about_camera']), world_offset_m=offset_mean)
    np.savez_compressed(OUT / f'human_joints_{side}.npz', **old)
    src = BASE / f'l20_{side}_stable'; dst = OUT / f'l20_{side}_stable'; dst.mkdir(exist_ok=True)
    for f in src.iterdir():
        if f.name != 'motion.npz': shutil.copy2(f, dst / f.name)
    mo = dict(np.load(src / 'motion.npz', allow_pickle=True)); assert np.abs(mo['wrist_world'] - J[b, :, 0]).max() < 1e-4, 'old motion wrist differs from prior joints'
    ww = Jw_new[b, :, 0].copy()
    if ALIGN == 'mcp':
        # place the L20 so its knuckle row (MCP centroid) coincides with the human's; the wrist retreats by the palm-length difference
        v = mo['actual'][:, [5, 9, 13, 17]].mean(1) - mo['actual'][:, 0]; mcp_h = Jw_new[b][:, [5, 9, 13, 17]].mean(1)
        ww = mcp_h - np.einsum('tij,tj->ti', mo['wrist_world_R'], v)
        report['hands'][side]['mcp_alignment_wrist_retreat_mm'] = {'median': float(np.median(np.linalg.norm(ww - Jw_new[b, :, 0], axis=1)) * 1000), 'max': float(np.max(np.linalg.norm(ww - Jw_new[b, :, 0], axis=1)) * 1000)}
    rot = mo['scene_rotation']; wp = (ww - ww[0]) @ rot.T + np.array([0, 0, .2])
    mo.update(wrist_world=ww, wrist_pos=wp, wrist_pos_unsmoothed=wp, scene_translation=np.array([0, 0, .2]) - rot @ ww[0], source=str(OUT / f'human_joints_{side}.npz'), l20_alignment=ALIGN)
    np.savez_compressed(dst / 'motion.npz', **mo)
    chk = (mo['wrist_pos'] - mo['scene_translation']) @ mo['scene_rotation']; assert np.abs(chk - ww).max() < 1e-6
(OUT / 'registration.json').write_text(json.dumps(report, indent=2)); print(json.dumps(report, indent=2))