"""'Curvy' connector variant: instead of a uniform rounded prism, loft the rounded cross-section through
N scaled sections so it stays FULL at both ends (bonds to holder + yoke) but WAISTS IN toward the middle
(smooth concave slope). Separate files (two_servo_curvy.urdf) so it can be discarded."""
import numpy as np
import trimesh
from shapely.geometry import box as sbox
import gen_two_servo as g

GLB = "two_servo_from_blender.glb"
R = 7.0          # corner radius (mm)
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           # loft sections along the length


def _stitch(rings):
    rings = np.asarray(rings)
    Nr, P, _ = rings.shape
    V = rings.reshape(-1, 3)
    F = []
    for k in range(Nr - 1):
        b0, b1 = k * P, (k + 1) * P
        for j in range(P):
            jn = (j + 1) % P
            F += [[b0 + j, b0 + jn, b1 + jn], [b0 + j, b1 + jn, b1 + j]]
    for ring_idx, flip in [(0, False), (Nr - 1, True)]:                # end caps (fan from centroid)
        cen = len(V); V = np.vstack([V, rings[ring_idx].mean(0)]); base = ring_idx * P
        for j in range(P):
            jn = (j + 1) % P
            F.append([cen, base + j, base + jn] if not flip else [cen, base + jn, base + j])
    m = trimesh.Trimesh(V, np.array(F), process=False); 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(GLB)
    g.CONNECTOR_FMT = "connector_curvy_{}.stl"
    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}")
    g.gen_urdf("two_servo_curvy.urdf")