"""Held-out pixel depth residual of yesterday's (20260915_171525) Dyn-HaMR hand trajectory.

Same method as scripts/compare_red_hand_depth.py: render MANO hand depth into the D405
camera, select visible skin pixels, fit a single camera-ray translation on one half of an
8x8 checkerboard, evaluate the residual on the other half. No smoothing, no articulation change.
"""
import os, sys, json
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'))
import numpy as np, cv2, torch, nvdiffrast.torch as dr
from body_model import MANO, run_mano
from Utils import nvdiffrast_render

SRC = ROOT / 'docs/20260915_171525'
BASE = ROOT / 'output/20260915_171525_dynhamr'
PRIOR = BASE / 'corrected/prior/20260915_171525_000000_world_results.npz'
OUT = ROOT / 'output/hand_depth_diagnosis_20260915'; OUT.mkdir(exist_ok=True)
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.]])
H, W = 480, 848
N = 352

p = dict(np.load(PRIOR))
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()          # (2,N,J,3), (2,N,778,3) world, metres
faces = None
for attr in ['faces', 'faces_tensor']:
    f = getattr(model, attr, None)
    if f is None: f = getattr(getattr(model, 'bm', None), attr, None)
    if f is not None:
        faces = np.asarray(f.cpu() if torch.is_tensor(f) else f).astype(np.int32); break
assert faces is not None and faces.max() < V.shape[2], 'MANO faces not found'
R = p['cam_R'][0]; ct = p['cam_t'][0]
Vc = np.einsum('tij,btvj->btvi', R, V) + ct[None, :, None, :]
Jc = np.einsum('tij,btkj->btki', R, J) + ct[None, :, None, :]
sides = ['right' if p['is_right'][b, 0] > .5 else 'left' for b in range(2)]

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=H, W=W, ob_in_cams=eye, glctx=ctx, mesh_tensors=mt)
    return d[0].cpu().numpy()

yy, xx = np.indices((H, W)); train = ((xx // 8 + yy // 8) % 2) == 0
def stat(x):
    x = np.asarray([v for v in x if v is not None], float); x = x[np.isfinite(x)]
    return {'n': int(len(x)), 'median': float(np.median(x)), 'p5': float(np.percentile(x, 5)), 'p95': float(np.percentile(x, 95)), 'mean': float(np.mean(x))} if len(x) else None

rows = {s: [] for s in sides}; deltas = np.zeros((2, N, 3)); supported = np.zeros((2, N), bool)
cap = cv2.VideoCapture(str(SRC / 'color.mp4')); keep = {0, 60, 120, 175, 240, 300, 351}
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']
    ycc = cv2.cvtColor(bgr, cv2.COLOR_BGR2YCrCb); hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    skin = cv2.inRange(ycc, np.array([0, 133, 77]), np.array([255, 173, 127])) > 0
    red = ((hsv[:, :, 0] < 10) | (hsv[:, :, 0] > 170)) & (hsv[:, :, 1] > 115)
    blue = (hsv[:, :, 0] > 95) & (hsv[:, :, 0] < 130) & (hsv[:, :, 1] > 115)
    skin &= ~red & ~blue
    d0 = [render(Vc[b, t]) for b in range(2)]
    vis = []
    for b in range(2):
        other = d0[1 - b]
        occluded = (other > .1) & (other < d0[b])
        interior = cv2.erode((d0[b] > .1).astype('uint8'), np.ones((5, 5), np.uint8)) > 0
        mask = skin & interior & (z > .1) & (z < .85) & ~occluded
        fit = mask & train; err = z[fit] - d0[b][fit]
        delta = float(np.median(err)) if len(err) else 0.; mad = float(np.median(np.abs(err - delta))) if len(err) else 1.
        wz = float(Jc[b, t, 0, 2])
        accept = len(err) >= 100 and mad < .025 and abs(delta) < .25 and wz > .1
        r = {'frame': t, 'side': sides[b], 'wrist_z_m': wz, 'mask_pixels': int(mask.sum()), 'fit_pixels': int(fit.sum()),
             'fit_mad_mm': mad * 1000, 'fit_shift_z_mm': delta * 1000, 'depth_supported': bool(accept),
             'passes_yesterday_gate': bool(len(err) >= 25 and mad < .010 and abs(delta) < .025)}
        test0 = mask & ~train & (d0[b] > .1)
        r['heldout_pixels'] = int(test0.sum())
        r['signed_residual_before_mm'] = float(np.median(z[test0] - d0[b][test0]) * 1000) if test0.any() else None
        r['abs_residual_before_mm'] = float(np.median(np.abs(z[test0] - d0[b][test0])) * 1000) if test0.any() else None
        d1 = d0[b]
        if accept:
            deltas[b, t] = Jc[b, t, 0] / Jc[b, t, 0, 2] * delta; supported[b, t] = True
            d1 = render(Vc[b, t] + deltas[b, t])
            test1 = mask & ~train & (d1 > .1)
            r['abs_residual_after_mm'] = float(np.median(np.abs(z[test1] - d1[test1])) * 1000) if test1.any() else None
            r['translation_norm_mm'] = float(np.linalg.norm(deltas[b, t]) * 1000)
        else:
            r['abs_residual_after_mm'] = None; r['translation_norm_mm'] = None
        rows[sides[b]].append(r); vis.append((d0[b], d1))
    if t in keep:
        img = bgr.copy()
        for b, col in enumerate([(255, 128, 0), (0, 200, 255)]):
            for d, thick in [(vis[b][0], 1), (vis[b][1], 2)]:
                cnts, _ = cv2.findContours((d > .1).astype('uint8'), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                cv2.drawContours(img, cnts, -1, col, thick)
        cv2.putText(img, f'frame {t}  thin=before thick=after  orange={sides[0]} yellow={sides[1]}', (8, 20), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 255, 255), 1)
        cv2.imwrite(str(OUT / f'overlay_{t:04d}.jpg'), img)
    if t % 50 == 0: print('frame', t, flush=True)
cap.release()

summary = {'method': 'Same as depth_ablation_red_20260916: rendered MANO depth vs D405 depth on visible skin pixels; single camera-ray translation fit on checkerboard half, held-out on other half. Hands rendered jointly, mutual occlusion excluded.',
           'input': str(PRIOR.relative_to(ROOT)), 'frames': N, 'hands': {}}
seg = [(0, 117), (118, 235), (236, 351)]
for b, s in enumerate(sides):
    rs = rows[s]; sup = [r for r in rs if r['depth_supported']]
    summary['hands'][s] = {
        'depth_supported_frames': len(sup),
        'frames_passing_yesterday_gate': int(sum(r['passes_yesterday_gate'] for r in rs)),
        'fit_pixels': stat([r['fit_pixels'] for r in rs]),
        'fit_mad_mm': stat([r['fit_mad_mm'] for r in rs]),
        'wrist_z_m': stat([r['wrist_z_m'] for r in rs]),
        'heldout_abs_residual_before_mm': stat([r['abs_residual_before_mm'] for r in rs]),
        'heldout_signed_residual_before_mm': stat([r['signed_residual_before_mm'] for r in rs]),
        'heldout_abs_residual_after_mm': stat([r['abs_residual_after_mm'] for r in sup]),
        'fit_shift_z_mm': stat([r['fit_shift_z_mm'] for r in sup]),
        'translation_norm_mm': stat([r['translation_norm_mm'] for r in sup]),
        'frames_shift_over_25mm': int(sum(abs(r['fit_shift_z_mm']) > 25 for r in sup)),
        'shift_z_by_segment_mm': {f'{a}-{c}': stat([r['fit_shift_z_mm'] for r in sup if a <= r['frame'] <= c]) for a, c in seg},
    }
np.savez_compressed(OUT / 'hand_depth_fit.npz', translation_camera=deltas, depth_supported=supported, sides=np.array(sides), K=K)
(OUT / 'frame_metrics.json').write_text(json.dumps(rows, indent=1))
(OUT / 'summary.json').write_text(json.dumps(summary, indent=2))
print(json.dumps(summary, indent=2, ensure_ascii=False))