"""OUR extrusion mount, designed from scratch (no GEM parts): 2020 V-slot + printed socket block +
REAL fastening: M5x16 socket screws through counterbored block holes, down the slot centerlines, into
drop-in T-nuts in the channel. Screw positions = exactly where you install them. -> my_mount.urdf"""
import numpy as np, trimesh
import xml.etree.ElementTree as ET
from gen_extrusion import extrusion
from gen_fastener import socket_screw, hex_nut, drop_in_tnut

SIZE, CLEAR = 20.0, 0.4
BLOCK, SOCK_DEPTH = 45.0, 38.0
CB_D, CB_DEPTH = 9.2, 6.0
def U(p): return trimesh.boolean.union(p, engine="manifold")
def D(a, b): return trimesh.boolean.difference([a, b], engine="manifold")
def cyl(r, h, axis, at):
    c = trimesh.creation.cylinder(radius=r, height=h, sections=32)
    if axis == "y": c.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2, [1, 0, 0]))
    if axis == "z": pass
    c.apply_translation(at); return c

# BLIND SOCKET layout: block x 0..45, mouth at x=45, channel floor at x=6 (extrusion seats there,
# 6mm above the block outer bottom face at x=0 -- the hard depth stop you feel on insertion)
FLOOR = 6.0
ext = extrusion(20, 300)
ext.apply_transform(trimesh.transformations.rotation_matrix(np.pi/2, [0, 1, 0]))
ext.apply_translation([FLOOR + 150, 0, 0])               # end face seats on the floor at x=FLOOR
ext.export("gnode_my_extrusion.stl")
blk = trimesh.creation.box(extents=[BLOCK, BLOCK, BLOCK])
blk.apply_translation([BLOCK/2, 0, 0])
chan = trimesh.creation.box(extents=[BLOCK - FLOOR + 1, SIZE + 2*CLEAR, SIZE + 2*CLEAR])
chan.apply_translation([FLOOR + (BLOCK - FLOOR + 1)/2, 0, 0])
blk = D(blk, chan)
SCREWS = [("y", 16.0), ("y", 34.0), ("z", 25.0)]          # stations inside the socket span (x 6..45)
for axis, xs in SCREWS:
    x = xs
    if axis == "y":
        blk = D(blk, cyl(2.7, 40, "y", [x, BLOCK/2 - 10, 0]))
        blk = D(blk, cyl(CB_D/2, CB_DEPTH*2, "y", [x, BLOCK/2, 0]))
    else:
        blk = D(blk, cyl(2.7, 40, "z", [x, 0, BLOCK/2 - 10]))
        blk = D(blk, cyl(CB_D/2, CB_DEPTH*2, "z", [x, 0, BLOCK/2]))
blk.export("gnode_my_block.stl")
print(f"block: wt={blk.is_watertight} vol={blk.volume/1000:.1f}cm3")
# screws + T-nuts at exact install positions
scr, tns = [], []
for axis, xs in SCREWS:
    x = xs
    s = socket_screw(5.0, 10.0)   # M5x10: ends in the cavity, no core bottom-out
    t = drop_in_tnut(5.0)
    if axis == "y":
        R = trimesh.transformations.rotation_matrix(np.pi/2, [1, 0, 0])       # +z -> -y
        for m in (s, t): m.apply_transform(R)
        s.apply_translation([x, BLOCK/2 - CB_DEPTH, 0])
        t.apply_translation([x, SIZE/2 - 1.8, 0])                              # nut top at the lip underside
        rot = trimesh.transformations.rotation_matrix(np.pi/2, [0, 0, 1])
    else:
        R = trimesh.transformations.rotation_matrix(np.pi, [1, 0, 0])          # +z -> -z
        for m in (s, t): m.apply_transform(R)
        s.apply_translation([x, 0, BLOCK/2 - CB_DEPTH])
        t.apply_translation([x, 0, SIZE/2 - 1.8])
    scr.append(s); tns.append(t)
trimesh.util.concatenate(scr).export("gnode_my_screws.stl")
trimesh.util.concatenate(tns).export("gnode_my_tnuts.stl")
# URDF
R = ET.Element("robot", {"name": "my_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 0"})
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, "gnode_my_block.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_extrusion", "base", "ext", "1 0 0", 0.1)
vis(l, "gnode_my_extrusion.stl", "0.66 0.68 0.72 1")
l = pjoint("explode_tnuts", "ext", "tnuts", "0 0.7 0.7", 0.04)
vis(l, "gnode_my_tnuts.stl", "0.55 0.42 0.25 1")
l = pjoint("explode_screws", "base", "screws", "0 0.7 0.7", 0.06)
vis(l, "gnode_my_screws.stl", "0.80 0.80 0.83 1")
t = ET.ElementTree(R); ET.indent(t, "  ")
t.write("my_mount.urdf", encoding="unicode", xml_declaration=True)
print("wrote my_mount.urdf")