"""[feetech_calib] Servo-zeroing fixture v2 — a 180deg SADDLE (not a full pocket).

Cameron's correction: the full-silhouette pocket enclosed the part 360deg, so you
couldn't place it over an assembled joint. Instead: a top-half shell that hugs the
servo+yoke over the yoke's screw-boss region and DROPS ON from above (open underneath
=> lifts straight off). It keys the yoke angle by conforming to the yoke arm + servo
body; the (limp) yoke only lets the saddle seat fully at q=0.

Build: voxel-offset shell of (servo U yoke) — outer skin minus inner skin — then keep
only the top hemisphere (z >= Z_CUT). Insertion/removal along +Z (the thin axis);
joint axis Y is horizontal, so the top-half wrap is 180deg around the Y-cylinder.
"""
import numpy as np, trimesh, scipy.ndimage as ndi
from trimesh.voxel import ops as vox_ops

CLEAR = 0.40   # inner clearance to the part (slip fit)
WALL  = 3.5    # shell wall thickness
P     = 0.6    # voxel pitch (mm) — resolution of the offset
Z_CUT = 0.0    # keep z >= this => top 180deg (open bottom, drops on/off)

servo = trimesh.load("st3215.stl", force="mesh")
yoke  = trimesh.load("yoke_fillet.stl", force="mesh")   # net-zero mate (servo frame)
asm   = trimesh.util.concatenate([servo, yoke])

vg = asm.voxelized(pitch=P).fill()
PAD = int(np.ceil((CLEAR + WALL) / P)) + 2             # empty margin so the offset can grow
S  = np.pad(vg.matrix, PAD, mode="constant", constant_values=False)
dist = ndi.distance_transform_edt(~S) * P              # mm from nearest solid voxel
inner = dist <= CLEAR
outer = dist <= (CLEAR + WALL)
shell_mat = outer & ~inner

def mat_to_world(mat):
    m = vox_ops.matrix_to_marching_cubes(mat, pitch=1.0)  # padded-index-space verts
    m.apply_transform(vg.transform)                       # P*index + origin
    m.apply_translation([-P * PAD, -P * PAD, -P * PAD])   # undo the pad shift
    return m

shell = mat_to_world(shell_mat)

# alignment sanity: outer skin of the servo should straddle servo.bounds +/- (CLEAR+WALL)
osk = mat_to_world(outer)
print(f"align chk: outer bounds {np.round(osk.bounds[0],1)}..{np.round(osk.bounds[1],1)} "
      f"vs servo {np.round(servo.bounds[0],1)}..{np.round(servo.bounds[1],1)} (grow {CLEAR+WALL:.1f})")

# keep top hemisphere only (z >= Z_CUT): intersect with a big top box
big = 500.0
topbox = trimesh.creation.box(extents=[big, big, big])
topbox.apply_translation([0, 0, Z_CUT + big/2.0])       # box spans z in [Z_CUT, ...]
saddle = trimesh.boolean.intersection([shell, topbox], engine="manifold")

# keep the single largest body (drop any voxel crumbs)
parts = saddle.split(only_watertight=False)
if len(parts) > 1:
    saddle = max(parts, key=lambda m: m.volume)

saddle.export("servo_zero_fixture.stl")
e = saddle.extents
print(f"saddle: wt={saddle.is_watertight} bodies={len(saddle.split(only_watertight=False))} "
      f"vol={saddle.volume/1000:.1f}cm3 bbox=({e[0]:.1f},{e[1]:.1f},{e[2]:.1f})mm "
      f"z=({saddle.bounds[0][2]:.1f},{saddle.bounds[1][2]:.1f})")

# ---- viewer URDFs ----
def vis(fn, rgba, org="0 0 0"):
    return (f'    <visual><origin xyz="{org}" rpy="0 0 0"/>'
            f'<geometry><mesh filename="{fn}" scale="0.001 0.001 0.001"/></geometry>'
            f'<material name="{fn}"><color rgba="{rgba}"/></material></visual>\n')

ASM = (vis("st3215.stl","0.20 0.20 0.22 1") + vis("holder_fillet.stl","0.55 0.58 0.62 1")
       + vis("yoke_fillet.stl","0.85 0.60 0.20 1"))
SAD = vis("servo_zero_fixture.stl","0.25 0.80 0.40 0.85")

open("servo_zero_fixture.urdf","w").write(
    f'<?xml version="1.0"?>\n<robot name="servo_zero_fixture">\n  <link name="saddle">\n{SAD}  </link>\n</robot>\n')

open("servo_zero_fixture_demo.urdf","w").write(
    f'<?xml version="1.0"?>\n<robot name="servo_zero_fixture_demo">\n'
    f'  <link name="assembly">\n{ASM}  </link>\n'
    f'  <link name="saddle">\n{SAD}  </link>\n'
    f'  <joint name="seated" type="fixed"><parent link="assembly"/><child link="saddle"/>'
    f'<origin xyz="0 0 0" rpy="0 0 0"/></joint>\n</robot>\n')

# install/explode: assembly fixed, saddle lifts straight up (+Z) off the joint
open("servo_zero_fixture_install.urdf","w").write(
    f'<?xml version="1.0"?>\n<robot name="servo_zero_fixture_install">\n'
    f'  <link name="assembly">\n{ASM}  </link>\n'
    f'  <link name="saddle">\n{SAD}  </link>\n'
    f'  <joint name="explode" type="prismatic"><parent link="assembly"/><child link="saddle"/>'
    f'<origin xyz="0 0 0" rpy="0 0 0"/><axis xyz="0 0 1"/>'
    f'<limit lower="0" upper="0.05" effort="1" velocity="1"/></joint>\n</robot>\n')
print("wrote servo_zero_fixture{,_demo,_install}.urdf (saddle = removable cap, explode 0..50mm)")