"""Optimize ROLLS[k] (free roll about each servo output axis) so each servo's YOKE points toward the
NEXT servo's holder -> short, clean bridges. Roll doesn't change the joint axis or servo position, so
SO-100 kinematics + poses are preserved; only the body/yoke orientation about the output spins.

Coordinate descent: ROLLS[k] affects gap(k-1) (servo_k is the target of servo_{k-1}'s yoke) and gap(k)
(servo_k's yoke aims at servo_{k+1}). Minimize their sum, sweep, repeat to convergence.
"""
import numpy as np
from so100_para import JOINTS, servo_pose, _T, _mat, LINKS

YOKE_TIP = np.array([-0.118, 0.0, 0.0])
HOLDER_C = np.array([0.0028, -0.0096, 0.0])
NK = len(JOINTS)                                             # servos 0..NK-1 (parametric)


def gap(k, rolls):
    """Distance (m) from servo_k's yoke tip to servo_{k+1}'s holder, in LINKS[k+1] frame."""
    if k + 1 >= NK:
        return 0.0
    pj = JOINTS[k]
    Rp, tp = servo_pose(pj[3], pj[4], pj[5], rolls[k])
    P = np.linalg.inv(_T(pj[3], pj[4])) @ _mat(Rp, tp)
    tip = (P @ np.append(YOKE_TIP, 1))[:3]
    cj = JOINTS[k + 1]
    Rc, tc = servo_pose(cj[3], cj[4], cj[5], rolls[k + 1])
    holder = Rc.apply(HOLDER_C) + tc
    return np.linalg.norm(tip - holder)


def cost_of(k, rolls):
    c = gap(k, rolls)
    if k >= 1:
        c += gap(k - 1, rolls)
    return c


def main():
    rolls = [0.0] * NK
    grid = np.linspace(0, 2 * np.pi, 144, endpoint=False)
    for sweep in range(6):
        for k in range(NK):
            best, br = 1e9, rolls[k]
            for r in grid:
                rolls[k] = r
                c = cost_of(k, rolls)
                if c < best:
                    best, br = c, r
            rolls[k] = br
    print("optimized ROLLS (rad):", [round(r, 4) for r in rolls])
    print("per-link yoke->holder gap (mm) after optimization:")
    for k in range(NK - 1):
        print(f"  {LINKS[k]:10s}-> {LINKS[k+1]:10s}: {gap(k, rolls)*1000:6.1f}")
    return rolls


if __name__ == "__main__":
    main()