# Written by yams_any4d agent, 2026-07-12. Probe: reproduce the GT-keypoint vs
# gripper-toolpoint misalignment. Projects GT EE (action r_xyz) onto the cached scene
# thumb via the prep K_scene/T_w2c_scene, overlays, and prints uv + principal-point math.
import argparse, glob, os
import numpy as np, cv2, torch
from custom.any4d.v4_head import V4Head


def draw(bg, uv, color, r=4):
    img = bg.copy()
    for t in range(len(uv)):
        x, y = int(round(uv[t, 0])), int(round(uv[t, 1]))
        cv2.circle(img, (x, y), r, color, -1, cv2.LINE_AA)
    return img


def main(a):
    head = V4Head(prep_path=a.prep, use_wrist=False, temporal_mode='stack_mlp', pred_size=128)
    files = sorted(glob.glob(f'{a.cache_dir}/ep*.pt'))
    os.makedirs(a.out_dir, exist_ok=True)
    K = head.K_scene.numpy(); swh = head.scene_save_wh.numpy()
    print(f'K_scene fx={K[0,0]:.1f} fy={K[1,1]:.1f} cx={K[0,2]:.2f} cy={K[1,2]:.2f}  save_wh={swh}')
    print(f'save_wh aspect w/h = {swh[0]/swh[1]:.4f}')
    for fi, f in enumerate(files[:a.n]):
        d = torch.load(f, map_location='cpu', weights_only=False)
        thumb = d['rgb0_thumb']            # (176,320,3) uint8 RGB
        H, W = thumb.shape[:2]
        act = d['action'].float()          # (41,20)
        gt_xyz = act[:, :3]                 # r_xyz world (right_arm_base)
        hw = (H, W)
        uv = head.project_world_to_pix(gt_xyz, hw).numpy()   # (41,2)
        Ks = head._scaled_K(head.K_scene, head.scene_save_wh, hw, 'cpu', torch.float32).numpy()
        bg = cv2.cvtColor(thumb, cv2.COLOR_RGB2BGR)
        img = draw(bg, uv, (0, 0, 255))    # GT red
        cv2.circle(img, (int(round(Ks[0,2])), int(round(Ks[1,2]))), 6, (0,255,0), 1, cv2.LINE_AA)  # scaled cx,cy green
        outp = f'{a.out_dir}/kpprobe_{fi:02d}_{d.get("ep","?")}.png'
        cv2.imwrite(outp, cv2.resize(img, (W*3, H*3), interpolation=cv2.INTER_NEAREST))
        if fi < 3:
            print(f'\n{os.path.basename(f)} img {W}x{H}  scaledK cx={Ks[0,2]:.1f} cy={Ks[1,2]:.1f}')
            print(f'  gt_xyz[0]={gt_xyz[0].numpy().round(3)}  gt_xyz[-1]={gt_xyz[-1].numpy().round(3)}')
            print(f'  uv[0]={uv[0].round(1)}  uv[mid]={uv[len(uv)//2].round(1)}  uv[-1]={uv[-1].round(1)}')
            print(f'  uv range: u[{uv[:,0].min():.0f},{uv[:,0].max():.0f}] v[{uv[:,1].min():.0f},{uv[:,1].max():.0f}]')
    print('\nsaved overlays to', a.out_dir)
    print('KP_PROBE EXIT OK')


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--prep', default=os.path.expanduser('~/any4d_work/v4head_prep.npz'))
    p.add_argument('--cache_dir', default=os.path.expanduser('~/any4d_work/feat_cache_gen'))
    p.add_argument('--out_dir', default='/tmp/kpprobe')
    p.add_argument('--n', type=int, default=6)
    main(p.parse_args())