"""Print ALL colliding part pairs at a given single-joint config. Usage: so100_at.py <joint> <angle>."""
import sys
import trimesh
import xml.etree.ElementTree as ET
from so100_check import fk_world, T, D

urdf = "so100_para_connected.urdf"
jn, ang = sys.argv[1], float(sys.argv[2])
root = ET.parse(D + urdf).getroot()
world = fk_world(root, {jn: ang})
mgr = trimesh.collision.CollisionManager()
for l in root.findall("link"):
    name = l.get("name")
    if name not in world:
        continue
    for vi, v in enumerate(l.findall("visual")):
        o = v.find("origin")
        xyz = [float(x) for x in (o.get("xyz").split() if o is not None else [0, 0, 0])]
        rpy = [float(x) for x in (o.get("rpy").split() if o is not None else [0, 0, 0])]
        geom = v.find("geometry")
        g, b = geom.find("mesh"), geom.find("box")
        if g is not None:
            m = trimesh.load(D + g.get("filename"), force="mesh").copy()
            m.apply_scale([float(x) for x in g.get("scale", "1 1 1").split()])
            part = g.get("filename").split("/")[-1].split(".")[0]
        elif b is not None:
            m = trimesh.creation.box(extents=[float(x) for x in b.get("size").split()])
            part = (v.find("material").get("name") if v.find("material") is not None else f"box{vi}")
        else:
            continue
        m.apply_transform(world[name] @ T(xyz, rpy))
        mgr.add_object(f"{name}:{part}", m)
_, pairs = mgr.in_collision_internal(return_names=True)
cross = [(a, b) for a, b in pairs if a.split(":")[0] != b.split(":")[0]]
print(f"{jn}={ang}: {len(cross)} CROSS-LINK pairs (of {len(pairs)} total)")
for a, b in sorted(cross):
    print(f"  {a}  <->  {b}")