"""Compose the replay scene: mj_arm world + UMI gripper model attached at its ESTIMATED pose + cameras at
the ESTIMATED capture poses. Manual XML merge with u_ prefixing (deterministic, no MjSpec API roulette)."""
import json, os, shutil
import numpy as np
import xml.etree.ElementTree as ET
import trimesh.transformations as tt
from calib_utils import cam_cv_to_mj

rl = json.load(open("relocalization.json"))
os.makedirs("mj_scene", exist_ok=True)
for src in ("mj_arm", "mj_umi"):
    for f in os.listdir(src):
        if not f.endswith(".xml"): shutil.copy(f"{src}/{f}", f"mj_scene/{f}")

arm = ET.parse("mj_arm/arm.xml"); ar = arm.getroot()
umi = ET.parse("mj_umi/umi_gripper.xml"); ur = umi.getroot()
# --- prefix all umi asset names + references ---
ren = {}
for a in ur.find("asset"):
    n = a.get("name") or (os.path.splitext(os.path.basename(a.get("file")))[0] if a.get("file") else None)
    if a.get("name") is None and n: a.set("name", n)
    if a.get("name"):
        ren[a.get("name")] = "u_" + a.get("name")
        a.set("name", "u_" + a.get("name"))
for a in ur.find("asset"):
    if a.tag == "material" and a.get("texture") in ren: a.set("texture", ren[a.get("texture")])
for e in ur.find("worldbody").iter():
    for attr in ("mesh", "material"):
        if e.get(attr) in ren: e.set(attr, ren[e.get(attr)])
    if e.get("name"): e.set("name", "u_" + e.get("name"))
# --- wrapper body at the estimated umi pose ---
T = np.array(rl["scene"]["T_root_w"])
q = tt.quaternion_from_matrix(T)                                  # wxyz
wrap = ET.Element("body", {"name": "u_root", "pos": " ".join(f"{v:.5f}" for v in T[:3, 3]),
                           "quat": " ".join(f"{v:.6f}" for v in q)})
for child in list(ur.find("worldbody")):
    if child.tag == "light": continue
    wrap.append(child)
# --- merge into arm scene ---
a_asset = ar.find("asset")
for a in ur.find("asset"): a_asset.append(a)
ar.find("worldbody").append(wrap)
# --- estimated cameras ---
wb = ar.find("worldbody")
for name, blob in [("scene_est", rl["scene"]), ("wrist_est", rl["wrist"])]:
    pos, xyax = cam_cv_to_mj(np.array(blob["T_camCV_w"]))
    fovy = 2 * np.degrees(np.arctan(blob["H"] / 2.0 / blob["K"][1][1]))
    ET.SubElement(wb, "camera", {"name": name, "pos": " ".join(f"{v:.5f}" for v in pos),
                                 "xyaxes": xyax, "fovy": f"{fovy:.3f}"})
for gl in ar.find("visual").iter("global"):
    gl.set("offwidth", "2200"); gl.set("offheight", "2200")
ET.indent(arm, "  "); arm.write("mj_scene/replay.xml", encoding="unicode")
print("wrote mj_scene/replay.xml")