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]); gv=np.array([0,-GRAV,0])
# pick shoulder = highest-torque joint
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))   # 1 = horizontal axis (pitch)
    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}")
# attach point A on connector_3 (world)
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=|joint->attach|={a*1000:.0f} mm")
def spec(name,off,offdir_desc):
    P=pj+off*up
    k=tau/(off*a)                 # zero-free-length: k*h*a = tau
    # ROM sweep: rotate A about the joint axis +-70deg, spring length s, force F=k*s
    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 {offdir_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  (travel {(max(ss)-min(ss))*1000:.0f} mm)")
    print(f"  spring force over ROM: {min(Fs):.1f} -> {max(Fs):.1f} N  ({max(Fs)/9.81*1000:.0f} gf peak)")
    print(f"  needs zero-free-length behavior (pretension/cable-pulley) for exact balance at all angles")
    return P,k
Po,ko=spec("OVERHEAD POST",0.20,"above shoulder (fixed to world)")
Pi,ki=spec("SHOULDER-INTEGRATED",0.08,"tower on base link (rotates with arm)")
print(f"\nNote: system is TUNABLE post-build (adjust anchor height or preload to zero out sag).")