"""umi_gripper_v2 -> MuJoCo bundle mj_umi_v2/ with real ArUco grid textures + wrist camera at the
NEW camera position (pose derived from camera_body geometry). v2 world mapping: rpy -pi/2 =>
world = (x_s, z_s, -y_s)."""
import os, shutil
import numpy as np, trimesh
import xml.etree.ElementTree as ET
import mujoco

root = ET.parse("umi_gripper_v2.urdf").getroot()
_has_cam = any("camera_mount" in (m.get("filename") or "") for m in root.iter("mesh"))
assert _has_cam, "STALE URDF - run gen_v2_finish.py first"
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/v2_raw.xml", model)
t = ET.parse("/tmp/v2_raw.xml"); r2 = t.getroot()
asset = r2.find("asset")
def W(p):
    return np.array([p[0], p[2], -p[1]]) / 1000.0
cube = trimesh.load("gnode_v2_aruco_cube.stl", force="mesh")
lo, hi = cube.bounds; cc = (lo + hi) / 2
ET.SubElement(asset, "texture", {"name": "ar_back", "type": "2d", "file": "aruco_umi_v2_grid.png"})
ET.SubElement(asset, "material", {"name": "ar_back", "texture": "ar_back"})
# slab face is 120 (y_s, world-vertical) x 80 (z_s): plane local x carries 120, local y carries 80
ET.SubElement(asset, "mesh", {"name": "aruco_plane", "file": "aruco_plane.obj",
                              "scale": f"{0.120/0.3:.5f} {0.080/0.3:.5f} 0.001", "inertia": "shell"})
wb = r2.find("worldbody")
back_pos = W([hi[0] + 0.7, cc[1], cc[2]])
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"})
# wrist camera from camera_body geometry (servo frame): board = dominant face normal; optical = away from board bulk
body = trimesh.load("gnode_v2_camera_body.stl", force="mesh")
areas, normals = body.area_faces, body.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 = n / np.linalg.norm(n)
mid = body.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                    # ROBUST: optical points AWAY from the mount plate
opt_s = n
tip = body.vertices[np.argmax(body.vertices @ opt_s)]
pos_s = mid + opt_s * ((tip - mid) @ opt_s + 2.0)
opt_w = np.array([opt_s[0], opt_s[2], -opt_s[1]])
pos_w = W(pos_s)
up_guess = np.array([0, 0, 1.0])
x_w = np.cross(up_guess, -opt_w); x_w /= np.linalg.norm(x_w)
y_w = np.cross(-opt_w, x_w)
ET.SubElement(wb, "camera", {"name": "wrist_cam", "pos": " ".join(f"{v:.4f}" for v in pos_w),
                             "xyaxes": " ".join(f"{v:.4f}" for v in np.concatenate([x_w, y_w])), "fovy": "47"})
print("wrist cam at", np.round(pos_w, 3).tolist(), "optical", np.round(opt_w, 2).tolist())
ET.SubElement(asset, "texture", {"name": "floor", "type": "2d", "builtin": "checker", "rgb1": "0.25 0.3 0.35",
                                 "rgb2": "0.18 0.22 0.26", "width": "300", "height": "300"})
ET.SubElement(asset, "material", {"name": "floor", "texture": "floor", "texuniform": "true", "texrepeat": "6 6"})
ET.SubElement(wb, "geom", {"type": "plane", "size": "1 1 0.05", "pos": "0 0 -0.25", "material": "floor"})
ET.SubElement(wb, "geom", {"type": "box", "size": "0.015 0.015 0.015", "pos": "-0.17 0 -0.02",
                           "rgba": "0.85 0.3 0.25 1", "contype": "0", "conaffinity": "0"})
ET.SubElement(wb, "light", {"pos": "0.3 0.3 0.8", "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.5 0.5 0.5"})
act = ET.SubElement(r2, "actuator")
ET.SubElement(act, "position", {"joint": "jaw", "kp": "5", "kv": "0.4", "ctrlrange": "-0.9 0.1", "name": "jaw"})
os.makedirs("mj_umi_v2", exist_ok=True)
ET.indent(t, "  "); t.write("mj_umi_v2/umi_gripper_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_umi_v2/" + os.path.basename(f))
for f in ("aruco_umi_v2_grid.png", "aruco_plane.obj"):
    shutil.copy(f, "mj_umi_v2/" + f)
print("wrote mj_umi_v2/umi_gripper_v2.xml")