"""Rigid transform mapping OUR servo mesh (st3215.stl) -> smith300's (sts3215_03a_v1.stl),
via multi-init ICP (try all 24 axis-aligned rotations as starts, ICP-refine each, keep best).
Outputs the rpy + translation to apply (composed with smith300's per-servo pose).
"""
import itertools
import numpy as np
import trimesh
from scipy.spatial import cKDTree
from scipy.spatial.transform import Rotation as Rot

D = "/data/cameron/repos/smith300_para_stuff"
ours = trimesh.load(f"{D}/st3215.stl", force="mesh")
ours.apply_scale(0.001)                                    # mm -> m
smith = trimesh.load(f"{D}/assets/sts3215_03a_v1.stl", force="mesh")
print("ours bbox(m)", np.round(ours.extents, 4), " smith bbox(m)", np.round(smith.extents, 4))

po = ours.sample(3000).astype(float)
ps = smith.sample(6000).astype(float)
oc, sc = po.mean(0), ps.mean(0)
tree_s = cKDTree(ps)


def kabsch(P, Q):                                          # R,t with R@P+t ~= Q
    pm, qm = P.mean(0), Q.mean(0)
    H = (P - pm).T @ (Q - qm)
    U, _, Vt = np.linalg.svd(H)
    d = np.sign(np.linalg.det(Vt.T @ U.T))
    R = Vt.T @ np.diag([1, 1, d]) @ U.T
    return R, qm - R @ pm


def icp(R, t, iters=50):
    for _ in range(iters):
        d, idx = tree_s.query(po @ R.T + t)
        R, t = kabsch(po, ps[idx])
    d, _ = tree_s.query(po @ R.T + t)
    return R, t, d.mean()


def cube_rotations():
    out = []
    for perm in itertools.permutations(range(3)):
        for sg in itertools.product([1, -1], repeat=3):
            M = np.zeros((3, 3))
            for i, p in enumerate(perm):
                M[i, p] = sg[i]
            if abs(np.linalg.det(M) - 1) < 1e-6:
                out.append(M)
    return out


results = []
for M in cube_rotations():
    R, t, cost = icp(M, sc - M @ oc)
    results.append((cost, R, t))
results.sort(key=lambda r: r[0])
print("\ntop 3 ICP fits (lower = better):")
for cost, R, t in results[:3]:
    print(f"  cost={cost*1000:.3f}mm  rpy={np.round(Rot.from_matrix(R).as_euler('XYZ'),4).tolist()}  t(m)={np.round(t,4).tolist()}")
cost, R, t = results[0]
rpy = Rot.from_matrix(R).as_euler("XYZ")
print(f"\nBEST (cost {cost*1000:.3f}mm):")
print(f"APPLY  _FLIP rpy(XYZ,rad) = [{rpy[0]:.6f}, {rpy[1]:.6f}, {rpy[2]:.6f}]")
print(f"APPLY  _T translation (m) = [{t[0]:.6f}, {t[1]:.6f}, {t[2]:.6f}]")