"""'Pretty' arm variant: same layout + parts, but the connectors get a ROUNDED-RECTANGLE cross-section
(rounded corners + softened look) instead of the hard square prism. Writes two_servo_pretty.urdf +
connector_pretty_*.stl. Does NOT change the default arm files."""
import os
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"
PRETTY_R = 7.0                                   # corner radius (mm)


def pretty_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)            # holder face corners (world)
    hc = hw.mean(0)
    e1, e2 = hw[1] - hw[0], hw[3] - hw[0]                      # in-plane edge vectors
    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)                # holder -> yoke direction
    depth = float((trimesh.transform_points(arm, Ty) @ bdir).max() - hc @ bdir) + 4.0
    r = float(min(PRETTY_R, h1 - 1.0, h2 - 1.0))
    poly = sbox(-h1 + r, -h2 + r, h1 - r, h2 - r).buffer(r, join_style=1, resolution=8)   # rounded rect
    prism = trimesh.creation.extrude_polygon(poly, 1.0)       # local prism z:0->1
    M = np.eye(4)
    M[:3, 0], M[:3, 1], M[:3, 2], M[:3, 3] = e1u, e2u, bdir * depth, hc   # local -> world
    prism.apply_transform(M)
    prism.apply_transform(g._mm(np.linalg.inv(g.LAYOUT[i])))  # world -> servo_i frame
    return prism


if __name__ == "__main__":
    g.LAYOUT = g.read_layout(GLB)
    g.CONNECTOR_FMT = "connector_pretty_{}.stl"
    for i in range(1, len(g.LAYOUT)):
        c = pretty_connector(i)
        c.export(f"connector_pretty_{i}.stl")
        print(f"connector_pretty_{i}: vol={c.volume/1000:.1f}cm3 wt={c.is_watertight}")
    g.gen_urdf("two_servo_pretty.urdf")