"""Real self-intersection checker for the SO-100 arm (FK + mesh-mesh collision).

The URDF carries only <visual> geometry, so MuJoCo reports no contacts. This parses the URDF, does
forward kinematics at a set of joint poses, places every visual mesh in the world, and runs true
mesh-mesh collision (trimesh CollisionManager / FCL). Reports colliding part pairs, flagging
SERVO<->SERVO overlaps (always unintended) separately from connector<->holder (intended mates).

Usage: so100_check.py <urdf> [q]   (q applied to all revolute joints; default 0)
"""
import sys
import numpy as np
import trimesh
import xml.etree.ElementTree as ET
from scipy.spatial.transform import Rotation as Rot

D = "/data/cameron/repos/smith300_para_stuff/"


def T(xyz, rpy):
    M = np.eye(4)
    M[:3, :3] = Rot.from_euler("xyz", rpy).as_matrix()
    M[:3, 3] = xyz
    return M


def fk_world(root, q):
    joints = []
    for j in root.findall("joint"):
        o = j.find("origin")
        xyz = [float(x) for x in o.get("xyz").split()]
        rpy = [float(x) for x in o.get("rpy").split()]
        axis = j.find("axis")
        ax = [float(x) for x in axis.get("xyz").split()] if axis is not None else [0, 0, 0]
        joints.append((j.get("name"), j.find("parent").get("link"), j.find("child").get("link"),
                       xyz, rpy, ax, j.get("type")))
    children = {c for _, _, c, _, _, _, _ in joints}
    roots = [l.get("name") for l in root.findall("link") if l.get("name") not in children]
    world = {r: np.eye(4) for r in roots}                   # seed every root (URDFs w/o a 'world' link)
    changed = True
    while changed:
        changed = False
        for jn, p, c, xyz, rpy, ax, jt in joints:
            if p in world and c not in world:
                M = world[p] @ T(xyz, rpy)
                if jt == "revolute":
                    ang = q.get(jn, 0.0) if isinstance(q, dict) else q   # per-joint dict OR scalar-all
                    R = np.eye(4)
                    R[:3, :3] = Rot.from_rotvec(np.array(ax) * ang).as_matrix()
                    M = M @ R
                world[c] = M
                changed = True
    return world


def check(urdf, q):
    """q: scalar (all revolute) or dict {joint_name: angle}. Returns (servo_pairs, bridge_servo_pairs)."""
    root = ET.parse(D + urdf).getroot()
    world = fk_world(root, q)
    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:
                sc = [float(x) for x in g.get("scale", "1 1 1").split()]
                m = trimesh.load(D + g.get("filename"), force="mesh").copy()
                m.apply_scale(sc)
                part = g.get("filename").split("/")[-1].split(".")[0]
            elif b is not None:
                ext = [float(x) for x in b.get("size").split()]
                m = trimesh.creation.box(extents=ext)
                mat = v.find("material")
                part = mat.get("name") if mat is not None else f"box{vi}"
            else:
                continue
            m.apply_transform(world[name] @ T(xyz, rpy))
            mgr.add_object(f"{name}:{part}", m)
    hit, pairs = mgr.in_collision_internal(return_names=True)
    servo_pairs = [(a, b) for a, b in pairs if "st3215" in a and "st3215" in b]
    # bridge ("br*") passing through a servo that is NOT its own link's servo (intended holder mate)
    bridge_servo = [(a, b) for a, b in pairs
                    if (("br" in a and "st3215" in b and a.split(":")[0] != b.split(":")[0])
                        or ("br" in b and "st3215" in a and a.split(":")[0] != b.split(":")[0]))]
    return servo_pairs, bridge_servo


def main():
    urdf = sys.argv[1] if len(sys.argv) > 1 else "so100_para_connected.urdf"
    arg = sys.argv[2] if len(sys.argv) > 2 else "0.0"
    q = {kv.split("=")[0].strip(): float(kv.split("=")[1]) for kv in arg.split(",")} if "=" in arg else float(arg)
    servo_pairs, bridge_servo = check(urdf, q)
    print(f"q={q}: SERVO-SERVO={len(servo_pairs)}  BRIDGE-thru-OTHER-SERVO={len(bridge_servo)}")
    for a, b in sorted(servo_pairs):
        print(f"  SERVO-SERVO  {a}  <->  {b}")
    for a, b in sorted(bridge_servo):
        print(f"  BRIDGE-SERVO {a}  <->  {b}")
    return servo_pairs


if __name__ == "__main__":
    main()