"""Use observed skin-surface depth to correct camera-space hand translation.

This is a robust translation fit, not depth-derived skeleton ground truth.
"""
from pathlib import Path
import json
import argparse
import numpy as np
import cv2
from scipy.ndimage import gaussian_filter1d
from scipy.spatial.transform import Rotation

ROOT=Path(__file__).resolve().parents[1];BASE=ROOT/'output/20260915_171525'
SRC=ROOT/'docs/20260915_171525'
parser=argparse.ArgumentParser();parser.add_argument('--left',action='store_true');args=parser.parse_args()
folder='handflow_left' if args.left else 'handflow'
suffix='_left' if args.left else ''
h=np.load(BASE/folder/'handflow_results.npz');j=dict(np.load(BASE/folder/'human_joints.npz'))
meta=json.loads((SRC/'intrinsics.json').read_text());camera=np.load(BASE/'rgbd_camera.npz')
K=np.array([meta[k] for k in ['fx','fy','cx','cy']]);N=len(h['pose'])
c2w=camera['c2w'];old_camera=h['c2w']
wrist_cam=np.einsum('tji,tj->ti',old_camera[:,:3,:3],j['wrist_world']-old_camera[:,:3,3])
shifts=[];records=[];valid=[];cap=cv2.VideoCapture(str(SRC/'color.mp4'))
for t,verts in enumerate(h['verts_cam']):
    ok,color=cap.read();assert ok
    dep=cv2.imread(str(SRC/'depth'/f'{t:06d}.png'),-1).astype(float)*meta['depth_scale_m']
    ycc=cv2.cvtColor(color,cv2.COLOR_BGR2YCrCb)
    skin=cv2.inRange(ycc,np.array([0,133,77]),np.array([255,173,127]))>0
    uv=np.rint(verts[:,:2]/verts[:,2,None]*K[:2]+K[2:]).astype(int)
    keep=(verts[:,2]>.05)&(uv[:,0]>=2)&(uv[:,0]<color.shape[1]-2)&(uv[:,1]>=2)&(uv[:,1]<color.shape[0]-2)
    # Keep nearest projected vertex per 4x4 pixel bin, reducing hidden-surface bias.
    closest={}
    for index in np.flatnonzero(keep):
        u,v=uv[index];key=(u//4,v//4)
        if key not in closest or verts[index,2]<verts[closest[key],2]:closest[key]=index
    errors=[]
    for index in closest.values():
        u,v=uv[index]
        if not skin[v,u]:continue
        patch=dep[v-1:v+2,u-1:u+2];values=patch[(patch>.1)&(patch<.85)]
        if len(values)<5 or np.ptp(values)>.04:continue
        errors.append(float(np.median(values)-verts[index,2]))
    dz=float(np.median(errors)) if errors else 0.
    mad=float(np.median(np.abs(np.asarray(errors)-dz))) if errors else 1.
    accepted=bool(len(errors)>=25 and mad<.025 and abs(dz)<.25 and h['pred_valid'][t])
    shifts.append(dz);valid.append(accepted)
    records.append(dict(frame=t,samples=len(errors),median_depth_residual_m=dz,mad_m=mad,accepted=accepted))
cap.release();valid=np.asarray(valid);idx=np.flatnonzero(valid)
assert len(idx)>=N*.5, f'Insufficient depth support: {len(idx)}/{N}'
filled=np.interp(np.arange(N),idx,np.asarray(shifts)[idx]);smooth=gaussian_filter1d(filled,2)
delta=wrist_cam/wrist_cam[:,2,None]*smooth[:,None]
corrected=wrist_cam+delta
world=np.einsum('tij,tj->ti',c2w[:,:3,:3],corrected)+c2w[:,:3,3]
root_world=Rotation.from_matrix(c2w[:,:3,:3]@Rotation.from_rotvec(h['pose'][:,:3]).as_matrix()).as_rotvec()
j.update(wrist_world=world,root_orient=root_world,time=camera['time'],
         before_depth_wrist_camera=wrist_cam,depth_translation_camera=delta,
         depth_valid=valid,depth_correction_interpolated=~valid,
         detection_valid=j['detection_valid']&valid&camera['valid'],
         source=str((BASE/folder/'handflow_results.npz').resolve()),
         metric_scale_provenance='D405 depth_scale from recording metadata; skin-surface translation fit, model size unchanged; RGBD static-background odometry world. Not calibrated motion ground truth.')
np.savez_compressed(BASE/f'human_joints_rgbd{suffix}.npz',**j)
np.savez_compressed(BASE/f'depth_hand_fit{suffix}.npz',raw_depth_shift=shifts,applied_depth_shift=smooth,
                    accepted=valid,wrist_camera=corrected,wrist_world=world,c2w=c2w,
                    corrected_verts_camera=h['verts_cam']+delta[:,None,:])
report=dict(frames=N,depth_supported_frames=int(valid.sum()),rejected_depth_frames=np.flatnonzero(~valid).tolist(),
            median_abs_raw_depth_residual_m=float(np.median(np.abs(np.asarray(shifts)[valid]))),
            median_abs_remaining_sample_residual_m=float(np.median(np.abs(np.asarray(shifts)[valid]-smooth[valid]))),
            applied_depth_shift_range_m=[float(smooth.min()),float(smooth.max())],frames_detail=records,
            boundary='Translation-only fit to selected visible skin depth; self-occlusion and skin-mask errors remain. No finger-joint depth truth. Rejected depth corrections interpolated and marked invalid.')
(BASE/f'depth_fit_validation{suffix}.json').write_text(json.dumps(report,indent=2));print(json.dumps({k:v for k,v in report.items() if k!='frames_detail'},indent=2))