"""Extract T_umi_board→TCP by replaying the viewer's Layer setup exactly.

Uses `Layer.T_board_w` (the transform the viewer computes via PnP on the
rendered ArUco board — the SAME transform that made the big-TCP render
look correct). Composes: tcp_in_umi_board = inv(T_board_w) @ tcp_in_layer_model.

Saves the resulting T_board_tcp to tcp_offset.json for the dataset loader.
Then validates by projecting the derived TCP into the wrist camera across
10 sample frames and comparing with the pixel location of the green sphere
in the previous render — should match visually.
"""
from __future__ import annotations

import json, os, sys
from pathlib import Path
import numpy as np

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

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")
AR_VIEW = Path("/data/cameron/claude_feetech_controller/ar_view")

# Use the SAME patched UMI model that produced the correct-looking render,
# so any coordinate convention baked into the mesh + euler stays consistent.
UMI_MODEL = Path("/tmp/mj_umi_v2_bigtcp/umi_gripper_v2.xml")


def main():
    sys.path.insert(0, str(AR_VIEW))
    import mujoco
    from render_core import Layer
    from calib_board import CALIB_BOARD  # noqa (unused but keeps sys.path warm)

    # Re-import UMI_V2 board spec from the viewer's module scope.
    import render_dataset_frame as R
    UMI_V2 = R.UMI_V2   # dict from render_dataset_frame

    # Build the SAME Layer the viewer builds for the UMI.
    layer = Layer("umi", str(UMI_MODEL), UMI_V2, "ar_back", lambda jn: None)
    T_board_w = np.asarray(layer.T_board_w, dtype=np.float64)   # (4,4) board→layer_model

    tcp_id = layer.model.site("tcp").id
    tcp_in_layer = np.asarray(layer.data.site_xpos[tcp_id], dtype=np.float64)
    print(f"tcp in layer model: {tcp_in_layer}  |= {np.linalg.norm(tcp_in_layer):.4f}m")
    print(f"T_board_w (board→layer):\n{T_board_w}")

    # Invert to get layer→board, then apply to tcp_in_layer.
    T_layer_from_board = T_board_w
    T_board_from_layer = np.linalg.inv(T_layer_from_board)
    tcp_in_board_h = T_board_from_layer @ np.append(tcp_in_layer, 1.0)
    tcp_in_board = tcp_in_board_h[:3]
    print(f"\ntcp in OpenCV umi-board frame: {tcp_in_board}  |= {np.linalg.norm(tcp_in_board):.4f}m")

    # Full 4x4 T_board→tcp with the site's rotation carried through
    tcp_xmat_in_layer = np.asarray(layer.data.site_xmat[tcp_id], dtype=np.float64).reshape(3, 3)
    T_layer_from_tcp = np.eye(4)
    T_layer_from_tcp[:3, :3] = tcp_xmat_in_layer
    T_layer_from_tcp[:3, 3] = tcp_in_layer
    T_board_from_tcp = T_board_from_layer @ T_layer_from_tcp
    print(f"\nT_board_from_tcp (constant, for use in dataset):\n{T_board_from_tcp}")

    OUT_JSON.write_text(json.dumps({
        "T_board_tcp": T_board_from_tcp.tolist(),
        "tcp_in_board_translation": tcp_in_board.tolist(),
        "derivation": "layer.T_board_w from render_core.measure_board_world, "
                       "then inv() applied to data.site('tcp').xpos",
        "source_xml": str(UMI_MODEL),
    }, indent=2))
    print(f"\nsaved → {OUT_JSON}")

    # ── Validate: project TCP into wrist across 10 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 = 500, 312
    print(f"\n── project TCP into wrist cam (500×{H_img}) using CURRENT dataset composition ──")
    print(f"     using T_world_wc = umi_pose @ inv(T_wc_umi) [as in dataset_arview.py]")
    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)
        tcp_world = (umi @ T_board_from_tcp @ np.array([0, 0, 0, 1.0]))[:3]
        T_w2c = np.linalg.inv(umi @ np.linalg.inv(T_wc_umi))
        cam = T_w2c[:3, :3] @ tcp_world + T_w2c[: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]
            print(f"  f={f}: z={z:+.4f}m  ({u:+.0f},{v:+.0f})  in_frame={0<=u<W and 0<=v<H_img}")
        else:
            print(f"  f={f}: z={z:+.4f}m  BEHIND camera")

    layer.close()


if __name__ == "__main__":
    main()