"""Diagnostic: world centroids + separation of the wrist & gripper servos (and min gap)."""
import sys
import numpy as np
import trimesh
import xml.etree.ElementTree as ET
from so100_check import fk_world, T, D

urdf = sys.argv[1] if len(sys.argv) > 1 else "so100_para_connected.urdf"
q = float(sys.argv[2]) if len(sys.argv) > 2 else 0.0
root = ET.parse(D + urdf).getroot()
world = fk_world(root, q)
meshes = {}
for l in root.findall("link"):
    name = l.get("name")
    if name not in world:
        continue
    for v in l.findall("visual"):
        g = v.find("geometry").find("mesh")
        if g is None or "st3215" not in g.get("filename"):
            continue
        sc = [float(x) for x in g.get("scale", "1 1 1").split()]
        o = v.find("origin")
        xyz = [float(x) for x in o.get("xyz").split()]
        rpy = [float(x) for x in o.get("rpy").split()]
        m = trimesh.load(D + g.get("filename"), force="mesh").copy()
        m.apply_scale(sc)
        m.apply_transform(world[name] @ T(xyz, rpy))
        meshes[name] = m

for a in ["wrist", "gripper", "lower_arm"]:
    if a in meshes:
        print(f"{a:10s} servo centroid (mm): {np.round(meshes[a].bounds.mean(0)*1000,1).tolist()}")
if "wrist" in meshes and "gripper" in meshes:
    cw, cg = meshes["wrist"].bounds.mean(0), meshes["gripper"].bounds.mean(0)
    print(f"wrist->gripper centre separation = {np.round((cg-cw)*1000,1).tolist()} mm  |d|={np.linalg.norm(cg-cw)*1000:.1f}")