"""Rubber-band gravity comp for the ARM shoulder (Cameron 2026-07-07): clip-on hook above output_1 on
servo_1 + clip-on hook on connector_2 (distal), elastic band between them. Stretched rubber approximates a
zero-free-length spring -> assist ~ sin(theta) like gravity. Writes hook STLs + two_servo_gravcomp.urdf."""
import numpy as np, trimesh
import xml.etree.ElementTree as ET
import gen_two_servo as g

GRAV=9.81; M_SERVO=0.060; RHO=1240.0; OUT_X=-0.0255
g.LAYOUT = g.read_layout("two_servo_from_blender.glb"); LAY = g.LAYOUT; N = len(LAY)
def load(f): return trimesh.load(f, force="mesh")
def box(x0,x1,y0,y1,z0,z1):
    b=trimesh.creation.box(extents=[x1-x0,y1-y0,z1-z0]); b.apply_translation([(x0+x1)/2,(y0+y1)/2,(z0+z1)/2]); return b
def cylb(p,q,r,sec=24):
    p=np.array(p,float); q=np.array(q,float); v=q-p; L=np.linalg.norm(v)
    c=trimesh.creation.cylinder(radius=r,height=L,sections=sec); zc=np.array([0,0,1.]); n=v/L
    ax=np.cross(zc,n); s=np.linalg.norm(ax)
    if s>1e-9: c.apply_transform(trimesh.transformations.rotation_matrix(np.arctan2(s,np.dot(zc,n)),ax/s))
    c.apply_translation((p+q)/2); return c
def U(p): return trimesh.boolean.union(p,engine="manifold")
def D(a,b): return trimesh.boolean.difference([a,b],engine="manifold")
def to_frame(mesh,k):
    m=mesh.copy(); m.apply_transform(np.linalg.inv(g._mm(LAY[k]))); return m

# shoulder torque (recompute)
mh=load("holder_fillet.stl").volume*1e-9*RHO; my=load("yoke_fillet.stl").volume*1e-9*RHO
def modmass(i):
    m=M_SERVO+mh+my
    try: m+=load(f"connector_beauty_{i}.stl").volume*1e-9*RHO
    except Exception: pass
    return m if i>0 else M_SERVO+mh+my
up=np.array([0,1.0,0])
ax1=LAY[1][:3,:3]@up; ax1/=np.linalg.norm(ax1)
pj=trimesh.transform_points([[OUT_X,0,0]],LAY[1])[0]*1000            # output_1 (mm, servo_0 frame)
Md=sum(modmass(j) for j in range(2,N))
com=sum(modmass(j)*LAY[j][:3,3] for j in range(2,N))/Md
d=com*1000-pj; rperp=np.linalg.norm(d-np.dot(d,ax1)*ax1)
tau=Md*GRAV*rperp/1000.0
# proximal hook: the module top is BELOW the axis here, so a short riser post lifts the eye above it.
mod1 = np.vstack([trimesh.transform_points(load(f).vertices/1000.0, LAY[1])*1000
                  for f in ("st3215.stl", "holder_fillet.stl")])
top_y = mod1[:,1].max()
near = mod1[np.abs(mod1[:,1]-top_y) < 3.0]                            # top surface points
i = np.argmin((near[:,0]-pj[0])**2 + (near[:,2]-pj[2])**2)            # closest to above-the-axis
foot = near[i].copy()
H_EYE = 30.0                                                           # riser height above module top
P = np.array([foot[0], top_y + H_EYE, foot[2]])
h = P[1] - pj[1]                                                       # effective vertical lever
# distal hook: collar on connector_2 at its centroid
c2 = load("connector_beauty_2.stl")
A = trimesh.transform_points([c2.centroid/1000.0], LAY[2])[0]*1000
a = np.linalg.norm(A-pj)
k_needed = tau/((h/1000.0)*(a/1000.0))
Lband = np.linalg.norm(P-A)
F_horiz = k_needed*Lband/1000.0
print(f"shoulder tau={tau:.2f} Nm | h={h:.0f}mm a={a:.0f}mm -> elastic k={k_needed:.0f} N/m, band length@rest-pose {Lband:.0f}mm, tension@horizontal ~{F_horiz:.1f} N")
# hooks: saddle clip over holder rim (prox) + collar clip on connector_2 (distal), each with a hook eye
s2=LAY[1][:3,3]*1000; s3=LAY[2][:3,3]*1000; cdir=(s3-s2)/np.linalg.norm(s3-s2)
prox = U([box(foot[0]-10,foot[0]+10,top_y-16,top_y+3,foot[2]-10,foot[2]+10),      # saddle base on module top
          cylb([foot[0],top_y,foot[2]], P, 5),                                     # riser post
          D(cylb(P,P+[0,8,0],7), cylb(P,P+[0,9,0],4))])                            # eye ring
prox = D(prox, box(foot[0]-7.5,foot[0]+7.5,top_y-17,top_y+0.6,foot[2]-7.5,foot[2]+7.5))  # clip slot
col = box(A[0]-12,A[0]+12,A[1]-16,A[1]+16,A[2]-12,A[2]+12)
col = D(col, cylb(A-30*cdir, A+30*cdir, 10.5))                                      # bore for connector_2
dist = U([col, D(cylb(A+[0,18,0],A+[0,26,0],7), cylb(A+[0,18,0],A+[0,27,0],4))])
Ae = A+[0,22,0]
band = cylb(P+[0,6,0], Ae, 2.5)                                                     # elastic viz (rest pose)
to_frame(prox,1).export("gc_hook_prox.stl")
to_frame(dist,2).export("gc_hook_dist.stl")
to_frame(band,1).export("gc_band.stl")
print(f"hooks: prox wt={prox.is_watertight} dist wt={dist.is_watertight}")
# URDF copy
t=ET.parse("two_servo_detach.urdf"); root=t.getroot()
links={l.get("name"):l for l in root.findall("link")}
for fn,parent,rgba in [("gc_hook_prox.stl","servo_1","0.85 0.35 0.25 1"),
                        ("gc_hook_dist.stl","servo_2","0.85 0.35 0.25 1"),
                        ("gc_band.stl","servo_1","0.95 0.75 0.15 1")]:
    v=ET.SubElement(links[parent],"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":"m"+fn}),"color",{"rgba":rgba})
ET.indent(t,"  "); t.write("two_servo_gravcomp.urdf",encoding="unicode",xml_declaration=True)
print("wrote two_servo_gravcomp.urdf")