"""Gripper (with camera) -> self-contained MuJoCo bundle mj_gripper/:
- gripper.xml: converted from gripper_testing.urdf, + position actuator on the jaw, + a <camera> named
  wrist_cam at the B0332 optical frame (pose read from the GLB camera_body node; fovy 47 = the 70deg-H
  low-distortion lens on the 1280x800 sensor), + light/visual.
- copies all referenced meshes (meshdir=.) so the folder runs anywhere:
    python -m mujoco.viewer --mjcf=gripper.xml"""
import os, shutil
import numpy as np, trimesh
import xml.etree.ElementTree as ET
import mujoco

# 1. URDF -> MJCF
root = ET.parse("gripper_testing.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/gripper_raw.xml", model)

# 2. camera pose from the GLB (world == servo frame; URDF world joint = +90deg X -> mujoco Z-up)
cam_mesh = trimesh.load("gnode_camera_body.stl", force="mesh")   # baked world-mm geometry
V = cam_mesh.vertices
# optical axis = dominant (area-weighted) face normal = the board plane normal; tilt-robust
n = cam_mesh.face_normals; a = cam_mesh.area_faces
key = np.round(n, 2)
uniq, inv = np.unique(key, axis=0, return_inverse=True)
w = np.bincount(inv, weights=a)
axis = uniq[np.argmax(w)]; axis = axis / np.linalg.norm(axis)
# lens side = the FAR side of the board plane (lens protrudes ~18mm vs ~4.5mm of bottom parts)
t = V @ axis
faces_dom = np.isclose(np.abs(cam_mesh.face_normals @ axis), 1.0, atol=0.05)
t_board = float(np.median(cam_mesh.triangles_center[faces_dom] @ axis))
optical_s = axis if (t.max() - t_board) > (t_board - t.min()) else -axis
tip = float((V @ optical_s).max())
mid = V.mean(0)
pos_s = mid + optical_s * (tip - mid @ optical_s + 2.0)          # 2mm past the lens tip, on-axis
Rx = trimesh.transformations.rotation_matrix(np.pi/2, [1, 0, 0])[:3, :3]
opt_w = Rx @ optical_s
pos_w = Rx @ (pos_s / 1000.0)
up_w = np.array([0, 0, 1.0]) - np.dot([0, 0, 1.0], opt_w) * opt_w
up_w /= np.linalg.norm(up_w)
x_w = np.cross(up_w, -opt_w)                    # image right (right-handed with look=-z)
x_w /= np.linalg.norm(x_w)
y_w = up_w
print("camera pos(m):", np.round(pos_w, 3).tolist(), " optical:", np.round(opt_w, 2).tolist())
t = ET.parse("/tmp/gripper_raw.xml"); r2 = t.getroot()
wb = r2.find("worldbody")
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 list(x_w) + list(y_w)), "fovy": "47"})
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": "1280", "offheight": "800"})
ET.SubElement(vis, "headlight", {"ambient": "0.45 0.45 0.45"})
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_gripper", exist_ok=True)
ET.indent(t, "  "); t.write("mj_gripper/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_gripper/" + os.path.basename(f))
print("wrote mj_gripper/gripper.xml (+meshes)")

# 3. renders: third person + THROUGH the wrist cam (closed + open)
m2 = mujoco.MjModel.from_xml_path("mj_gripper/gripper.xml")
d2 = mujoco.MjData(m2)
from PIL import Image
ims = []
for q in (0.0, 0.7):
    d2.qpos[:] = q; d2.ctrl[:] = q; mujoco.mj_forward(m2, d2)
    with mujoco.Renderer(m2, 400, 640) as r:
        cam = mujoco.MjvCamera(); cam.lookat[:] = [-0.03, 0, 0.02]; cam.distance = 0.35; cam.azimuth = 40; cam.elevation = -20
        r.update_scene(d2, cam); ims.append(r.render())
    with mujoco.Renderer(m2, 400, 640) as r:
        r.update_scene(d2, "wrist_cam"); ims.append(r.render())
Image.fromarray(np.vstack([np.hstack(ims[:2]), np.hstack(ims[2:])])).save("mj_gripper_views.png")
print("rendered mj_gripper_views.png (rows: closed/open; cols: third-person / wrist_cam)")