"""Curvy 'beauty' connectors: loft a rounded-rect cross-section along the connector axis, blending from the
HOLDER face size down to the YOKE-arm size, with a waist pinch mid-span. EXTRUDE_FRAC=0.5 = reach only
halfway toward the yoke (Cameron 2026-07-02).

RECOVERY NOTE (2026-07-06): reconstructed from session context. Constants + _rrect + curvy_connector are
verbatim; _stitch was REWRITTEN from its documented behavior (stitch rings -> watertight loft) — verify
against the lab original when phe108 is back."""
import numpy as np
import trimesh
from shapely.geometry import box as sbox
import gen_two_servo as g

R = 7.0          # cross-section corner radius
WAIST = 0.30     # fraction the cross-section narrows at the middle (0 = uniform)
WAIST_BY_CONN = {3: 0.425}   # connector_3: halfway between default 0.30 and 0.55
EXTRUDE_FRAC = 0.5           # how far the connector reaches out toward the yoke (1.0 = old; 0.5 = half)
N = 16


def _stitch(rings):
    """Stitch a list of same-length rings (K points each, matched order) into a closed watertight mesh:
    side quads between consecutive rings + triangle-fan end caps. [RECONSTRUCTED]"""
    rings = [np.asarray(r, float) for r in rings]
    K = len(rings[0])
    V = np.vstack(rings)
    F = []
    for i in range(len(rings) - 1):
        a0, b0 = i * K, (i + 1) * K
        for j in range(K):
            j2 = (j + 1) % K
            F += [[a0 + j, a0 + j2, b0 + j2], [a0 + j, b0 + j2, b0 + j]]
    c0, c1 = len(V), len(V) + 1                       # cap centroids
    V = np.vstack([V, rings[0].mean(0), rings[-1].mean(0)])
    last = (len(rings) - 1) * K
    for j in range(K):
        j2 = (j + 1) % K
        F.append([c0, j2, j])                          # start cap (reversed winding)
        F.append([c1, last + j, last + j2])            # end cap
    m = trimesh.Trimesh(V, np.array(F), process=True)
    m.fix_normals()
    return m


def _rrect(a, b, r, nc=8):
    """Rounded-rectangle outline with a FIXED point count (4*nc) so lofted rings always align."""
    r = max(1e-3, min(r, a - 1e-3, b - 1e-3))
    pts = []
    for cx, cy, a0 in [(a - r, b - r, 0.0), (-(a - r), b - r, np.pi / 2),
                       (-(a - r), -(b - r), np.pi), (a - r, -(b - r), 3 * np.pi / 2)]:
        for k in range(nc):
            ang = a0 + (np.pi / 2) * k / (nc - 1)
            pts.append((cx + r * np.cos(ang), cy + r * np.sin(ang)))
    return np.asarray(pts)


def curvy_connector(i):
    yoke = trimesh.load("yoke_exact.stl", force="mesh")
    holder = trimesh.load("holder_approx_drilled.stl", force="mesh")
    Ty, Th = g._mm(g.LAYOUT[i - 1]), g._mm(g.LAYOUT[i])
    arm = yoke.vertices[yoke.vertices[:, 0] < g.YOKE_ARM_X]
    yfaces, hfaces = g._box_faces(arm.min(0), arm.max(0)), g._box_faces(*holder.bounds)
    a, b = g._closest_faces(yfaces, hfaces, Ty, Th)
    hw = trimesh.transform_points(hfaces[b][1], Th); hc = hw.mean(0)
    e1, e2 = hw[1] - hw[0], hw[3] - hw[0]
    h1, h2 = np.linalg.norm(e1) / 2, np.linalg.norm(e2) / 2
    e1u, e2u = e1 / np.linalg.norm(e1), e2 / np.linalg.norm(e2)
    yc = trimesh.transform_points([yfaces[a][0]], Ty)[0]
    bdir = (yc - hc) / np.linalg.norm(yc - hc)
    proj = trimesh.transform_points(arm, Ty) @ bdir
    near = float(proj.min()) - float(hc @ bdir)                        # connector first reaches the arm tip here
    far = float(proj.max()) - float(hc @ bdir)                         # ...and the yoke-head end here
    depth = near + (far - near + 4.0) * EXTRUDE_FRAC                   # how far it reaches out toward the yoke
    aend = trimesh.transform_points(yfaces[a][1], Ty)                  # yoke arm END face (world)
    ya1 = ((aend @ e1u).max() - (aend @ e1u).min()) / 2 + 1.5          # arm half-extents in e1/e2 (+margin to engulf)
    ya2 = ((aend @ e2u).max() - (aend @ e2u).min()) / 2 + 1.5
    w = WAIST_BY_CONN.get(i, WAIST)
    rings = []
    for k in range(N):
        t = k / (N - 1)
        pinch = 1.0 - w * np.sin(np.pi * t)
        aa = (h1 * (1 - t) + ya1 * t) * pinch                          # holder face -> yoke arm cross-section (+ waist)
        bb = (h2 * (1 - t) + ya2 * t) * pinch
        p2 = _rrect(aa, bb, min(R, aa - 0.5, bb - 0.5))
        rings.append(hc + bdir * depth * t + np.outer(p2[:, 0], e1u) + np.outer(p2[:, 1], e2u))
    conn = _stitch(rings)
    conn.apply_transform(g._mm(np.linalg.inv(g.LAYOUT[i])))
    return conn


if __name__ == "__main__":
    g.LAYOUT = g.read_layout("two_servo_from_blender.glb")
    for i in range(1, len(g.LAYOUT)):
        c = curvy_connector(i)
        c.export(f"connector_curvy_{i}.stl")
        print(f"connector_curvy_{i}: vol={c.volume/1000:.1f}cm3 wt={c.is_watertight}")