"""Keypoint projection viz — TCP-target version.

Rerun of preprocess/kp_projection_diff_viz.py but with future keypoints
= TCP fingertip positions (umi_pose @ T_board_tcp @ origin), not the
ArUco board center. Uses the same wrist-cam composition the dataset
currently uses (umi_pose @ inv(T_wrist_cam_umi)) — which is what
Cameron confirmed correct via the big-TCP-sphere render.

Output: /data/cameron/repos/smithbot_v4/preprocess/kp_tcp_viz/
        f{frame:06d}.jpg  — 2-panel: scene | wrist  (both with TCP keypoints)
        all_frames_stacked.jpg
"""
from __future__ import annotations
import json
from pathlib import Path
import cv2
import numpy as np


DATASET = Path("/data/cameron/puget_ar_datasets/testing_cube_in_bowl")
EPISODES_JSON = Path("/data/cameron/ar_renders/testing_cube_in_bowl/frames/episodes.json")
CAL = Path("/data/cameron/claude_feetech_controller/ar_view/wrist_umi_calibration.json")
TCP_OFF = Path("/data/cameron/repos/smithbot_v4/preprocess/tcp_offset.json")
OUT_DIR = Path("/data/cameron/repos/smithbot_v4/preprocess/kp_tcp_viz")

N_WINDOW = 32
N_SAMPLE_FRAMES = 10
EPISODE_INDEX = 1


def project(xyz_world, K, T_w2c):
    homo = np.concatenate([xyz_world, np.ones((len(xyz_world), 1))], axis=1)
    cam = (T_w2c @ homo.T).T[:, :3]
    z = cam[:, 2:3]
    front = z[:, 0] > 1e-3
    z_safe = np.where(z <= 1e-3, 1e-3, z)
    pix = (K @ (cam / z_safe).T).T[:, :2]
    return pix, front


def draw_keypoints(img, pix, front, color_start=(255, 255, 0),
                    color_end=(0, 128, 255), radius=5):
    h, w = img.shape[:2]
    prev = None
    for i, (x, y) in enumerate(pix):
        if not front[i] or not (0 <= x < w and 0 <= y < h):
            prev = None
            continue
        alpha = i / max(1, len(pix) - 1)
        color = tuple(int(color_start[c] * (1 - alpha) + color_end[c] * alpha)
                       for c in range(3))
        cv2.circle(img, (int(x), int(y)), radius, color, -1, cv2.LINE_AA)
        if prev is not None:
            cv2.line(img, prev, (int(x), int(y)), color, 2, cv2.LINE_AA)
        prev = (int(x), int(y))
    return img


def label(img, txt, color=(255, 255, 255)):
    cv2.rectangle(img, (0, 0), (img.shape[1], 24), (0, 0, 0), -1)
    cv2.putText(img, txt, (6, 17), cv2.FONT_HERSHEY_SIMPLEX,
                 0.55, color, 1, cv2.LINE_AA)
    return img


def main():
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    meta = json.loads((DATASET/"meta.json").read_text())
    K_scene = np.asarray(meta["scene"]["K_save"], dtype=np.float64)
    K_wrist = np.asarray(meta["wrist"]["K_save"], dtype=np.float64)
    T_wc_umi = np.asarray(json.loads(CAL.read_text())["T_wrist_cam_umi"], dtype=np.float64)
    T_board_tcp = np.asarray(json.loads(TCP_OFF.read_text())["T_board_tcp"], dtype=np.float64)
    print(f"TCP offset: |{T_board_tcp[:3, 3]}| = {np.linalg.norm(T_board_tcp[:3, 3]):.3f}m")

    eps = json.loads(EPISODES_JSON.read_text())["episodes"]
    ep = eps[EPISODE_INDEX]
    e_start, e_end = int(ep["start"]), int(ep["end"])
    print(f"episode {ep['id']}  frames [{e_start}..{e_end}]  len={e_end-e_start+1}")

    def _load(i):
        d = np.load(DATASET/"state"/f"{i:06d}.npz")
        return {
            "scene_cam_pose": np.asarray(d["scene_cam_pose"], dtype=np.float64),
            "umi_pose":       np.asarray(d["umi_pose"],       dtype=np.float64),
        }
    total = int(meta["num_frames"])
    state = [_load(i) for i in range(total)]

    frame_ids = np.linspace(e_start, e_end - 1, N_SAMPLE_FRAMES).round().astype(int).tolist()
    for f in frame_ids:
        future_ids = np.clip(np.arange(1, N_WINDOW + 1) + f, 0, e_end)
        # TCP world = umi_pose @ T_board_tcp @ origin
        tcp_world = np.stack([
            (state[i]["umi_pose"] @ T_board_tcp @ np.array([0, 0, 0, 1.0]))[:3]
            for i in future_ids
        ], axis=0)

        scene_bgr = cv2.imread(str(DATASET/"scene"/f"{f:06d}.jpg"))
        wrist_bgr = cv2.imread(str(DATASET/"wrist"/f"{f:06d}.jpg"))
        umi_pose = state[f]["umi_pose"]
        T_w2c_scene = np.linalg.inv(state[f]["scene_cam_pose"])
        # Wrist cam pose per current dataset composition (with inv, per Cameron's verify)
        T_world_wrist = umi_pose @ T_wc_umi   # ALT — matches big-TCP render
        T_w2c_wrist = np.linalg.inv(T_world_wrist)

        p_scene, f_scene = project(tcp_world, K_scene, T_w2c_scene)
        p_wrist, f_wrist = project(tcp_world, K_wrist, T_w2c_wrist)

        panel_scene = draw_keypoints(scene_bgr.copy(), p_scene, f_scene)
        panel_wrist = draw_keypoints(wrist_bgr.copy(), p_wrist, f_wrist)
        panel_scene = label(panel_scene, f"scene  f={f}  (TCP fingertip targets)", (255, 255, 255))
        panel_wrist = label(panel_wrist, f"wrist  f={f}  (TCP fingertip targets)", (255, 255, 255))

        H = max(panel_scene.shape[0], panel_wrist.shape[0])
        def pad(im):
            if im.shape[0] == H: return im
            return np.vstack([im, np.zeros((H - im.shape[0], im.shape[1], 3), dtype=np.uint8)])
        comp = np.hstack([pad(panel_scene), pad(panel_wrist)])
        cv2.imwrite(str(OUT_DIR / f"f{f:06d}.jpg"), comp, [cv2.IMWRITE_JPEG_QUALITY, 92])

        # Count in-frame future keypoints per view for a sanity print
        n_scene = sum(1 for p, fr in zip(p_scene, f_scene) if fr and 0 <= p[0] < scene_bgr.shape[1] and 0 <= p[1] < scene_bgr.shape[0])
        n_wrist = sum(1 for p, fr in zip(p_wrist, f_wrist) if fr and 0 <= p[0] < wrist_bgr.shape[1] and 0 <= p[1] < wrist_bgr.shape[0])
        print(f"  f={f}: {n_scene}/32 in scene, {n_wrist}/32 in wrist")

    stack = np.vstack([cv2.imread(str(OUT_DIR/f"f{f:06d}.jpg")) for f in frame_ids])
    cv2.imwrite(str(OUT_DIR / "all_frames_stacked.jpg"), stack, [cv2.IMWRITE_JPEG_QUALITY, 92])
    print(f"\nwrote → {OUT_DIR}/")
    print(f"browse: https://omidlab.net/browse/repos/smithbot_v4/preprocess/kp_tcp_viz/")


if __name__ == "__main__":
    main()