"""Compare overall envelope (q=0 world bounding box of all visual geometry) of our arm vs the real
SO-100 reference — a quantitative resemblance check for the 'loosely resembles SO-100' goal.
"""
import numpy as np
import trimesh
import xml.etree.ElementTree as ET
from so100_check import fk_world, T


def envelope(urdf_path, base):
    root = ET.parse(urdf_path).getroot()
    world = fk_world(root, 0.0)
    lo, hi = np.array([1e9] * 3), np.array([-1e9] * 3)
    for l in root.findall("link"):
        name = l.get("name")
        if name not in world:
            continue
        for v in 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(f"{base}/{g.get('filename')}", force="mesh").copy()
                m.apply_scale([float(x) for x in g.get("scale", "1 1 1").split()])
            elif b is not None:
                m = trimesh.creation.box(extents=[float(x) for x in b.get("size").split()])
            else:
                continue
            m.apply_transform(world[name] @ T(xyz, rpy))
            lo, hi = np.minimum(lo, m.bounds[0]), np.maximum(hi, m.bounds[1])
    return (hi - lo) * 1000.0, lo * 1000.0, hi * 1000.0


D = "/data/cameron/repos/smith300_para_stuff"
SO = "/data/cameron/repos/SO-ARM100/Simulation/SO100"
for label, urdf, base in [("OURS    ", f"{D}/so100_para_connected.urdf", D),
                          ("SO-100  ", f"{SO}/so100.urdf", SO)]:
    size, lo, hi = envelope(urdf, base)
    print(f"{label} q=0 envelope WxDxH (mm) = {np.round(size,0).tolist()}   diag={np.linalg.norm(size):.0f}")