"""UMI gripper -> MuJoCo bundle mj_umi/ with REAL ArUco textures: the two marker planes become native box
geoms with 2d textures (MuJoCo maps 2d textures onto the local +Z face of boxes), replacing the plain STL
plates. Textures = the actual print PNGs (id240 top, id241 back). Verification: render + cv2 detect."""
import os, shutil
import numpy as np, trimesh
import xml.etree.ElementTree as ET
import mujoco

root = ET.parse("umi_gripper.urdf").getroot()
for l in root.findall("link"):
    if l.find("inertial") is None:
        i = ET.SubElement(l, "inertial")
        ET.SubElement(i, "mass", {"value": "0.05"})
        ET.SubElement(i, "inertia", dict(ixx="1e-5", ixy="0", ixz="0", iyy="1e-5", iyz="0", izz="1e-5"))
mj = ET.Element("mujoco"); c = ET.SubElement(mj, "compiler")
c.set("discardvisual", "false"); c.set("meshdir", "."); root.insert(0, mj)
model = mujoco.MjModel.from_xml_string(ET.tostring(root, encoding="unicode"))
mujoco.mj_saveLastXML("/tmp/umi_raw.xml", model)

t = ET.parse("/tmp/umi_raw.xml"); r2 = t.getroot()
# drop the plain STL plates (replaced by textured boxes)
for parent in r2.iter():
    for g in list(parent.findall("geom")):
        if "aruco_planes" in (g.get("mesh") or ""): parent.remove(g)
for a in r2.find("asset").findall("mesh"):
    if "aruco_planes" in (a.get("file") or ""): r2.find("asset").remove(a)
# textures + materials
asset = r2.find("asset")
for name, png in [("ar_top", "aruco_umi_top_grid.png"), ("ar_back", "aruco_umi_back_grid.png")]:
    ET.SubElement(asset, "texture", {"name": name, "type": "2d", "file": png})
    ET.SubElement(asset, "material", {"name": name, "texture": name, "texuniform": "false", "texrepeat": "1 1"})
# plate poses from the cube bounds (servo frame mm -> mujoco world: (x, -z, y)/1000)
cb = trimesh.load("gnode_aruco_cube.stl", force="mesh")
lo, hi = cb.bounds; cc = (lo + hi) / 2; sx, sy, sz = hi - lo
wb = r2.find("worldbody")
top_pos = np.array([cc[0], -cc[2], hi[1] + 0.7]) / 1000.0
back_pos = np.array([hi[0] + 0.7, -cc[2], cc[1]]) / 1000.0
s080 = 0.08 / 0.3                                              # exo plane.obj is 0.3m, faces +Z, has UVs
ET.SubElement(asset, "mesh", {"name": "aruco_plane", "file": "aruco_plane.obj",
                              "scale": f"{s080:.5f} {s080:.5f} 0.001", "inertia": "shell"})
ET.SubElement(wb, "geom", {"type": "mesh", "mesh": "aruco_plane", "material": "ar_top",
                           "pos": " ".join(f"{v:.4f}" for v in top_pos), "contype": "0", "conaffinity": "0"})
ET.SubElement(wb, "geom", {"type": "mesh", "mesh": "aruco_plane", "material": "ar_back",
                           "pos": " ".join(f"{v:.4f}" for v in back_pos), "euler": "0 1.5708 0", "contype": "0", "conaffinity": "0"})
ET.SubElement(wb, "light", {"pos": "0.3 0.3 0.6", "dir": "-0.4 -0.4 -1"})
vis = ET.SubElement(r2, "visual"); ET.SubElement(vis, "global", {"offwidth": "1600", "offheight": "1200"})
ET.SubElement(vis, "headlight", {"ambient": "0.5 0.5 0.5"})
act = ET.SubElement(r2, "actuator")
ET.SubElement(act, "position", {"joint": "jaw", "kp": "5", "kv": "0.4", "ctrlrange": "-0.1 0.9", "name": "jaw"})
os.makedirs("mj_umi", exist_ok=True)
ET.indent(t, "  "); t.write("mj_umi/umi_gripper.xml", encoding="unicode")
for a in r2.iter("mesh"):
    f = a.get("file")
    if f and os.path.exists(f): shutil.copy(f, "mj_umi/" + os.path.basename(f))
for f in ("aruco_umi_top_grid.png", "aruco_umi_back_grid.png", "aruco_plane.obj"):
    shutil.copy(f, "mj_umi/" + f)
print("wrote mj_umi/umi_gripper.xml (+meshes+textures)")