import numpy as np, trimesh
import shapely.vectorized as sv
from shapely.ops import unary_union
import gen_two_servo as g
GLB="two_servo_from_blender.glb"; GRAV=9.81; M_SERVO=0.060; RHO=1240.0; E=3500.0; YIELD=50.0; BACKLASH_DEG=0.5
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
holder=load("holder_fillet.stl"); yoke=load("yoke_fillet.stl")
mh=holder.volume*1e-9*RHO; my=yoke.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 secprops(mesh,origin,normal):
    sec=mesh.section(plane_origin=origin,plane_normal=normal)
    if sec is None: return None
    try: p2,_=sec.to_planar()
    except Exception: return None
    polys=list(p2.polygons_full)
    if not polys: return None
    poly=unary_union(polys); minx,miny,maxx,maxy=poly.bounds
    n=160; gx,gy=np.meshgrid(np.linspace(minx,maxx,n),np.linspace(miny,maxy,n))
    da=((maxx-minx)/(n-1))*((maxy-miny)/(n-1))
    m=sv.contains(poly,gx.ravel(),gy.ravel()); pts=np.c_[gx.ravel()[m],gy.ravel()[m]]
    if len(pts)<8: return None
    cen=pts.mean(0); d=pts-cen
    return len(pts)*da, min(np.sum(d[:,1]**2)*da,np.sum(d[:,0]**2)*da), np.sqrt((d**2).sum(1)).max()
print(f"{'conn':<7}{'len mm':>7}{'minA mm2':>9}{'Imin mm4':>10}{'load g':>8}{'arm mm':>8}{'sig MPa':>9}{'SF':>6}{'flex deg':>9}{'wob mm':>8}")
tipwob_conn=0
for i in range(1,N):
    mesh=load(f"connector_beauty_{i}.stl")
    if mesh is None: continue
    # axis from the LAYOUT: servo_{i-1} direction in servo_i frame
    pprev=(np.linalg.inv(LAY[i])@LAY[i-1])[:3,3]*1000.0
    axis=pprev/np.linalg.norm(pprev); L=np.linalg.norm(pprev)
    o0=mesh.vertices.mean(0); t=(mesh.vertices-o0)@axis
    best=None
    for f in np.linspace(0.15,0.85,16):
        s=t.min()+f*(t.max()-t.min()); p=secprops(mesh,o0+s*axis,axis)
        if p and (best is None or p[0]<best[0]): best=p
    if best is None: continue
    A,Imin,cmax=best
    Mc=sum(modmass(j) for j in range(i,N))
    coms=np.array([LAY[j][:3,3] for j in range(i,N)]); ms=np.array([modmass(j) for j in range(i,N)])
    com=(coms*ms[:,None]).sum(0)/ms.sum(); arm=np.linalg.norm(com-LAY[i-1][:3,3])*1000
    M=Mc*GRAV*arm; sig=M*cmax/Imin; th=M*L/(E*Imin); wob=th*arm
    tipwob_conn+=wob
    print(f"conn_{i:<2}{L:>7.0f}{A:>9.0f}{Imin:>10.0f}{Mc*1000:>8.0f}{arm:>8.0f}{sig:>9.1f}{YIELD/sig:>6.1f}{np.degrees(th):>9.2f}{wob:>8.2f}")
# wobble budget: servo backlash per joint x distance to tip
print(f"\n--- WOBBLE BUDGET (tip) ---")
tip=LAY[N-1][:3,3]; tb=0
for i in range(N-1):
    d=np.linalg.norm(tip-LAY[i][:3,3])*1000; w=np.radians(BACKLASH_DEG)*d; tb+=w
    print(f"output_{i} backlash {BACKLASH_DEG}deg x {d:.0f}mm arm = {w:.1f}mm")
print(f"total servo-backlash tip wobble ~ {tb:.0f} mm  (@ {BACKLASH_DEG}deg/joint)")
print(f"total connector-flex tip wobble ~ {tipwob_conn:.0f} mm")
print(f"PLA E=3500 yield~50MPa; SF>3 desirable. backlash est {BACKLASH_DEG}deg/joint.")