"""Smith300-FAITHFUL modular arm from OUR LOCKED parts (wrap_holder + yoke_exact + st3215).

CORRECT approach (Cameron 2026-06-24), NOT the rejected hack:
  - Every joint rotates about OUR servo's OWN output axis (+Y through (-25.5,0)mm) -> every servo
    is mounted physically sensibly (output shaft IS the joint).
  - We reproduce smith300's joint-axis RELATIONSHIPS (how each DOF is oriented vs the previous) by
    CONJUGATING smith300's joint rpy into our +Y convention:  R_our = C^-1 @ R_smith @ C,
    C = Rx(90deg) (our +Y -> smith300 +Z). See compute_bends.py (verified: consecutive joint-axis
    angles match smith300 exactly: 90,0,0,90,90,0 deg). smith300 is only a GUIDE for DOF variety.
  - Connection geometry is the clean MODULAR chain (yoke grips a servo's output and reaches the next
    servo seat). Whole arm = f(holder, yoke, connector, chain). No mesh registration, no smith300 mesh.
"""
import json
import sys
import xml.etree.ElementTree as ET

HOLDER, YOKE, SERVO = "wrap_holder.stl", "yoke_arm.stl", "st3215.stl"
SCALE = "0.001 0.001 0.001"

# --- connector geometry (parametric; f(yoke link length)) ---
ARM_YOKE_LEN = 45.0                              # mm; compact inter-servo connector
OUT_X = -0.0255                                  # servo base -> its output axis (joint), m
SERVO_OFF = -(34.7 + ARM_YOKE_LEN) / 1000.0      # output frame -> next servo base, m
YOKE_OFF = -OUT_X                                # render the yoke back in the gripped servo's frame

# --- chain: per-joint BENDS reproducing smith300's joint-axis relationships in OUR +Y convention ---
# (from compute_bends.py: R_our = C^-1 @ R_smith @ C). Joint axis stays +Y (our real output shaft).
BENDS = [
    (-0.00000, +0.00000, -3.14159),   # j1 <- smith300 (pi,0,pi)
    (-1.57080, +0.00000, +1.57080),   # j2 <- smith300 (0,-pi/2,-pi/2)   [perpendicular: 90deg]
    (+3.14159, +1.57079, +3.14159),   # j3 <- smith300 (0,0,pi/2)        [collinear]
    (+0.00000, +0.00000, +0.00000),   # j4 <- smith300 (0,0,0)           [collinear]
    (+1.57080, +0.00000, +0.00000),   # j5 <- smith300 (pi/2,0,0)        [perpendicular: 90deg]
    (+3.14159, +0.00000, -1.57079),   # j6 <- smith300 (0,pi/2,-pi)      [perpendicular: 90deg]
    (+0.00000, +0.00000, +0.00000),   # j7 <- smith300 (0,0,0)           [collinear -> wrist]
]
N = len(BENDS)

COLORS = {"servo": "0.32 0.32 0.36 1", "holder": "0.20 0.52 0.85 1", "yoke": "0.85 0.60 0.20 1"}
ROOT_RPY = (3.14159, 0, 0)                        # base plate sits on the table

# fiducial (ArUco) exoskeleton — faithful to medium_single_redo: base 3x3@50mm id0, links 2x2@15mm.
FIDEX = True
ARUCO_BASE_M, ARUCO_LINK_M = 0.050, 0.015
FIDEX_CFG = []


def _box(link, xyz, size, rgba, name):
    v = ET.SubElement(link, "visual")
    ET.SubElement(v, "origin", {"xyz": xyz, "rpy": "0 0 0"})
    ET.SubElement(ET.SubElement(v, "geometry"), "box", {"size": size})
    ET.SubElement(ET.SubElement(v, "material", {"name": name + str(id(link))}), "color", {"rgba": rgba})


def _vis(link, fname, cname, xyz):
    v = ET.SubElement(link, "visual")
    ET.SubElement(v, "origin", {"xyz": xyz, "rpy": "0 0 0"})
    ET.SubElement(ET.SubElement(v, "geometry"), "mesh", {"filename": fname, "scale": SCALE})
    ET.SubElement(ET.SubElement(v, "material", {"name": cname + str(id(link))}), "color", {"rgba": COLORS[cname]})


def _aruco(link, link_name, sx, board_m, grid, ids):
    if not FIDEX:
        return
    sx = float(sx)
    zf = 0.0158
    board_z = zf + 0.0038
    hstl = "fidex_holder_bg.stl" if board_m > 0.03 else "fidex_holder_sm.stl"
    v = ET.SubElement(link, "visual")
    ET.SubElement(v, "origin", {"xyz": f"{sx} 0 {zf}", "rpy": "0 0 0"})
    ET.SubElement(ET.SubElement(v, "geometry"), "mesh", {"filename": hstl, "scale": SCALE})
    ET.SubElement(ET.SubElement(v, "material", {"name": "fxh" + str(id(link))}), "color", {"rgba": "0.16 0.16 0.18 1"})
    _box(link, f"{sx} 0 {board_z}", f"{board_m:.4f} {board_m:.4f} 0.0016", "0.96 0.96 0.96 1", "fxmark")
    FIDEX_CFG.append({
        "link": link_name, "mujoco_name": link_name,
        "robot_mesh_path": HOLDER, "exo_mesh_path": hstl,
        "aruco_offset_pos": [round(sx * 1000, 2), 0.0, round(board_z * 1000, 2)],
        "aruco_offset_rot": [0.0, 0.0, 0.0],
        "aruco_board_name": f"board_{link_name}", "board_length": round(board_m, 4),
        "grid": grid, "ids": ids, "dict": "DICT_4X4_250",
        "marker_image": f"fidex_markers/{link_name}.png",
    })


def gen_urdf():
    robot = ET.Element("robot", {"name": "smith300_faithful"})
    ET.SubElement(robot, "link", {"name": "world"})
    bj = ET.SubElement(robot, "joint", {"name": "base_fix", "type": "fixed"})
    ET.SubElement(bj, "parent", {"link": "world"})
    ET.SubElement(bj, "child", {"link": "L0"})
    ET.SubElement(bj, "origin", {"xyz": "0 0 0", "rpy": f"{ROOT_RPY[0]} {ROOT_RPY[1]} {ROOT_RPY[2]}"})

    L0 = ET.SubElement(robot, "link", {"name": "L0"})
    _box(L0, "-0.01 0 -0.024", "0.075 0.075 0.006", "0.90 0.75 0.15 1", "baseplate")
    _vis(L0, SERVO, "servo", "0 0 0")
    _vis(L0, HOLDER, "holder", "0 0 0")
    _aruco(L0, "L0", "0", ARUCO_BASE_M, [3, 3], list(range(9)))

    prev = "L0"
    for i in range(1, N + 1):
        jx = OUT_X if i == 1 else (SERVO_OFF + OUT_X)
        br, bp, by = BENDS[i - 1]
        j = ET.SubElement(robot, "joint", {"name": f"j{i}", "type": "revolute"})
        ET.SubElement(j, "parent", {"link": prev})
        ET.SubElement(j, "child", {"link": f"L{i}"})
        ET.SubElement(j, "origin", {"xyz": f"{jx} 0 0", "rpy": f"{br} {bp} {by}"})
        ET.SubElement(j, "axis", {"xyz": "0 1 0"})           # OUR servo output shaft = the joint
        ET.SubElement(j, "limit", {"lower": "-2", "upper": "2", "effort": "2.94", "velocity": "3.0"})

        Li = ET.SubElement(robot, "link", {"name": f"L{i}"})
        _vis(Li, YOKE, "yoke", f"{YOKE_OFF} 0 0")
        _vis(Li, SERVO, "servo", f"{SERVO_OFF} 0 0")
        _vis(Li, HOLDER, "holder", f"{SERVO_OFF} 0 0")
        base_id = 20 + (i - 1) * 4
        _aruco(Li, f"L{i}", f"{SERVO_OFF}", ARUCO_LINK_M, [2, 2], [base_id + k for k in range(4)])
        prev = f"L{i}"

    # actuated 2-finger gripper at the wrist
    gx = SERVO_OFF + OUT_X
    gj = ET.SubElement(robot, "joint", {"name": "gripper_mount", "type": "fixed"})
    ET.SubElement(gj, "parent", {"link": prev})
    ET.SubElement(gj, "child", {"link": "gripper"})
    ET.SubElement(gj, "origin", {"xyz": f"{gx} 0 0", "rpy": "0 0 0"})
    gl = ET.SubElement(robot, "link", {"name": "gripper"})
    _box(gl, "-0.012 0 0", "0.022 0.038 0.024", "0.85 0.60 0.20 1", "gpalm")
    for fname, z0 in [("finger_l", 0.018), ("finger_r", -0.018)]:
        fj = ET.SubElement(robot, "joint", {"name": fname, "type": "prismatic"})
        ET.SubElement(fj, "parent", {"link": "gripper"})
        ET.SubElement(fj, "child", {"link": fname})
        ET.SubElement(fj, "origin", {"xyz": f"-0.026 0 {z0}", "rpy": "0 0 0"})
        ET.SubElement(fj, "axis", {"xyz": f"0 0 {-1 if z0 > 0 else 1}"})
        ET.SubElement(fj, "limit", {"lower": "0", "upper": "0.018", "effort": "5", "velocity": "0.1"})
        fl = ET.SubElement(robot, "link", {"name": fname})
        _box(fl, "-0.018 0 0", "0.036 0.014 0.006", "0.30 0.30 0.34 1", "fp_" + fname)
    return robot


if __name__ == "__main__":
    tree = ET.ElementTree(gen_urdf())
    ET.indent(tree, "  ")
    tree.write("smith300_faithful.urdf", encoding="unicode", xml_declaration=True)
    with open("fidex_config_faithful.json", "w") as f:
        json.dump(FIDEX_CFG, f, indent=2)
    print(f"wrote smith300_faithful.urdf (N={N} joints, our +Y output as every joint) + "
          f"{len(FIDEX_CFG)} fidex boards")
