# Written by yams_any4d agent, 2026-07-12. Compare prep-baked scene calib vs the
# per-episode raw calibration_results.json: project GT EE with both, overlay on the
# cached thumb (prep=RED, fresh-raw=GREEN) to see which aligns with the gripper.
import argparse, glob, os, json
import numpy as np, cv2, torch
from custom.any4d.v4_head import V4Head


def scaled_K(K, save_wh, hw):
    (h, w) = hw
    K = K.copy().astype(np.float64)
    K[0] *= w / float(save_wh[0]); K[1] *= h / float(save_wh[1])
    return K


def project(K, T_w2c, xyz):
    pts = xyz.astype(np.float64)
    cam = (T_w2c[:3, :3] @ pts.T).T + T_w2c[:3, 3]
    z = np.clip(cam[:, 2:3], 1e-3, None)
    uvw = (K @ cam.T).T
    return uvw[:, :2] / z


def draw(img, uv, color, r=4):
    for t in range(len(uv)):
        cv2.circle(img, (int(round(uv[t, 0])), int(round(uv[t, 1]))), 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)
    K_prep = head.K_scene.numpy(); swh = head.scene_save_wh.numpy()
    T_prep = head.T_w2c_scene.numpy()

    ep = a.ep
    cj = json.load(open(f'{a.raw}/{ep:04d}/calibration_results.json'))
    K_raw = np.array(cj['scene_camera']['intrinsics_save_res'])
    swh_raw = np.array(cj['scene_camera']['save_wh'])
    ext_c2w = np.array(cj['scene_camera']['extrinsics_right_arm_base'])  # camera-to-world
    T_w2c_raw = np.linalg.inv(ext_c2w)

    print('save_wh prep', swh, ' raw', swh_raw)
    print('T_w2c_prep=\n', np.round(T_prep, 4))
    print('T_w2c_raw (inv of ep%d c2w)=\n' % ep, np.round(T_w2c_raw, 4))
    print('||T_prep - T_w2c_raw|| =', np.abs(T_prep - T_w2c_raw).max())
    print('cam pos (c2w t) prep-implied=', np.round(np.linalg.inv(T_prep)[:3, 3], 4),
          ' raw=', np.round(ext_c2w[:3, 3], 4))

    f = sorted(glob.glob(f'{a.cache_dir}/ep{ep:04d}_w*.pt'))[0]
    d = torch.load(f, map_location='cpu', weights_only=False)
    thumb = d['rgb0_thumb']; H, W = thumb.shape[:2]
    gt_xyz = d['action'].float().numpy()[:, :3]
    hw = (H, W)
    uv_prep = project(scaled_K(K_prep, swh, hw), T_prep, gt_xyz)
    uv_raw = project(scaled_K(K_raw, swh_raw, hw), T_w2c_raw, gt_xyz)
    print('uv_prep[0]', uv_prep[0].round(1), ' uv_raw[0]', uv_raw[0].round(1))
    print('uv_prep[-1]', uv_prep[-1].round(1), ' uv_raw[-1]', uv_raw[-1].round(1))

    img = cv2.cvtColor(thumb, cv2.COLOR_RGB2BGR)
    img = draw(img, uv_prep, (0, 0, 255))     # prep RED
    img = draw(img, uv_raw, (0, 255, 0))      # raw GREEN
    outp = f'{a.out_dir}/kp_prep_vs_raw_ep{ep:04d}.png'
    os.makedirs(a.out_dir, exist_ok=True)
    cv2.imwrite(outp, cv2.resize(img, (W * 3, H * 3), interpolation=cv2.INTER_NEAREST))
    print('saved', outp)
    print('KP_PROBE2 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('--raw', default=os.path.expanduser('~/any4d_work/raw/pickplace_70'))
    p.add_argument('--out_dir', default='/tmp/kpprobe')
    p.add_argument('--ep', type=int, default=0)
    main(p.parse_args())