"""Relocalize the real captures: estimate pinhole intrinsics (Zhang, pp fixed at center, zero distortion),
solve board poses, compose the UMI gripper + camera poses in the mj_arm world frame."""
import json
import numpy as np, cv2
from calib_utils import detect_boards, estimate_focal, solve_pose
from aruco_boards import ARM_BASE, UMI_TOP, UMI_BACK

sf = {k: np.array(v) for k, v in json.load(open("sim_frames.json")).items()}
T_armB_w = sf["T_armB_in_world"]

def angle_deg(R):
    return float(np.degrees(np.arccos(np.clip((np.trace(R) - 1) / 2, -1, 1))))

# ---------- scene capture ----------
img = cv2.imread("captures/scene_view_umi_capture.png")
H, W = img.shape[:2]
g = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
dets = detect_boards(g, [ARM_BASE, UMI_TOP, UMI_BACK])
views = [dets[n] for n in ("arm_base", "umi_top", "umi_back")]
K_s, poses, err = estimate_focal(views, (W, H), f_guess=0.8 * max(W, H))
f_s = K_s[0, 0]
print(f"SCENE: {W}x{H}, estimated f={f_s:.0f}px (fovy={2*np.degrees(np.arctan(H/2/f_s)):.1f}deg), calib reproj {err:.2f}px")
T = {}
for name in ("arm_base", "umi_top", "umi_back"):
    T[name], e = solve_pose(*dets[name], K_s)
    print(f"  {name}: PnP reproj {e:.2f}px, dist {np.linalg.norm(T[name][:3,3]):.3f}m")
# camera in world
T_camCV_w = T_armB_w @ np.linalg.inv(T["arm_base"])
# umi root in world via BOTH boards -> consistency check
roots = {}
for bname, key in [("umi_top", "T_umiTop_in_root"), ("umi_back", "T_umiBack_in_root")]:
    T_UB_root = sf[key]
    roots[bname] = T_camCV_w @ T[bname] @ np.linalg.inv(T_UB_root)
dt = np.linalg.norm(roots["umi_top"][:3, 3] - roots["umi_back"][:3, 3]) * 1000
dR = angle_deg(roots["umi_top"][:3, :3].T @ roots["umi_back"][:3, :3])
print(f"  umi root: top-vs-back consistency = {dt:.1f}mm / {dR:.2f}deg")
T_root_w = roots["umi_top"]
# ---------- wrist capture ----------
wimg = cv2.imread("captures/wrist_cam_capture.png")
Hw, Ww = wimg.shape[:2]
wdets = detect_boards(cv2.cvtColor(wimg, cv2.COLOR_BGR2GRAY), [ARM_BASE])
K_w, _, werr = estimate_focal([wdets["arm_base"]], (Ww, Hw), f_guess=900.0)
f_w = K_w[0, 0]
print(f"WRIST: {Ww}x{Hw}, estimated f={f_w:.0f}px (expected ~914 for the 70deg B0332 lens), calib reproj {werr:.2f}px")
T_wB_cam, we = solve_pose(*wdets["arm_base"], K_w)
T_wcamCV_w = T_armB_w @ np.linalg.inv(T_wB_cam)
print(f"  arm_base PnP reproj {we:.2f}px, wrist cam {np.linalg.norm(T_wB_cam[:3,3]):.3f}m from board")
json.dump({"scene": {"K": K_s.tolist(), "W": W, "H": H, "T_camCV_w": T_camCV_w.tolist(),
                     "T_root_w": T_root_w.tolist()},
           "wrist": {"K": K_w.tolist(), "W": Ww, "H": Hw, "T_camCV_w": T_wcamCV_w.tolist()}},
          open("relocalization.json", "w"), indent=1)
print("saved relocalization.json")