"""v2 UMI finish (run AFTER gen_umi_v2.py + gen_v2_umi.py): shell+recess the cube, fuse the umi fixed
print piece (gripper fixed side + camera bracket + mount + cube), add transplants to umi_gripper_v2.urdf."""
import numpy as np, trimesh
import xml.etree.ElementTree as ET
from gen_curvy import _rrect, _stitch

import trimesh.repair
def L(f):
    m = trimesh.load(f, force="mesh")
    if not m.is_volume:
        m.merge_vertices(merge_tex=True, merge_norm=True)
        m.update_faces(m.nondegenerate_faces()); m.update_faces(m.unique_faces())
        trimesh.repair.fill_holes(m); m.fix_normals()
    return m
def D(a, b): return trimesh.boolean.difference([a, b], engine="manifold")
def U(parts):
    m = trimesh.boolean.union(parts, engine="manifold")
    keep = [b for b in m.split(only_watertight=False) if b.volume > 100.0]
    return keep[0] if len(keep) == 1 else trimesh.util.concatenate(keep)

# --- shell + marker recesses (port of the v1 print treatment) ---
cube = L("gnode_v2_aruco_cube.stl")
lo, hi = cube.bounds
if hi[0] - lo[0] > 6.5:                                    # thin the slab to 6mm, marker face kept
    cube = trimesh.intersections.slice_mesh_plane(cube, plane_normal=[1, 0, 0],
                                                  plane_origin=[hi[0] - 6.0, 0, 0], cap=True)
    cube.export("gnode_v2_aruco_cube.stl")
    lo, hi = cube.bounds
# cube prints SOLID (closed on all faces) — slicer infill handles the material saving; the old
# open-face shell left a gaping pane and a shelled-closed roof would need 74mm bridging
cx, cy, cz = (lo[0]+hi[0])/2, (lo[1]+hi[1])/2, (lo[2]+hi[2])/2
def rec(center, size, ax, depth, face):
    e = [size, size, size]; e[ax] = depth * 2
    r = trimesh.creation.box(extents=e)
    p = list(center); p[ax] = face
    r.apply_translation(p); return r
# single-board slab, FLAT face (Cameron 2026-07-08): no recess/notch — sheet glues flat on +x face
cube.export("gnode_v2_aruco_cube_shelled.stl")

# --- bond pad cube<->riser/plate east faces ---
riser = L("gnode_v2_fixed_riser.stl")
re = riser.bounds[1][0]
pad = None
def prism_aa(lo, hi, r=4.0, N=12):
    Ln = hi - lo; k = int(np.argmax(Ln)); lat = [i for i in range(3) if i != k]
    e1, e2 = np.eye(3)[lat[0]], np.eye(3)[lat[1]]
    h1, h2 = Ln[lat[0]] / 2, Ln[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.4, h2 - 0.4))
        p = c.copy(); p[k] = lo[k] + Ln[k] * t
        rings.append(p + np.outer(q[:, 0], e1) + np.outer(q[:, 1], e2))
    return _stitch(rings)
rb, cb2 = riser.bounds, cube.bounds
yo = (max(rb[0][1], cb2[0][1]), min(rb[1][1], cb2[1][1]))
zo = (max(rb[0][2], cb2[0][2]), min(rb[1][2], cb2[1][2]))
if yo[1] - yo[0] < 8.0:                                   # degenerate overlap: cube slid past the riser
    yc = (yo[0] + yo[1]) / 2
    yo = (max(cb2[0][1], yc - 7.0), min(rb[1][1], yc + 7.0))
if zo[1] - zo[0] < 8.0:
    zc = (zo[0] + zo[1]) / 2
    zo = (zc - 7.0, zc + 7.0)
pad = prism_aa(np.array([re - 5.0, yo[0], zo[0] + 1]),
               np.array([min(cb2[0][0] + 4.0, cb2[1][0] - 2.0), yo[1], zo[1] - 1]))   # never through the marker face
print(f"cube bridge: x {re-5:.0f}..{cb2[0][0]+7:.0f} y {yo[0]:.0f}..{yo[1]:.0f} z {zo[0]:.0f}..{zo[1]:.0f}")
umi_fixed = U([L("gripper_v2_print_fixed.stl"), cube, pad, L("gnode_v2_cam_conn.stl"), L("gnode_v2_camera_mount.stl")])
umi_fixed.export("gripper_v2_print_fixed_umi.stl")
bodies = len(umi_fixed.split(only_watertight=False))
print(f"gripper_v2_print_fixed_umi: wt={umi_fixed.is_watertight} bodies={bodies} vol={umi_fixed.volume/1000:.1f}cm3 bbox={[round(v) for v in umi_fixed.extents]}mm{'  <-- DISJOINT' if bodies>1 else ''}")

# --- URDF: add transplants to base link ---
t = ET.parse("umi_gripper_v2.urdf")
base = next(l for l in t.getroot().findall("link") if l.get("name") == "base")
adds = [("gnode_v2_aruco_cube_shelled.stl", "0.92 0.92 0.92 1"), ("gnode_v2_cam_conn.stl", "0.93 0.60 0.20 1"),
        ("gnode_v2_camera_mount.stl", "0.93 0.60 0.20 1"), ("gnode_v2_camera_body.stl", "0.15 0.35 0.15 1"),
        ("gnode_v2_camera_cover.stl", "0.55 0.85 0.65 1")]
for fn, rgba in adds:
    v = ET.SubElement(base, "visual"); ET.SubElement(v, "origin", {"xyz": "0 0 0", "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": "mu" + fn}), "color", {"rgba": rgba})
ET.indent(t, "  ")
t.write("umi_gripper_v2.urdf", encoding="unicode", xml_declaration=True)
print("umi_gripper_v2.urdf updated with camera + cube")