"""arm_with_umi_v2.urdf -> MuJoCo bundle mj_arm_v2/: full arm + arm-variant gripper, fused small-board
base (textured 80mm ids-100-108 board), position actuators on all joints, wrist cam on the gripper."""
import os, shutil
import numpy as np, trimesh
import xml.etree.ElementTree as ET
import mujoco

root = ET.parse("arm_with_umi_v2.urdf").getroot()
# swap base visual to the fused small-board part
for l in root.findall("link"):
    if l.get("name") == "servo_0":
        for v in l.findall("visual"):
            m = v.find("geometry/mesh")
            if m is not None and "base_holder" in m.get("filename"):
                m.set("filename", "base_fused_smallboard.stl")
for l in root.findall("link"):
    if l.find("inertial") is None:
        i = ET.SubElement(l, "inertial")
        ET.SubElement(i, "mass", {"value": "0.08"})
        ET.SubElement(i, "inertia", dict(ixx="5e-5", ixy="0", ixz="0", iyy="5e-5", iyz="0", izz="5e-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/armv2_raw.xml", model)
t = ET.parse("/tmp/armv2_raw.xml"); r2 = t.getroot()
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"})
s = 0.08 / 0.3
ET.SubElement(asset, "mesh", {"name": "aruco_plane80", "file": "aruco_plane.obj",
                              "scale": f"{s:.5f} {s:.5f} 0.001", "inertia": "shell"})
ET.SubElement(asset, "texture", {"name": "floor", "type": "2d", "builtin": "checker", "rgb1": "0.22 0.27 0.32",
                                 "rgb2": "0.16 0.2 0.24", "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: fused base board x -2..78, z -180..-100, top y=-28 (servo_0 frame); world=(x,-z,y+lift)
lift_j = next(j for j in ET.parse("arm_with_umi_v2.urdf").getroot().findall("joint") if j.get("name") == "world_to_servo_0")
lift = float(lift_j.find("origin").get("xyz").split()[2])
pos = [0.038, 0.140, (-28.0) / 1000.0 + lift + 0.0007]
ET.SubElement(wb, "geom", {"type": "mesh", "mesh": "aruco_plane80", "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"})
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"):
    n = j.get("name", "")
    if n.startswith("output"):
        ET.SubElement(act, "position", {"joint": n, "kp": "40", "kv": "2", "ctrlrange": "-1.74 1.74", "name": n})
    elif n == "gripper_jaw":
        ET.SubElement(act, "position", {"joint": n, "kp": "5", "kv": "0.4", "ctrlrange": "-0.9 0.1", "name": n})
# wrist cam on the gripper base body (pose derived from camera body geometry in that link frame)
body_m = trimesh.load("gnode_v2_camera_body.stl", force="mesh")
areas, normals = body_m.area_faces, body_m.face_normals
key = np.round(normals, 2)
uniq, inv = np.unique(key, axis=0, return_inverse=True)
sums = np.zeros(len(uniq))
for i, a in zip(inv, areas): sums[i] += a
n = uniq[np.argmax(sums)]; n /= np.linalg.norm(n)
mid = body_m.vertices.mean(0)
mount_c = trimesh.load("gnode_v2_camera_mount.stl", force="mesh").vertices.mean(0)
if n @ (mount_c - mid) > 0: n = -n
tip = body_m.vertices[np.argmax(body_m.vertices @ n)]
cpos = (mid + n * ((tip - mid) @ n + 2.0)) / 1000.0
up_l = np.array([0, -1.0, 0])                       # device up in the gripper link frame
x_l = np.cross(up_l, -n); x_l /= np.linalg.norm(x_l)
y_l = np.cross(-n, x_l)
cpos = cpos + np.array([-0.0464, 0.0, 0.0])   # gripper base is welded into yoke_5 (fixed joint merge)
gbody = next(b for b in r2.iter("body") if b.get("name") == "yoke_5")
ET.SubElement(gbody, "camera", {"name": "wrist_cam", "pos": " ".join(f"{v:.4f}" for v in cpos),
                                "xyaxes": " ".join(f"{v:.4f}" for v in np.concatenate([x_l, y_l])), "fovy": "47"})
os.makedirs("mj_arm_v2", exist_ok=True)
ET.indent(t, "  "); t.write("mj_arm_v2/arm_umi_v2.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_v2/" + os.path.basename(f))
for f in ("aruco_arm_base_grid.png", "aruco_plane.obj"):
    shutil.copy(f, "mj_arm_v2/" + f)
print("wrote mj_arm_v2/arm_umi_v2.xml")