"""Derive the constant T_board→tcp offset from the UMI MuJoCo model
(with the tcp site cad added 2026-07-14) and validate that it puts
the TCP inside the wrist camera's FOV on real data.

Outputs:
  1. Numeric T_board→tcp printed to stdout
  2. Saves to /data/cameron/repos/smithbot_v4/preprocess/tcp_offset.json
  3. Per-frame projection depths for TCP into wrist cam (should be
     positive z, in-frame u/v)
"""
from __future__ import annotations
import json, os
from pathlib import Path
import numpy as np

os.environ.setdefault("MUJOCO_GL", "osmesa")

MODEL_XML = "/data/cameron/cad_recovery/mj_umi_v2/umi_gripper_v2.xml"
DATASET = Path("/data/cameron/puget_ar_datasets/testing_cube_in_bowl")
CAL = json.loads(Path("/data/cameron/claude_feetech_controller/ar_view/wrist_umi_calibration.json").read_text())
OUT_JSON = Path("/data/cameron/repos/smithbot_v4/preprocess/tcp_offset.json")


def _quat_to_mat(q_wxyz):
    """MuJoCo quat is (w, x, y, z)."""
    import mujoco
    m = np.zeros(9, dtype=np.float64)
    mujoco.mju_quat2Mat(m, np.asarray(q_wxyz, dtype=np.float64))
    return m.reshape(3, 3)


def main():
    import mujoco
    model = mujoco.MjModel.from_xml_path(MODEL_XML)
    data = mujoco.MjData(model)
    mujoco.mj_forward(model, data)

    # TCP site — this is what cad added.
    tcp_id = model.site("tcp").id
    tcp_xpos = np.asarray(data.site_xpos[tcp_id], dtype=np.float64)   # world pos
    tcp_xmat = np.asarray(data.site_xmat[tcp_id], dtype=np.float64).reshape(3, 3)
    T_world_tcp = np.eye(4)
    T_world_tcp[:3, :3] = tcp_xmat
    T_world_tcp[:3, 3] = tcp_xpos

    # ArUco board geom — no name in XML, so find it by mesh reference.
    # The aruco_plane geom is the only one whose mesh id matches "aruco_plane".
    aruco_mesh_id = model.mesh("aruco_plane").id
    aruco_geom_id = None
    for gi in range(model.ngeom):
        if model.geom_type[gi] == mujoco.mjtGeom.mjGEOM_MESH and model.geom_dataid[gi] == aruco_mesh_id:
            aruco_geom_id = gi
            break
    assert aruco_geom_id is not None, "aruco_plane geom not found in UMI XML"
    board_xpos = np.asarray(data.geom_xpos[aruco_geom_id], dtype=np.float64)
    board_xmat = np.asarray(data.geom_xmat[aruco_geom_id], dtype=np.float64).reshape(3, 3)
    T_world_board = np.eye(4)
    T_world_board[:3, :3] = board_xmat
    T_world_board[:3, 3] = board_xpos
    T_board_world = np.linalg.inv(T_world_board)

    # Constant offset from board frame to TCP frame.
    T_board_tcp = T_board_world @ T_world_tcp

    print("── T_board_tcp (offset from ArUco board center to TCP fingertip) ──")
    print(T_board_tcp)
    print(f"  translation (m): ({T_board_tcp[0, 3]:+.4f}, {T_board_tcp[1, 3]:+.4f}, {T_board_tcp[2, 3]:+.4f})")
    dist = float(np.linalg.norm(T_board_tcp[:3, 3]))
    print(f"  |offset| = {dist:.4f} m ({dist*1000:.1f} mm)")

    OUT_JSON.parent.mkdir(parents=True, exist_ok=True)
    OUT_JSON.write_text(json.dumps({
        "T_board_tcp": T_board_tcp.tolist(),
        "source_xml": MODEL_XML,
        "site_pos_in_model": tcp_xpos.tolist(),
        "board_pos_in_model": board_xpos.tolist(),
        "note": "constant rigid offset — the UMI board frame is what OpenCV "
                "PnP returns for umi_pose in state.npz; TCP is at "
                "T_world_tcp = umi_pose @ T_board_tcp."
    }, indent=2))
    print(f"\nsaved → {OUT_JSON}")

    # ── VALIDATION: project TCP into wrist camera across our sample frames ──
    T_wc_umi = np.asarray(CAL["T_wrist_cam_umi"], dtype=np.float64)
    meta = json.loads((DATASET/"meta.json").read_text())
    K_wrist = np.asarray(meta["wrist"]["K_save"], dtype=np.float64)
    W, H_img = int(meta["wrist"]["save_width"]), 281
    print(f"\n── validation on 10 sample frames (should be z>0, in-frame) ──")
    print("f  | tcp_z_in_wrist_cam | (u, v) | in_frame?")
    for f in [91, 95, 99, 103, 107, 111, 115, 119, 123, 127]:
        st = np.load(DATASET/"state"/f"{f:06d}.npz")
        umi = np.asarray(st["umi_pose"], dtype=np.float64)
        T_world_tcp = umi @ T_board_tcp
        tcp_world = T_world_tcp[:3, 3]
        # wrist cam pose per current dataset code
        T_world_wrist = umi @ np.linalg.inv(T_wc_umi)
        T_w2c_wrist = np.linalg.inv(T_world_wrist)
        cam = T_w2c_wrist[:3, :3] @ tcp_world + T_w2c_wrist[:3, 3]
        z = cam[2]
        if z > 1e-3:
            u = K_wrist[0, 0]*cam[0]/z + K_wrist[0, 2]
            v = K_wrist[1, 1]*cam[1]/z + K_wrist[1, 2]
            in_frame = 0 <= u < W and 0 <= v < H_img
            print(f"  {f}: z={z:+.4f}m  ({u:+.0f}, {v:+.0f})  in={in_frame}")
        else:
            print(f"  {f}: z={z:+.4f}m  (BEHIND camera)")


if __name__ == "__main__":
    main()