"""ALIGNED variant (separate 'branch'): snap the layout so connectors go orthogonal instead of slightly
diagonal. Orientations -> nearest 90deg; position offsets between consecutive servos with |component| <
POS_THRESH_MM -> 0 (kept: the big intentional offsets). Rounded connectors. Writes two_servo_aligned.urdf
+ connector_aligned_*.stl. Does NOT touch the default arm files."""
import numpy as np
import trimesh
from scipy.spatial.transform import Rotation as R
import gen_two_servo as g

GLB = "two_servo_from_blender.glb"
POS_THRESH_MM = 15.0                         # off-axis offsets smaller than this are placement error -> snap to 0


def snap_layout(lay):
    snapped = [lay[0].copy()]
    prev = lay[0][:3, 3].copy()
    for i in range(1, len(lay)):
        M = lay[i].copy()
        rpy = R.from_matrix(M[:3, :3]).as_euler("xyz", degrees=True)
        M[:3, :3] = R.from_euler("xyz", np.round(rpy / 90.0) * 90.0, degrees=True).as_matrix()   # -> nearest 90
        off_mm = (lay[i][:3, 3] - prev) * 1000.0
        off_mm[np.abs(off_mm) < POS_THRESH_MM] = 0.0                                              # kill small offsets
        prev = prev + off_mm / 1000.0
        M[:3, 3] = prev
        snapped.append(M)
    return snapped


if __name__ == "__main__":
    g.LAYOUT = snap_layout(g.read_layout(GLB))
    g.CONNECTOR_FMT = "connector_aligned_{}.stl"
    for i in range(1, len(g.LAYOUT)):
        c = g.connector_mesh(i)
        c.export(f"connector_aligned_{i}.stl")
        p = g.LAYOUT[i][:3, 3] * 1000
        print(f"servo_{i}: snapped pos={np.round(p,1).tolist()}  connector vol={c.volume/1000:.1f}cm3")
    g.gen_urdf("two_servo_aligned.urdf")