import numpy as np, trimesh
import xml.etree.ElementTree as ET
import gen_two_servo as g
GLB="two_servo_from_blender.glb"; OUT_X=-0.0255
g.LAYOUT=g.read_layout(GLB); LAY=g.LAYOUT
def load(s): return trimesh.load(s,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=20):
    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):                                        # servo_0-frame (mm) mesh -> servo_k frame
    m=mesh.copy(); m.apply_transform(np.linalg.inv(g._mm(LAY[k]))); return m

pj=trimesh.transform_points([[OUT_X,0,0]],LAY[1])[0]*1000
c3=load("connector_beauty_3.stl"); A=trimesh.transform_points([c3.center_mass/1000.0],LAY[3])[0]*1000
up=np.array([0,1.0,0]); Po=pj+200*up; Pi=pj+80*up
s2=LAY[2][:3,3]*1000; s3=LAY[3][:3,3]*1000; cdir=(s3-s2)/np.linalg.norm(s3-s2)

def arm_clip(at):
    col=box(at[0]-15,at[0]+15,at[1]-13,at[1]+13,at[2]-11,at[2]+11)
    col=D(col,cylb(at-30*cdir,at+30*cdir,8))
    peg=cylb(at+np.array([0,16,-12]),at+np.array([0,16,12]),3)
    post=box(at[0]-4,at[0]+4,at[1]+9,at[1]+18,at[2]-4,at[2]+4)
    return U([col,post,peg])
def hook_block(at):
    return U([box(at[0]-12,at[0]+12,at[1]-8,at[1]+8,at[2]-10,at[2]+10),
              cylb(at+np.array([0,10,-14]),at+np.array([0,10,14]),3)])

clip=arm_clip(A)                                             # -> rides servo_3 (connector_3, distal)
# OVERHEAD gantry (ground-fixed -> servo_0)
zoff=110
gantry=U([cylb([pj[0],-32,zoff],[pj[0],Po[1]+8,zoff],8),
          cylb([pj[0],Po[1]+8,zoff],[Po[0],Po[1]+8,Po[2]],8),
          hook_block(np.array([Po[0],Po[1],Po[2]])),
          D(box(pj[0]-20,pj[0]+20,-34,-20,zoff-18,zoff+18),cylb([pj[0],-40,zoff],[pj[0],-15,zoff],8.5))])
spring_o=cylb(Po,A,4)
# SHOULDER tower (proximal -> servo_1)
tower=U([box(pj[0]-14,pj[0]+14,pj[1]-6,pj[1]+2,pj[2]-14,pj[2]+14),
         cylb([pj[0],pj[1],pj[2]],[Pi[0],Pi[1],Pi[2]],7),
         hook_block(np.array([Pi[0],Pi[1],Pi[2]]))])
spring_i=cylb(Pi,A,4)

# export in each part's PARENT frame
to_frame(gantry,0).export("cb_gantry.stl"); to_frame(spring_o,0).export("cb_spring_o.stl")
to_frame(tower,1).export("cb_tower.stl");   to_frame(spring_i,1).export("cb_spring_i.stl")
to_frame(clip,3).export("cb_clip.stl")
print("tower->servo_1, clip->servo_3, gantry->servo_0(ground)")

def build(parts,out):   # parts=[(fn,parent,rgba),...]
    t=ET.parse("two_servo_beauty.urdf"); root=t.getroot()
    links={l.get("name"):l for l in root.findall("link")}
    for fn,parent,rgba in parts:
        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":fn}),"color",{"rgba":rgba})
    ET.indent(t,"  "); t.write(out,encoding="unicode",xml_declaration=True); print("wrote",out)
build([("cb_gantry.stl","servo_0","0.55 0.6 0.65 1"),("cb_spring_o.stl","servo_0","0.95 0.45 0.1 1"),("cb_clip.stl","servo_3","0.55 0.6 0.65 1")],"two_servo_cb_overhead.urdf")
build([("cb_tower.stl","servo_1","0.55 0.6 0.65 1"),("cb_spring_i.stl","servo_1","0.95 0.45 0.1 1"),("cb_clip.stl","servo_3","0.55 0.6 0.65 1")],"two_servo_cb_shoulder.urdf")