import numpy as np, trimesh
import gen_two_servo as g
GLB="two_servo_from_blender.glb"; GRAV=9.81; M_SERVO=0.060; RHO=1240.0; STALL=2.94; OUT_X=-0.0255
def load(s):
    try: return trimesh.load(s, force="mesh")
    except Exception: return None
def pm(s):
    m=load(s); return m.volume*1e-9*RHO if m else 0.0
g.LAYOUT=g.read_layout(GLB); LAY=g.LAYOUT; N=len(LAY)
servo=load("st3215.stl"); holder=load("holder_fillet.stl"); yoke=load("yoke_fillet.stl")
mh=pm("holder_fillet.stl"); my=pm("yoke_fillet.stl")
scom,hcom,ycom=servo.center_mass,holder.center_mass,yoke.center_mass
print(f"N servos={N}")
print(f"{'part':<20}{'vol cm3':>9}{'mass g':>9}")
print(f"{'servo (datasheet)':<20}{'--':>9}{M_SERVO*1000:>9.0f}")
print(f"{'holder_fillet':<20}{holder.volume/1000:>9.1f}{mh*1000:>9.1f}")
print(f"{'yoke_fillet':<20}{yoke.volume/1000:>9.1f}{my*1000:>9.1f}")
for i in range(1,N):
    c=load(f"connector_beauty_{i}.stl")
    if c is not None: print(f"{'connector_'+str(i):<20}{c.volume/1000:>9.1f}{c.volume*1e-9*RHO*1000:>9.1f}")
mods=[]
for i in range(N):
    parts=[(M_SERVO,scom),(mh,hcom),(my,ycom)]
    if i>=1:
        c=load(f"connector_beauty_{i}.stl")
        if c is not None: parts.append((c.volume*1e-9*RHO,c.center_mass))
    if i==0:
        for bs in ["base_holder_detach.stl","fiducial_mount.stl"]:
            b=load(bs)
            if b is not None: parts.append((b.volume*1e-9*RHO,b.center_mass))
    M=sum(p[0] for p in parts)
    coml=sum(p[0]*np.array(p[1]) for p in parts)/M
    com0=trimesh.transform_points([coml/1000.0],LAY[i])[0]
    mods.append((M,com0))
tot=sum(m for m,_ in mods)
print(f"\ntotal arm mass = {tot*1000:.0f} g")
base=LAY[0][:3,3]
reach=max(np.linalg.norm(LAY[i][:3,3]-base) for i in range(N))*1000
print(f"max servo reach from base = {reach:.0f} mm")
gvec=np.array([0,-GRAV,0])
print(f"\n{'joint':<9}{'downstr g':>10}{'lever mm':>10}{'tau_worst Nm':>14}{'%stall':>8}")
for i in range(N-1):
    axis=LAY[i][:3,:3]@np.array([0,1.0,0]); axis/=np.linalg.norm(axis)
    pj=trimesh.transform_points([[OUT_X,0,0]],LAY[i])[0]
    Md=0; comd=np.zeros(3)
    for j in range(i+1,N):
        M,com=mods[j]; Md+=M; comd+=M*com
    comd/=Md
    d=comd-pj; rperp=np.linalg.norm(d-np.dot(d,axis)*axis)
    tmax=Md*GRAV*rperp
    print(f"output_{i:<3}{Md*1000:>10.0f}{rperp*1000:>10.0f}{tmax:>14.2f}{100*tmax/STALL:>7.0f}%")
print(f"\nSTS3215 stall @12V = {STALL:.2f} Nm (continuous-safe holding ~1/3)")