# Written by yams_any4d agent, 2026-07-12. Test the TCP-offset hypothesis: is the ~10cm
# vertical keypoint error because FK EE is the wrist/flange, not the fingertip toolpoint?
# Reconstruct R from rot6d (first two ROWS), offset EE along each tool axis, overlay.
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, xyz):
    cam = (T[:3, :3] @ xyz.T).T + T[:3, 3]
    z = np.clip(cam[:, 2:3], 1e-3, None)
    return (K @ cam.T).T[:, :2] / z


def R_from_rot6d(r6):
    # rot6d = first two ROWS of R (per repo convention)
    a, b = r6[:3], r6[3:6]
    a = a / (np.linalg.norm(a) + 1e-9)
    b = b - (a @ b) * a; b = b / (np.linalg.norm(b) + 1e-9)
    c = np.cross(a, b)
    return np.stack([a, b, c], axis=0)  # rows are the two given + cross => full R (rows)


def draw(img, uv, color, r=3):
    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):
    ep = a.ep
    cj = json.load(open(f'{a.raw}/{ep:04d}/calibration_results.json'))['scene_camera']
    K = np.array(cj['intrinsics_save_res']); swh = np.array(cj['save_wh'])
    T = np.linalg.inv(np.array(cj['extrinsics_right_arm_base']))
    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]
    act = d['action'].float().numpy()
    xyz = act[:, :3]
    Ks = scaled_K(K, swh, (H, W))

    img = cv2.cvtColor(thumb, cv2.COLOR_RGB2BGR)
    img = draw(img, project(Ks, T, xyz), (0, 255, 0), 4)   # GREEN = EE (flange)
    L = a.offset
    # offset along each row-axis of R (per frame) by +/-L; pick tool approach axis (row2)
    colors = {2: (255, 0, 0), 1: (0, 255, 255), 0: (255, 0, 255)}  # blue=axis2, yellow=axis1, magenta=axis0
    for axis in (0, 1, 2):
        off = np.zeros_like(xyz)
        for t in range(len(xyz)):
            R = R_from_rot6d(act[t, 3:9])
            off[t] = xyz[t] + L * R[axis]
        img = draw(img, project(Ks, T, off), colors[axis], 2)
    outp = f'{a.out_dir}/kp_tcp_ep{ep:04d}_L{int(L*100)}.png'
    os.makedirs(a.out_dir, exist_ok=True)
    cv2.imwrite(outp, cv2.resize(img, (W * 3, H * 3), interpolation=cv2.INTER_NEAREST))
    print(f'GREEN=EE flange, BLUE=+{L}*Rrow2, YELLOW=+{L}*Rrow1, MAGENTA=+{L}*Rrow0')
    print('saved', outp); print('KP_PROBE3 EXIT OK')


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    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)
    p.add_argument('--offset', type=float, default=0.10)
    main(p.parse_args())