"""Empirically measure board frames in the sims (absorbs ALL texture/UV/frame conventions):
- T_armB_in_world  : arm base GridBoard frame in mj_arm world
- T_umiTop_in_root, T_umiBack_in_root : UMI board frames in the umi model root frame
Loop-closure check: recovered camera pose must match the ground-truth render camera."""
import json
import numpy as np, cv2, mujoco
import xml.etree.ElementTree as ET
from calib_utils import detect_boards, solve_pose, K_from_fovy, cam_mj_to_cv, cam_cv_to_mj
from aruco_boards import ARM_BASE, UMI_TOP, UMI_BACK

def measure(xml_path, cam_pos, cam_xyaxes, fovy, specs, W=1400, H=1400):
    t = ET.parse(xml_path); r = t.getroot()
    wb = r.find("worldbody")
    ET.SubElement(wb, "camera", {"name": "calib", "pos": " ".join(map(str, cam_pos)),
                                 "xyaxes": cam_xyaxes, "fovy": str(fovy)})
    for g in list(r.find("visual")):
        if g.tag == "global": g.set("offwidth", "2000"); g.set("offheight", "2000")
    tmp = xml_path.replace(".xml", "_calib.xml"); t.write(tmp)
    m = mujoco.MjModel.from_xml_path(tmp); d = mujoco.MjData(m); mujoco.mj_forward(m, d)
    i = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_CAMERA, "calib")
    T_camCV_w = cam_mj_to_cv(d.cam_xpos[i], d.cam_xmat[i])
    with mujoco.Renderer(m, H, W) as ren:
        ren.update_scene(d, "calib"); img = ren.render()
    K = K_from_fovy(fovy, W, H)
    dets = detect_boards(cv2.cvtColor(img, cv2.COLOR_RGB2GRAY), specs)
    out = {}
    for name, (obj, imgp) in dets.items():
        T_B_cam, err = solve_pose(obj, imgp, K)
        T_B_w = T_camCV_w @ T_B_cam
        out[name] = (T_B_w, err)
        print(f"  {name}: reproj {err:.2f}px")
    return out

res = {}
print("mj_arm (board in world):")
arm = measure("mj_arm/arm.xml", [0.15, 0.35, 0.55], "-1 0.6 0 -0.35 -0.6 0.75", 55, [ARM_BASE])
res["T_armB_in_world"] = arm["arm_base"][0].tolist()
print("mj_umi (boards in model root frame):")
umi = measure("mj_umi/umi_gripper.xml", [0.35, 0.25, 0.35], "-0.6 0.85 0 -0.5 -0.35 0.85", 55, [UMI_TOP, UMI_BACK])
for k, key in [("umi_top", "T_umiTop_in_root"), ("umi_back", "T_umiBack_in_root")]:
    if k in umi: res[key] = umi[k][0].tolist()
json.dump(res, open("sim_frames.json", "w"), indent=1)
print("saved sim_frames.json:", sorted(res))