"""GEM extrusion-mount visualization (files from joeclinton.me/gem, CC BY-NC-SA): the two printed
clamp mouths (mount_base + motor0_mount, in their shared CAD frame) + the square tube they sandwich.
NOTE: sockets MEASURE 30x30 (docs say 20x20 — geometry wins). Explode sliders show installation.
Representative M5 screw+nut sets (their exact clamp-hole positions TBD from the lab print)."""
import numpy as np, trimesh
import xml.etree.ElementTree as ET

def box(e, p):
    b = trimesh.creation.box(extents=e); b.apply_translation(p); return b
from gen_extrusion import square_tube
tube = square_tube(30, 2.0, 280)                         # HOLLOW SQUARE TUBE (GEM design intent: through-bolts pass thin walls)
tube.apply_transform(trimesh.transformations.rotation_matrix(np.pi / 2, [0, 1, 0]))
tube.apply_translation([240, 1.5, 0])
tube.export("gnode_gem_tube.stl")
from gen_fastener import socket_screw, hex_nut
def screwset(x):
    """Two REAL M5x45 socket screws entering from -y, hex nuts on +y (measured hole positions z=+-10)."""
    Ry = trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0])   # screw +z -> +y
    parts = []
    for z in (10.0, -10.0):
        s = socket_screw(5.0, 45.0); s.apply_transform(Ry); s.apply_translation([x, -22.5, z])
        n = hex_nut(5.0); n.apply_transform(Ry); n.apply_translation([x, 22.5 + 4.7, z])
        parts += [s, n]
    return trimesh.util.concatenate(parts)
screwset(119.0).export("gnode_gem_screws_base.stl")
screwset(331.0).export("gnode_gem_screws_top.stl")
R = ET.Element("robot", {"name": "gem_mount"}); ET.SubElement(R, "link", {"name": "world"})
wj = ET.SubElement(R, "joint", {"name": "world_to_base", "type": "fixed"})
ET.SubElement(wj, "parent", {"link": "world"}); ET.SubElement(wj, "child", {"link": "base"})
ET.SubElement(wj, "origin", {"xyz": "0 0 0", "rpy": "0 0 -1.5708"})
def vis(link, fn, rgba):
    v = ET.SubElement(link, "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": "m" + fn}), "color", {"rgba": rgba})
base = ET.SubElement(R, "link", {"name": "base"})
vis(base, "gem_mount/mount_base.stl", "0.25 0.55 0.85 1")
def pjoint(name, parent, child, axis, upper):
    j = ET.SubElement(R, "joint", {"name": name, "type": "prismatic"})
    ET.SubElement(j, "parent", {"link": parent}); ET.SubElement(j, "child", {"link": child})
    ET.SubElement(j, "origin", {"xyz": "0 0 0", "rpy": "0 0 0"}); ET.SubElement(j, "axis", {"xyz": axis})
    ET.SubElement(j, "limit", {"lower": "0", "upper": str(upper), "effort": "1", "velocity": "1"})
    return ET.SubElement(R, "link", {"name": child})
l = pjoint("explode_tube", "base", "tube", "1 0 0", 0.12)
vis(l, "gnode_gem_tube.stl", "0.62 0.65 0.70 1")
l = pjoint("explode_top_mouth", "tube", "top_mouth", "1 0 0", 0.1)
vis(l, "gem_mount/motor0_mount.stl", "0.30 0.70 0.45 1")
l = pjoint("explode_screws_base", "base", "screws_base", "0 1 0", 0.06)
vis(l, "gnode_gem_screws_base.stl", "0.75 0.75 0.78 1")
l = pjoint("explode_screws_top", "top_mouth", "screws_top", "0 1 0", 0.06)
vis(l, "gnode_gem_screws_top.stl", "0.75 0.75 0.78 1")
t = ET.ElementTree(R); ET.indent(t, "  ")
t.write("gem_mount.urdf", encoding="unicode", xml_declaration=True)
print("wrote gem_mount.urdf")