"""NEMA17 + mount + T8 leadscrew/nut-block + MGN12 rails + GT2 pulleys + GT2 belt loops. -> gantry_arm_motorized.urdf"""
import numpy as np, trimesh
import xml.etree.ElementTree as ET
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 cylZ(x,y,r,z0,z1):
    c=trimesh.creation.cylinder(radius=r,height=z1-z0,sections=40); c.apply_translation([x,y,(z0+z1)/2]); return c
def cylY(x,z,r,y0,y1):
    c=trimesh.creation.cylinder(radius=r,height=y1-y0,sections=40); c.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2,[1,0,0])); c.apply_translation([x,(y0+y1)/2,z]); return c
def U(p): return trimesh.boolean.union(p,engine="manifold")
def D(a,b): return trimesh.boolean.difference([a,b],engine="manifold")
trimesh.util.concatenate([box(-21.15,21.15,-21.15,21.15,-40,0),cylZ(0,0,11,0,2),cylZ(0,0,2.5,0,24)]).export("nema17.stl")
plate=box(-25,25,-25,25,0,5); plate=D(plate,cylZ(0,0,11.5,-1,6))
for x in (-15.5,15.5):
    for y in (-15.5,15.5): plate=D(plate,cylZ(x,y,1.7,-1,6))
foot=box(-25,25,-25,-20,-34,5)
for fx in (-16,16): foot=D(foot,cylY(fx,-14,2.6,-26,-19))
U([plate,foot]).export("nema17_mount.stl")
cylY(0,0,4,-500,0).export("leadscrew_t8.stl")
nb=box(-12,12,-7,7,-12,12); nb=D(nb,cylY(0,0,4.1,-8,8))
for hx in (-8,8): nb=D(nb,cylZ(hx,0,1.7,-13,13))
nb.export("nut_block.stl")
box(-6,6,-4,4,-250,250).export("mgn12_rail.stl")
trimesh.util.concatenate([cylZ(0,0,6,-4,4),cylZ(0,0,8,-5,-4),cylZ(0,0,8,4,5)]).export("gt2_pulley.stl")
trimesh.util.concatenate([box(-280,280,997,1003,5,7),box(-280,280,997,1003,-7,-5)]).export("belt_x.stl")
trimesh.util.concatenate([box(5,7,-48,-42,-270,270),box(-7,-5,-48,-42,-270,270)]).export("belt_z.stl")
print("meshes ok")
t=ET.parse("gantry_arm.urdf"); root=t.getroot(); links={l.get("name"):l for l in root.findall("link")}
def add(parent,fn,xyz,rpy,rgba):
    v=ET.SubElement(links[parent],"visual"); ET.SubElement(v,"origin",{"xyz":xyz,"rpy":rpy})
    ET.SubElement(ET.SubElement(v,"geometry"),"mesh",{"filename":fn,"scale":"0.001 0.001 0.001"})
    ET.SubElement(ET.SubElement(v,"material",{"name":"g"+fn+xyz}),"color",{"rgba":rgba})
MOT="0.15 0.15 0.17 1"; MNT="0.3 0.6 0.45 1"; LS="0.75 0.72 0.5 1"; RL="0.55 0.55 0.6 1"; PL="0.2 0.2 0.22 1"; BT="0.07 0.07 0.07 1"
add("base_link","nema17.stl","0.30 1.0 0","0 -1.5708 0",MOT); add("base_link","nema17_mount.stl","0.275 1.0 0","0 -1.5708 0",MNT)
add("x_carriage","nema17.stl","0 -0.045 0.29","0 3.14159 0",MOT); add("z_carriage","nema17.stl","0 0.02 0","1.5708 0 0",MOT)
add("z_carriage","leadscrew_t8.stl","0 0 0","0 0 0",LS); add("y_carriage","nut_block.stl","0 0 0","0 0 0",MNT)
add("base_link","mgn12_rail.stl","0 0.965 0","0 1.5708 0",RL); add("x_carriage","mgn12_rail.stl","0 -0.08 0","0 0 0",RL); add("z_carriage","mgn12_rail.stl","0.03 -0.25 0","1.5708 0 0",RL)
add("base_link","gt2_pulley.stl","0.28 1.0 0","0 1.5708 0",PL); add("base_link","gt2_pulley.stl","-0.28 1.0 0","0 1.5708 0",PL)
add("x_carriage","gt2_pulley.stl","0 -0.045 0.27","0 0 0",PL); add("x_carriage","gt2_pulley.stl","0 -0.045 -0.27","0 0 0",PL)
add("base_link","belt_x.stl","0 0 0","0 0 0",BT); add("x_carriage","belt_z.stl","0 0 0","0 0 0",BT)
ET.indent(t,"  "); t.write("gantry_arm_motorized.urdf",encoding="unicode",xml_declaration=True); print("wrote motorized")
# --- integrate real printed parts: carriage plates + wrist bracket ---
CP="0.4 0.55 0.45 1"
add("x_carriage","carriage_plate.stl","0 -0.02 0","1.5708 0 0",CP)
add("z_carriage","carriage_plate.stl","0 -0.02 0","1.5708 0 0",CP)
add("y_carriage","gantry_wrist_bracket.stl","0 -0.02 0","0 0 0",CP)
ET.indent(t,"  "); t.write("gantry_arm_motorized.urdf",encoding="unicode",xml_declaration=True); print("re-wrote with printed parts")