"""Parametric build123d connector for the SO-100 modular arm (Phase 2 / structural).

One printable solid per inter-servo link, built directly in the CHILD link frame (mm):
  - COLLAR: cylinder around the PARENT servo output axis at the link origin (grips the driven horn).
  - BEAM:   thick rod from the parent output (origin) to THIS link's servo BODY (so it grips the body
            it carries; the servo OUTPUT stays free to drive the next joint).
  - BOSS:   cylinder at the servo body, embracing the holder so the connection actually mates.
Union -> watertight solid -> STL (mm). Length adapts per joint. Thick by design (bears load).

Target is the servo BODY CENTRE (not the output point): the output horn sticks ~18mm beyond the
holder, so a beam to the output would leave a gap. Body centre guarantees the connector grips the holder.
"""
import math
import numpy as np
from build123d import Vector, Cylinder, Cone, Axis, export_stl
from so100_para import JOINTS, ROLLS, servo_pose          # reuse the exact servo placement

HOLDER_C = np.array([2.8, -9.6, 0.0]) / 1000.0           # wrap_holder bbox centre in servo frame (m)
COLLAR_R, COLLAR_H, BOSS_R, BOSS_H = 15.0, 30.0, 16.0, 22.0
BEAM_R_ROOT, BEAM_R_TIP = 13.0, 9.0   # tapered beam: thick at the collar root (max bending), thinner at boss
CHILD_OF = {1: "shoulder", 2: "upper_arm", 3: "lower_arm", 4: "wrist", 5: "gripper"}


def _orient(solid, axis_vec):
    z = Vector(0, 0, 1)
    a = Vector(*axis_vec).normalized()
    ang = math.degrees(math.acos(max(-1.0, min(1.0, z.dot(a)))))
    ax = z.cross(a)
    if ax.length < 1e-9:
        ax = Vector(1, 0, 0)
    return solid.rotate(Axis((0, 0, 0), tuple(ax)), ang)


def _cone(p1, p2, r1, r2):
    """Truncated cone with radius r1 at p1 -> r2 at p2 (load-graded beam)."""
    p1, p2 = Vector(*p1), Vector(*p2)
    v = p2 - p1
    cone = _orient(Cone(r1, r2, v.length), tuple(v))      # r1 at -h/2 (p1), r2 at +h/2 (p2)
    return cone.translate(((p1.X + p2.X) / 2, (p1.Y + p2.Y) / 2, (p1.Z + p2.Z) / 2))


def _cyl_at(center, axis, r, h):
    return _orient(Cylinder(r, h), axis).translate(tuple(center))


def build(k):
    """k = child link index (1..5); connector built in child link k frame (mm)."""
    _, _, _, xyz, rpy, axis, lo, hi = JOINTS[k]
    pa = JOINTS[k - 1][5]                                  # parent output axis (child frame)
    R_s, t_s = servo_pose(xyz, rpy, axis, ROLLS[k])        # this link's servo pose (child frame, m)
    B = (R_s.apply(HOLDER_C) + t_s) * 1000.0               # holder centre (child frame, mm) -> true mate
    A = (0.0, 0.0, 0.0)
    collar = _cyl_at(A, pa, COLLAR_R, COLLAR_H)           # solid hub gripping the parent output
    beam = _cone(A, tuple(B), BEAM_R_ROOT, BEAM_R_TIP)    # load-graded: thick root -> thinner tip
    u = Vector(*B).normalized()
    boss = _cyl_at(tuple(B), tuple(u), BOSS_R, BOSS_H)    # mates into the holder
    return collar + beam + boss


if __name__ == "__main__":
    for k, child in CHILD_OF.items():
        solid = build(k)
        f = f"so100_conn_{child}.stl"
        export_stl(solid, f)
        print(f"wrote {f}  ({'valid' if solid.is_valid else 'INVALID'})")