"""Counterbalance spring math for the arm. Zero-free-length principle: spring anchored h above the joint,
attached a out on the arm, with force ∝ length gives torque k*h*a*sin(theta) = same shape as gravity's
m*g*L*sin(theta). Set k*h*a = tau_gravity -> balanced at every angle (does NOT fight downward motion).
Computed 2026-07-04 for our layout: shoulder=output_1 tau=1.63Nm, a=261mm;
overhead h=200mm -> k=31 N/m but 6x length swing (needs cable+pulley);
shoulder-integrated h=80mm -> k=78 N/m (0.45 lbf/in), 181-288mm travel — a REAL spring works. RECOMMENDED."""
import numpy as np, trimesh
from scipy.spatial.transform import Rotation as R
import gen_two_servo as g
GLB="two_servo_from_blender.glb"; GRAV=9.81; M_SERVO=0.060; RHO=1240.0; OUT_X=-0.0255
g.LAYOUT=g.read_layout(GLB); LAY=g.LAYOUT; N=len(LAY)
def load(s):
    try: return trimesh.load(s,force="mesh")
    except Exception: return None
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; c=load(f"connector_beauty_{i}.stl")
    if i>=1 and c is not None: m+=c.volume*1e-9*RHO
    return m
def modcom(i): return LAY[i][:3,3]
up=np.array([0,1.0,0])
best=None
for i in range(N-1):
    ax=LAY[i][:3,:3]@up; ax/=np.linalg.norm(ax)
    pj=trimesh.transform_points([[OUT_X,0,0]],LAY[i])[0]
    Md=sum(modmass(j) for j in range(i+1,N))
    com=sum(modmass(j)*modcom(j) for j in range(i+1,N))/Md
    d=com-pj; rperp=np.linalg.norm(d-np.dot(d,ax)*ax); tau=Md*GRAV*rperp
    horiz=1-abs(np.dot(ax,up))
    if best is None or tau>best[1]: best=(i,tau,ax,pj,Md,com,rperp,horiz)
i,tau,ax,pj,Md,com,rperp,horiz=best
print(f"SHOULDER = output_{i}: tau_worst={tau:.2f} Nm, downstream={Md*1000:.0f} g, axis-horizontality={horiz:.2f}")
c3=load("connector_beauty_3.stl"); A=trimesh.transform_points([c3.center_mass/1000.0],LAY[3])[0]
a=np.linalg.norm(A-pj)
print(f"joint@{np.round(pj*1000).astype(int)} mm  connector_3 attach@{np.round(A*1000).astype(int)} mm  a={a*1000:.0f} mm")
def spec(name,off,desc):
    P=pj+off*up
    k=tau/(off*a)
    Fs=[]; ss=[]
    for th in np.linspace(-70,70,29):
        Ar=pj+R.from_rotvec(np.radians(th)*ax).apply(A-pj)
        s=np.linalg.norm(P-Ar); ss.append(s); Fs.append(k*s)
    print(f"\n[{name}] anchor {off*1000:.0f} mm {desc}")
    print(f"  spring rate k = {k:.0f} N/m  ({k*0.00571:.2f} lbf/in)")
    print(f"  spring length over +-70deg ROM: {min(ss)*1000:.0f} -> {max(ss)*1000:.0f} mm")
    print(f"  spring force over ROM: {min(Fs):.1f} -> {max(Fs):.1f} N")
spec("OVERHEAD POST",0.20,"above shoulder (fixed to world)")
spec("SHOULDER-INTEGRATED",0.08,"tower on servo_1 (rotates with arm)")