import numpy as np, trimesh
import xml.etree.ElementTree as ET
from gen_curvy import _rrect, _stitch

s = trimesh.load("umi_v2_from_blender.glb")
def wmesh(node):
    T, gn = s.graph[node]
    m = trimesh.Trimesh(trimesh.transform_points(s.geometry[gn].vertices, np.array(T, float)) * 100.0,
                        s.geometry[gn].faces)
    if np.linalg.det(np.array(T, float)[:3, :3]) < 0:
        m.invert()
    return m
def U(p): return trimesh.boolean.union(p, engine="manifold")

M = {n: wmesh(n) for n in s.graph.nodes_geometry}
servo_hull = trimesh.PointCloud(trimesh.load("st3215.stl", force="mesh").vertices).convex_hull

def prism_aa(lo, hi, r=6.0, N=14):
    L = hi - lo; k = int(np.argmax(L)); lat = [i for i in range(3) if i != k]
    e1, e2 = np.eye(3)[lat[0]], np.eye(3)[lat[1]]
    h1, h2 = L[lat[0]] / 2, L[lat[1]] / 2
    c = (lo + hi) / 2
    rings = []
    for i in range(N):
        t = i / (N - 1)
        q = _rrect(h1, h2, min(r, h1 - 0.5, h2 - 0.5))
        p = c.copy(); p[k] = lo[k] + L[k] * t
        rings.append(p + np.outer(q[:, 0], e1) + np.outer(q[:, 1], e2))
    return _stitch(rings)

r_key = "Cube.001" if "Cube.001" in M else "fixed_riser"
p_key = "Cube" if "Cube" in M else "fixed_plate"
riser = prism_aa(M[r_key].bounds[0], M[r_key].bounds[1])
plate = prism_aa(M[p_key].bounds[0], M[p_key].bounds[1])
riser.export("gnode_v2_fixed_riser.stl"); plate.export("gnode_v2_fixed_plate.stl")

mv_key = "enpire_moving.001" if "enpire_moving.001" in M else "enpire_moving"
yk, mv = M["yoke"], M[mv_key]
d = trimesh.proximity.ProximityQuery(yk).signed_distance(trimesh.sample.sample_surface(mv, 8000)[0])
overlap = float(d.max())
print("moving finger vs yoke overlap/gap (mm):", round(overlap, 2))
rb_lo, rb_hi = mv.bounds
root_x1 = rb_hi[0]
root = mv.vertices[mv.vertices[:, 0] > root_x1 - 30.0]
rc = (root.min(0) + root.max(0)) / 2
yb_lo, yb_hi = yk.bounds
lo = np.array([root_x1 - 6.0, max(rb_lo[1], yb_lo[1]) + 4.0, yb_hi[2] - 6.0])
hi = np.array([yb_lo[0] + 6.0, min(rb_hi[1], yb_hi[1]) - 4.0, rb_lo[2] + 8.0])
lo2, hi2 = np.minimum(lo, hi), np.maximum(lo, hi)
conn = prism_aa(lo2, hi2, r=4.0)
conn.export("gnode_v2_yoke_conn.stl")
print("yoke_conn envelope:", np.round(lo2, 1).tolist(), np.round(hi2, 1).tolist())
exports = [("holder", "gnode_v2_holder.stl"), ("yoke", "gnode_v2_yoke.stl"),
           ("enpire_fixed", "gnode_v2_enpire_fixed.stl")]
mv_name = "enpire_moving.001" if "enpire_moving.001" in M else "enpire_moving"
exports.append((mv_name, "gnode_v2_enpire_moving.stl"))
for cn in ("thumb_cradle", "finger_cradle"):
    if cn in M: exports.append((cn, f"gnode_v2_{cn}.stl"))
for node, out in exports:
    M[node].export(out)
import shutil as _sh
_sh.copy("yoke_gen_thin.stl", "gnode_v2_yoke.stl")   # thin-plate yoke (gen_yoke.py) overrides the GLB yoke

for name, mesh in [("riser", riser), ("plate", plate), ("yoke_conn", conn),
                   ("enpire_fixed", M["enpire_fixed"]), ("enpire_moving", M[mv_key])]:
    pts, _ = trimesh.sample.sample_surface(mesh, 8000)
    pen = float(trimesh.proximity.ProximityQuery(servo_hull).signed_distance(pts).max())
    print(name, "wt=", mesh.is_watertight, " vs servo penetration mm:", round(max(0.0, pen), 2))

AX = -0.0255
R = ET.Element("robot", {"name": "umi_v2"}); ET.SubElement(R, "link", {"name": "world"})
wj = ET.SubElement(R, "joint", {"name": "world_to_base", "type": "fixed"})
ET.SubElement(wj, "parent", {"link": "world"}); ET.SubElement(wj, "child", {"link": "base"})
ET.SubElement(wj, "origin", {"xyz": "0 0 0", "rpy": "-1.5708 0 0"})   # display flip: camera/box side up
COLS = {"gnode_v2_thumb_cradle.stl": "0.90 0.40 0.30 1", "gnode_v2_handle_column.stl": "0.30 0.55 0.90 1",
        "st3215.stl": "0.35 0.35 0.40 1", "gnode_v2_holder.stl": "0.20 0.52 0.85 1",
        "gnode_v2_enpire_fixed.stl": "0.55 0.75 0.60 1", "gnode_v2_fixed_riser.stl": "0.30 0.66 0.45 1",
        "gnode_v2_fixed_plate.stl": "0.30 0.66 0.45 1", "gnode_v2_yoke.stl": "0.90 0.65 0.20 1",
        "gnode_v2_enpire_moving.stl": "0.55 0.60 0.85 1", "gnode_v2_yoke_conn.stl": "0.85 0.30 0.30 1"}
def vis(link, fn, 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": COLS[fn]})
base = ET.SubElement(R, "link", {"name": "base"})
for fn in ("st3215.stl", "gnode_v2_holder.stl", "gnode_v2_enpire_fixed.stl", "gnode_v2_fixed_riser.stl", "gnode_v2_fixed_plate.stl", "gnode_v2_handle_column.stl"):
    vis(base, fn)
jj = ET.SubElement(R, "joint", {"name": "jaw", "type": "revolute"})
ET.SubElement(jj, "parent", {"link": "base"}); ET.SubElement(jj, "child", {"link": "jaw_link"})
ET.SubElement(jj, "origin", {"xyz": str(AX) + " 0 0", "rpy": "0 0 0"}); ET.SubElement(jj, "axis", {"xyz": "0 1 0"})
ET.SubElement(jj, "limit", {"lower": "-0.9", "upper": "0.1", "effort": "3", "velocity": "3"})
jl = ET.SubElement(R, "link", {"name": "jaw_link"})
for fn in ("gnode_v2_yoke.stl", "gnode_v2_yoke_conn.stl", "gnode_v2_enpire_moving.stl", "gnode_v2_thumb_cradle.stl"):
    vis(jl, fn, off=str(-AX) + " 0 0")
ET.indent(ET.ElementTree(R), "  "); ET.ElementTree(R).write("umi_gripper_v2.urdf", encoding="unicode", xml_declaration=True)
out = trimesh.Scene()
for node, fn in [("servo", "st3215.stl"), ("holder", "gnode_v2_holder.stl"), ("yoke", "gnode_v2_yoke.stl"),
                 ("thumb_cradle", "gnode_v2_thumb_cradle.stl"), ("finger_cradle", "gnode_v2_finger_cradle.stl"),
                 ("enpire_fixed", "gnode_v2_enpire_fixed.stl"), ("enpire_moving", "gnode_v2_enpire_moving.stl"),
                 ("fixed_riser", "gnode_v2_fixed_riser.stl"), ("fixed_plate", "gnode_v2_fixed_plate.stl"),
                 ("yoke_conn", "gnode_v2_yoke_conn.stl"),
                 ("cam_conn", "gnode_v2_cam_conn.stl"), ("aruco_cube", "gnode_v2_aruco_cube.stl")] + \
                ([("handle_column", "gnode_v2_handle_column.stl")] if __import__("os").path.exists("gnode_v2_handle_column.stl") else []):
    m = trimesh.load(fn, force="mesh"); m.apply_scale(0.01)
    out.add_geometry(m, node_name=node, geom_name=node)
# camera group as a TREE: camera_mount is the parent, body/cover children (move the mount, all follow)
m = trimesh.load("gnode_v2_camera_mount.stl", force="mesh"); m.apply_scale(0.01)
out.add_geometry(m, node_name="camera_mount", geom_name="camera_mount")
for node, fn in [("camera_body", "gnode_v2_camera_body.stl"), ("camera_cover", "gnode_v2_camera_cover.stl")]:
    m = trimesh.load(fn, force="mesh"); m.apply_scale(0.01)
    out.add_geometry(m, node_name=node, geom_name=node, parent_node_name="camera_mount")
out.export("umi_gripper_v2.glb")
print("wrote umi_gripper_v2.urdf + glb")