"""UMI handheld variant v2: ORIENTED rounded handle (keeps Cameron rake, buried into the holder), clean
finger ring + lever arm (moving side), ArUco marker cube passed through EXACTLY as placed (flat faces for
markers — never rounded). Servo stays as unpowered pivot + encoder.
Reads umi_from_blender.glb -> gnode_handle/finger_ring/aruco_cube.stl + umi_gripper.urdf."""
import numpy as np, trimesh
from gen_curvy import _rrect, _stitch
import xml.etree.ElementTree as ET

s = trimesh.load("umi_from_blender.glb")
def wpts(node):
    T, gn = s.graph[node]
    return trimesh.transform_points(s.geometry[gn].vertices, np.array(T, float)) * 100.0
def U(p): return trimesh.boolean.union(p, engine="manifold")

hold = wpts("holder"); hc = (hold.min(0) + hold.max(0)) / 2
yoke_pts = wpts("yoke")

# --- handle: ORIENTED rounded loft over his (raked) envelope, buried 6mm into the holder end ---
Tn, gn = s.graph["handle"]
handle = trimesh.Trimesh(vertices=trimesh.transform_points(s.geometry[gn].vertices, np.array(Tn, float)) * 100.0,
                         faces=s.geometry[gn].faces)      # his placement of the already-polished grip: passthrough
# weld tab: bury the grip 6mm into the holder east wall where they meet
hhi = hold.max(0); hlo2 = hold.min(0)
near = handle.vertices[handle.vertices[:, 0] < hhi[0] + 8.0]
if len(near) > 3:
    y0, y1 = max(near[:, 1].min(), hlo2[1]), min(near[:, 1].max(), hhi[1])
    z0, z1 = max(near[:, 2].min(), hlo2[2]), min(near[:, 2].max(), hhi[2])
    if y1 > y0 + 4 and z1 > z0 + 4:
        tab = trimesh.creation.box(extents=[14.0, y1 - y0, z1 - z0])
        tab.apply_translation([hhi[0] - 1.0, (y0 + y1) / 2, (z0 + z1) / 2])
        handle = U([handle, tab])
        print(f"weld tab y[{y0:.0f},{y1:.0f}] z[{z0:.0f},{z1:.0f}]")
handle.export("gnode_handle.stl")

# --- finger ring + lever arm (unchanged logic) ---
rv = wpts("finger_ring"); rlo, rhi = rv.min(0), rv.max(0)
rc = (rlo + rhi) / 2; ext = rhi - rlo
ax = int(np.argmin(ext))
r_minor = ext[ax] / 2
r_major = (ext[[i for i in range(3) if i != ax]].mean() / 2) - r_minor
ring = trimesh.creation.torus(major_radius=r_major, minor_radius=r_minor)
Rt = trimesh.geometry.align_vectors([0, 0, 1], np.eye(3)[ax])
ring.apply_transform(Rt); ring.apply_translation(rc)
yc = yoke_pts.mean(0)
d = yc - rc; d[ax] = 0.0; dn = d / np.linalg.norm(d)
startp = rc + dn * (r_major - 2.0)
endp = rc + dn * (np.linalg.norm(d) - 8.0)
arm = trimesh.creation.cylinder(radius=4.0, height=np.linalg.norm(endp - startp), sections=32)
v = endp - startp; zc = np.array([0, 0, 1.0]); nrm = v / np.linalg.norm(v)
axc = np.cross(zc, nrm); sn = np.linalg.norm(axc)
if sn > 1e-9: arm.apply_transform(trimesh.transformations.rotation_matrix(np.arctan2(sn, np.dot(zc, nrm)), axc / sn))
arm.apply_translation((startp + endp) / 2)
ringasm = U([ring, arm]); ringasm.export("gnode_finger_ring.stl")

# --- aruco cube: EXACT passthrough (flat faces) ---
cube = wpts("Cube")
cm = trimesh.PointCloud(cube).convex_hull
cm.export("gnode_aruco_cube.stl")
print(f"handle wt={handle.is_watertight} vol={handle.volume/1000:.1f}cm3 | ring wt={ringasm.is_watertight} | aruco cube {np.round(cm.extents,0).tolist()}mm")

# --- umi_gripper.urdf ---
t = ET.parse("gripper_testing.urdf"); root = t.getroot()
links = {l.get("name"): l for l in root.findall("link")}
def vis(link, fn, rgba, off="0 0 0"):
    v = ET.SubElement(link, "visual"); ET.SubElement(v, "origin", {"xyz": off, "rpy": "0 0 0"})
    ET.SubElement(ET.SubElement(v, "geometry"), "mesh", {"filename": fn, "scale": "0.001 0.001 0.001"})
    ET.SubElement(ET.SubElement(v, "material", {"name": "m" + fn}), "color", {"rgba": rgba})
vis(links["base"], "gnode_handle.stl", "0.25 0.30 0.38 1")
vis(links["base"], "gnode_aruco_cube.stl", "0.92 0.92 0.92 1")
# marker plane positions as plain colored plates (textures live in the MuJoCo build, not the URDF viewer)
cb = trimesh.load("gnode_aruco_cube.stl", force="mesh")
clo, chi = cb.bounds
ccx, ccy, ccz = (clo + chi) / 2
sx, sy, sz = chi - clo
p1 = trimesh.creation.box(extents=[sx, 1.2, sz]); p1.apply_translation([ccx, chi[1] + 0.6, ccz])   # top, full face
p2 = trimesh.creation.box(extents=[1.2, sy, sz]); p2.apply_translation([chi[0] + 0.6, ccy, ccz])   # back, full face
trimesh.util.concatenate([p1, p2]).export("gnode_aruco_planes.stl")
vis(links["base"], "gnode_aruco_planes.stl", "0.10 0.10 0.10 1")
vis(links["jaw_link"], "gnode_finger_ring.stl", "0.85 0.35 0.25 1", off="0.0255 0 0")
ET.indent(t, "  "); t.write("umi_gripper.urdf", encoding="unicode", xml_declaration=True)
print("wrote umi_gripper.urdf")
