"""Arm (detach base) -> MuJoCo bundle mj_arm/ with the base fiducial board textured (ids 100-108 GridBoard),
position actuators on all joints (old make_mjcf convention kp40/kv2), floor + light. Detection-verified."""
import os, shutil
import numpy as np, trimesh
import xml.etree.ElementTree as ET
import mujoco

root = ET.parse("two_servo_detach.urdf").getroot()
lift = float(next(j for j in root.findall("joint") if j.get("name") == "world_to_servo_0").find("origin").get("xyz").split()[2])
for l in root.findall("link"):
    if l.find("inertial") is None:
        i = ET.SubElement(l, "inertial")
        ET.SubElement(i, "mass", {"value": "0.1"})
        ET.SubElement(i, "inertia", dict(ixx="1e-4", ixy="0", ixz="0", iyy="1e-4", iyz="0", izz="1e-4"))
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/arm_raw.xml", model)
t = ET.parse("/tmp/arm_raw.xml"); r2 = t.getroot()
# drop the dark viz plate (replaced by the textured plane)
for parent in r2.iter():
    for g in list(parent.findall("geom")):
        if "arm_marker_plate" in (g.get("mesh") or ""): parent.remove(g)
for a in r2.find("asset").findall("mesh"):
    if "arm_marker_plate" in (a.get("file") or ""): r2.find("asset").remove(a)
asset = r2.find("asset")
ET.SubElement(asset, "texture", {"name": "armboard", "type": "2d", "file": "aruco_arm_base_grid.png"})
ET.SubElement(asset, "material", {"name": "armboard", "texture": "armboard"})
s160 = 0.160 / 0.3
ET.SubElement(asset, "mesh", {"name": "aruco_plane160", "file": "aruco_plane.obj",
                              "scale": f"{s160:.5f} {s160:.5f} 0.001", "inertia": "shell"})
ET.SubElement(asset, "texture", {"name": "sky", "type": "skybox", "builtin": "gradient",
                                 "rgb1": "0.35 0.45 0.55", "rgb2": "0.05 0.05 0.1", "width": "256", "height": "256"})
ET.SubElement(asset, "texture", {"name": "floor", "type": "2d", "builtin": "checker", "rgb1": "0.2 0.25 0.3",
                                 "rgb2": "0.15 0.18 0.22", "width": "300", "height": "300"})
ET.SubElement(asset, "material", {"name": "floor", "texture": "floor", "texuniform": "true", "texrepeat": "4 4"})
wb = r2.find("worldbody")
ET.SubElement(wb, "geom", {"type": "plane", "size": "1.2 1.2 0.05", "material": "floor"})
# board top plane (servo frame (x,-z,y) + lift)
mp = trimesh.load("gnode_arm_marker_plate.stl", force="mesh")
lo, hi = mp.bounds
cc = (lo + hi) / 2
pos = np.array([cc[0], -cc[2], lo[1] + 0.1 + lift * 1000.0]) / 1000.0
pos[2] = (lo[1] + 0.1) / 1000.0 + lift
ET.SubElement(wb, "geom", {"type": "mesh", "mesh": "aruco_plane160", "material": "armboard",
                           "pos": " ".join(f"{v:.4f}" for v in pos), "contype": "0", "conaffinity": "0"})
ET.SubElement(wb, "light", {"pos": "0.4 0.4 1.0", "dir": "-0.3 -0.3 -1", "directional": "false"})
vis = ET.SubElement(r2, "visual"); ET.SubElement(vis, "global", {"offwidth": "1600", "offheight": "1200"})
ET.SubElement(vis, "headlight", {"ambient": "0.45 0.45 0.45"})
act = ET.SubElement(r2, "actuator")
for j in r2.iter("joint"):
    if j.get("type") != "free" and j.get("name", "").startswith("output"):
        ET.SubElement(act, "position", {"joint": j.get("name"), "kp": "40", "kv": "2",
                                        "ctrlrange": "-1.74 1.74", "name": j.get("name")})
os.makedirs("mj_arm", exist_ok=True)
ET.indent(t, "  "); t.write("mj_arm/arm.xml", encoding="unicode")
for a in r2.iter("mesh"):
    f = a.get("file")
    if f and os.path.exists(f): shutil.copy(f, "mj_arm/" + os.path.basename(f))
shutil.copy("aruco_arm_base_grid.png", "mj_arm/aruco_arm_base_grid.png")
shutil.copy("aruco_plane.obj", "mj_arm/aruco_plane.obj")
print(f"wrote mj_arm/arm.xml (lift={lift:.3f})")