"""[feetech_calib] Universal servo-zeroing fixture for the smith300 arm.

One reusable bench tool that holds an STS3215 + yoke_fillet assembly at the URDF
q=0 pose so `INST_OFSCAL` can set that pose to tick 2048 (zero) on every servo.

Principle: the pocket is the *silhouette of the servo+yoke at q=0*, extruded
straight down. The assembly only drops in when the (limp) yoke is rotated to q=0
— at any other angle the yoke's shadow doesn't match the pocket and it won't seat.
Keys off the servo BODY (st3215) + yoke, both identical at every joint, so ONE
fixture zeroes all six servos (and any future swap). Absorbs horn re-clocking:
INST_OFSCAL doesn't care what the raw tick was, only that this pose reads 2048.

Module frame (mm): servo = st3215.stl at identity. yoke_fillet.stl mates the servo
at ZERO net offset (verified: 0.00mm surface gap, wraps the servo). In
two_servo_beauty.urdf this same pose is expressed as joint output_N origin
(-25.5,0,0) PLUS a yoke visual origin (+25.5,0,0) that cancels it (see
gen_two_servo.py:40) — net zero, so the mesh is effectively already in the servo
frame. q=0 => yoke at its natural orientation. This fixture nests that mated pose.
Insertion axis = +Z (assembly's thin 26mm dimension); joint axis Y lies horizontal.
"""
import numpy as np, trimesh
from shapely.ops import unary_union
from trimesh.creation import extrude_polygon

CLEAR = 0.40   # cavity outward buffer: drop-in slip fit across FDM swell
WALL  = 6.0    # wall thickness around the silhouette
BASE  = 4.0    # solid floor below the part
Z_LO  = -13.0  # part bottom (both servo & yoke reach ~-13 in Z)
Z_HI  =  13.0  # part top / wall top

def foot(m):
    """True XY silhouette (shadow along +Z) as a shapely polygon."""
    p = m.projected([0, 0, 1])
    return unary_union(list(p.polygons_full))

servo = trimesh.load("st3215.stl", force="mesh")
yoke  = trimesh.load("yoke_fillet.stl", force="mesh")   # already in servo frame; zero-offset mate (q=0)

sil = unary_union([foot(servo), foot(yoke)])       # combined q=0 silhouette
cav2d = sil.buffer(CLEAR, join_style=2)            # pocket outline (mitred), may be 2 islands
out2d = sil.buffer(WALL, join_style=2)             # block outline (WALL buffer bridges the gap)

def polys(geom):
    return list(geom.geoms) if geom.geom_type == "MultiPolygon" else [geom]

# Block: Z from (Z_LO-BASE) up to Z_HI.  Cavity: open-topped pocket(s), floor at Z_LO.
block_parts = []
for g in polys(out2d):
    m = extrude_polygon(g, height=(Z_HI - (Z_LO - BASE)))
    m.apply_translation([0, 0, Z_LO - BASE]); block_parts.append(m)
block = block_parts[0] if len(block_parts) == 1 else trimesh.boolean.union(block_parts, engine="manifold")

cavities = []
for g in polys(cav2d):
    m = extrude_polygon(g, height=80.0)
    m.apply_translation([0, 0, Z_LO]); cavities.append(m)

fix = trimesh.boolean.difference([block] + cavities, engine="manifold")

# Thumb scallop on the exposed servo-front exterior wall (+X, past the yoke) for lift-out.
sb = servo.bounds
scallop = trimesh.creation.cylinder(radius=9.0, height=(Z_HI - Z_LO) + 2, sections=48)
scallop.apply_translation([sb[1][0] + WALL, (sb[0][1] + sb[1][1]) / 2.0, (Z_LO + Z_HI) / 2.0])
fix = trimesh.boolean.difference([fix, scallop], engine="manifold")

fix.export("servo_zero_fixture.stl")
print(f"n_silhouette_islands: {len(polys(cav2d))}  block_islands: {len(polys(out2d))}")

b = len(fix.split(only_watertight=False))
e = fix.extents
print(f"servo_zero_fixture: wt={fix.is_watertight} bodies={b} vol={fix.volume/1000:.1f}cm3 "
      f"bbox=({e[0]:.1f},{e[1]:.1f},{e[2]:.1f})mm")
print(f"silhouette bbox xy: {np.round(sil.bounds,1)}")

# --- bare-part URDF (for printing / viewing the fixture itself) ---
def urdf_one(stl, out, name):
    open(out, "w").write(f'''<?xml version="1.0"?>
<robot name="{name}">
  <link name="{name}">
    <visual><origin xyz="0 0 0" rpy="0 0 0"/>
      <geometry><mesh filename="{stl}" scale="0.001 0.001 0.001"/></geometry>
      <material name="jig"><color rgba="0.25 0.55 0.85 1"/></material>
    </visual>
  </link>
</robot>
''')
urdf_one("servo_zero_fixture.stl", "servo_zero_fixture.urdf", "servo_zero_fixture")

# --- demo URDF: servo + yoke seated in the fixture (fit check) ---
open("servo_zero_fixture_demo.urdf", "w").write('''<?xml version="1.0"?>
<robot name="servo_zero_fixture_demo">
  <link name="fixture">
    <visual><origin xyz="0 0 0" rpy="0 0 0"/>
      <geometry><mesh filename="servo_zero_fixture.stl" scale="0.001 0.001 0.001"/></geometry>
      <material name="jig"><color rgba="0.25 0.55 0.85 0.55"/></material></visual></link>
  <link name="servo">
    <visual><origin xyz="0 0 0" rpy="0 0 0"/>
      <geometry><mesh filename="st3215.stl" scale="0.001 0.001 0.001"/></geometry>
      <material name="srv"><color rgba="0.2 0.2 0.22 1"/></material></visual></link>
  <link name="yoke">
    <visual><origin xyz="0 0 0" rpy="0 0 0"/>
      <geometry><mesh filename="yoke_fillet.stl" scale="0.001 0.001 0.001"/></geometry>
      <material name="yk"><color rgba="0.8 0.5 0.2 1"/></material></visual></link>
  <joint name="f2s" type="fixed"><parent link="fixture"/><child link="servo"/>
    <origin xyz="0 0 0" rpy="0 0 0"/></joint>
  <joint name="s2y" type="fixed"><parent link="servo"/><child link="yoke"/>
    <origin xyz="0 0 0" rpy="0 0 0"/></joint>
</robot>
''')
print("wrote servo_zero_fixture.urdf + servo_zero_fixture_demo.urdf")

# --- install / exploded-view URDF: servo+holder+yoke lift out of the jig on a slider ---
# Root = jig (fixed bench fixture). Child = the servo assembly, on a prismatic "explode"
# joint along +Z (the insertion axis). explode=0 => seated in the pocket; increase the
# slider to lift the assembly straight up out of the jig and see how it registers.
open("servo_zero_fixture_install.urdf", "w").write('''<?xml version="1.0"?>
<robot name="servo_zero_fixture_install">
  <link name="jig">
    <visual><origin xyz="0 0 0" rpy="0 0 0"/>
      <geometry><mesh filename="servo_zero_fixture.stl" scale="0.001 0.001 0.001"/></geometry>
      <material name="jig"><color rgba="0.25 0.55 0.85 0.55"/></material></visual></link>
  <link name="assembly">
    <visual><origin xyz="0 0 0" rpy="0 0 0"/>
      <geometry><mesh filename="st3215.stl" scale="0.001 0.001 0.001"/></geometry>
      <material name="srv"><color rgba="0.20 0.20 0.22 1"/></material></visual>
    <visual><origin xyz="0 0 0" rpy="0 0 0"/>
      <geometry><mesh filename="holder_fillet.stl" scale="0.001 0.001 0.001"/></geometry>
      <material name="hold"><color rgba="0.55 0.58 0.62 1"/></material></visual>
    <visual><origin xyz="0 0 0" rpy="0 0 0"/>
      <geometry><mesh filename="yoke_fillet.stl" scale="0.001 0.001 0.001"/></geometry>
      <material name="yk"><color rgba="0.85 0.60 0.20 1"/></material></visual></link>
  <joint name="explode" type="prismatic">
    <parent link="jig"/><child link="assembly"/>
    <origin xyz="0 0 0" rpy="0 0 0"/><axis xyz="0 0 1"/>
    <limit lower="0" upper="0.06" effort="1" velocity="1"/></joint>
</robot>
''')
print("wrote servo_zero_fixture_install.urdf (explode slider 0..60mm)")