#!/usr/bin/env python3
"""Generate simplified Gen2 collision proxies and URDF variants.

The script does not modify the original robot.urdf. It writes:

- source/gen2_lab/assets/collision_simplified/*.stl
- source/gen2_lab/assets/robot_simplified_collision.urdf
- source/gen2_lab/assets/robot_simplified_collision_mesh.urdf
- source/gen2_lab/assets/collision_simplified_audit.csv

The recommended URDF is robot_simplified_collision.urdf. It uses URDF box
primitives for collision. The generated STL files and mesh URDF are useful for
visual overlay checks in tools such as Blender.
"""

from __future__ import annotations

import csv
import math
import struct
import xml.etree.ElementTree as ET
from pathlib import Path
from xml.dom import minidom


REPO_ROOT = Path(__file__).resolve().parents[1]
ASSET_DIR = REPO_ROOT / "source" / "gen2_lab" / "assets"
INPUT_URDF = ASSET_DIR / "robot.urdf"
OUTPUT_DIR = ASSET_DIR / "collision_simplified"
OUTPUT_PRIMITIVE_URDF = ASSET_DIR / "robot_simplified_collision.urdf"
OUTPUT_MESH_URDF = ASSET_DIR / "robot_simplified_collision_mesh.urdf"
AUDIT_CSV = ASSET_DIR / "collision_simplified_audit.csv"


def _parse_floats(text: str | None, default: tuple[float, float, float]) -> tuple[float, float, float]:
    if not text:
        return default
    values = [float(v) for v in text.split()]
    if len(values) != 3:
        raise ValueError(f"Expected 3 floats, got: {text}")
    return values[0], values[1], values[2]


def _rpy_to_matrix(rpy: tuple[float, float, float]) -> list[list[float]]:
    roll, pitch, yaw = rpy
    cr, sr = math.cos(roll), math.sin(roll)
    cp, sp = math.cos(pitch), math.sin(pitch)
    cy, sy = math.cos(yaw), math.sin(yaw)
    return [
        [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
        [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
        [-sp, cp * sr, cp * cr],
    ]


def _mat_vec_mul(mat: list[list[float]], vec: tuple[float, float, float]) -> tuple[float, float, float]:
    return tuple(sum(mat[i][j] * vec[j] for j in range(3)) for i in range(3))


def _fmt(values: tuple[float, ...] | list[float]) -> str:
    return " ".join(f"{v:.9g}" for v in values)


def _binary_stl_bbox(path: Path) -> tuple[tuple[float, float, float], tuple[float, float, float], int]:
    data = path.read_bytes()
    if len(data) < 84:
        raise ValueError(f"STL too small: {path}")
    tri_count = struct.unpack("<I", data[80:84])[0]
    if 84 + 50 * tri_count != len(data):
        raise ValueError(f"Only binary STL is supported by this script: {path}")

    mins = [float("inf"), float("inf"), float("inf")]
    maxs = [float("-inf"), float("-inf"), float("-inf")]
    offset = 84
    for _ in range(tri_count):
        values = struct.unpack("<12f", data[offset : offset + 48])
        for start in (3, 6, 9):
            for axis in range(3):
                value = values[start + axis]
                mins[axis] = min(mins[axis], value)
                maxs[axis] = max(maxs[axis], value)
        offset += 50

    extent = tuple(maxs[i] - mins[i] for i in range(3))
    center = tuple((maxs[i] + mins[i]) * 0.5 for i in range(3))
    return extent, center, tri_count


def _link_scale(name: str) -> tuple[float, float, float]:
    if name == "body_link":
        return 0.74, 0.74, 0.74
    if name.startswith("waist_link"):
        return 0.22, 0.22, 0.22
    if name.endswith("leg_link_6"):
        return 0.90, 0.80, 0.45
    if name.endswith("leg_link_4"):
        return 0.40, 0.44, 0.58
    if name.endswith("leg_link_5"):
        return 0.42, 0.42, 0.50
    if "leg_link_2" in name or "leg_link_3" in name:
        return 0.30, 0.30, 0.30
    if "leg_link_1" in name:
        return 0.22, 0.22, 0.22
    if "arm_link_7" in name:
        return 0.20, 0.20, 0.20
    if "arm_link_5" in name or "arm_link_6" in name:
        return 0.22, 0.22, 0.22
    if "arm_link_1" in name or "arm_link_2" in name:
        return 0.30, 0.30, 0.30
    if "arm_link" in name:
        return 0.36, 0.36, 0.36
    return 0.46, 0.46, 0.46


def _proxy_box(name: str, extent: tuple[float, float, float], center: tuple[float, float, float]):
    scale = _link_scale(name)
    size = tuple(max(extent[i] * scale[i], 0.01) for i in range(3))
    proxy_center = center

    # Keep the sole close to the original mesh bottom while reducing foot height.
    if name.endswith("leg_link_6"):
        min_z = center[2] - extent[2] * 0.5
        size = (max(extent[0] * 0.90, 0.08), max(extent[1] * 0.80, 0.04), min(max(extent[2] * 0.45, 0.025), 0.045))
        proxy_center = (center[0], center[1], min_z + size[2] * 0.5 + 0.002)

    return proxy_center, size, scale


def _normal(a, b, c):
    ux, uy, uz = b[0] - a[0], b[1] - a[1], b[2] - a[2]
    vx, vy, vz = c[0] - a[0], c[1] - a[1], c[2] - a[2]
    nx, ny, nz = uy * vz - uz * vy, uz * vx - ux * vz, ux * vy - uy * vx
    length = math.sqrt(nx * nx + ny * ny + nz * nz)
    if length <= 1e-12:
        return 0.0, 0.0, 0.0
    return nx / length, ny / length, nz / length


def _write_binary_stl(path: Path, triangles: list[tuple[tuple[float, float, float], ...]]) -> None:
    header = b"generated simplified collision proxy".ljust(80, b" ")
    with path.open("wb") as file:
        file.write(header)
        file.write(struct.pack("<I", len(triangles)))
        for tri in triangles:
            n = _normal(*tri)
            file.write(struct.pack("<3f", *n))
            for vertex in tri:
                file.write(struct.pack("<3f", *vertex))
            file.write(struct.pack("<H", 0))


def _box_triangles(center: tuple[float, float, float], size: tuple[float, float, float]):
    cx, cy, cz = center
    sx, sy, sz = (size[0] * 0.5, size[1] * 0.5, size[2] * 0.5)
    v = [
        (cx - sx, cy - sy, cz - sz),
        (cx + sx, cy - sy, cz - sz),
        (cx + sx, cy + sy, cz - sz),
        (cx - sx, cy + sy, cz - sz),
        (cx - sx, cy - sy, cz + sz),
        (cx + sx, cy - sy, cz + sz),
        (cx + sx, cy + sy, cz + sz),
        (cx - sx, cy + sy, cz + sz),
    ]
    faces = [
        (0, 2, 1),
        (0, 3, 2),
        (4, 5, 6),
        (4, 6, 7),
        (0, 1, 5),
        (0, 5, 4),
        (1, 2, 6),
        (1, 6, 5),
        (2, 3, 7),
        (2, 7, 6),
        (3, 0, 4),
        (3, 4, 7),
    ]
    return [(v[a], v[b], v[c]) for a, b, c in faces]


def _replace_collision_with_box(collision: ET.Element, origin_xyz, origin_rpy, size) -> None:
    for child in list(collision):
        collision.remove(child)

    origin = ET.SubElement(collision, "origin")
    origin.set("xyz", _fmt(origin_xyz))
    origin.set("rpy", _fmt(origin_rpy))
    geometry = ET.SubElement(collision, "geometry")
    box = ET.SubElement(geometry, "box")
    box.set("size", _fmt(size))


def _replace_collision_with_mesh(collision: ET.Element, mesh_path: str) -> None:
    geometry = collision.find("geometry")
    if geometry is None:
        geometry = ET.SubElement(collision, "geometry")
    for child in list(geometry):
        geometry.remove(child)
    mesh = ET.SubElement(geometry, "mesh")
    mesh.set("filename", mesh_path)


def _write_pretty_xml(tree: ET.ElementTree, path: Path) -> None:
    rough = ET.tostring(tree.getroot(), encoding="utf-8")
    pretty = minidom.parseString(rough).toprettyxml(indent="  ")
    lines = [line for line in pretty.splitlines() if line.strip()]
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def main() -> None:
    if not INPUT_URDF.exists():
        raise FileNotFoundError(INPUT_URDF)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    for stale_proxy in OUTPUT_DIR.glob("*_box_collision.stl"):
        stale_proxy.unlink()

    primitive_tree = ET.parse(INPUT_URDF)
    mesh_tree = ET.parse(INPUT_URDF)
    audit_rows = []

    mesh_links = {link.attrib["name"]: link for link in mesh_tree.getroot().findall("link")}

    for primitive_link in primitive_tree.getroot().findall("link"):
        name = primitive_link.attrib["name"]
        collision = primitive_link.find("collision")
        mesh_link = mesh_links[name]
        mesh_collision = mesh_link.find("collision")
        if collision is None or mesh_collision is None:
            continue

        mesh = collision.find("./geometry/mesh")
        origin = collision.find("origin")
        if mesh is None or origin is None:
            continue

        mesh_file = mesh.attrib["filename"]
        extent, center, tri_count = _binary_stl_bbox(ASSET_DIR / mesh_file)
        proxy_center, size, scale = _proxy_box(name, extent, center)

        original_xyz = _parse_floats(origin.attrib.get("xyz"), (0.0, 0.0, 0.0))
        original_rpy = _parse_floats(origin.attrib.get("rpy"), (0.0, 0.0, 0.0))
        rotated_center = _mat_vec_mul(_rpy_to_matrix(original_rpy), proxy_center)
        proxy_origin_xyz = tuple(original_xyz[i] + rotated_center[i] for i in range(3))

        stl_name = f"{name}_box_collision.stl"
        stl_rel_path = f"collision_simplified/{stl_name}"
        _write_binary_stl(OUTPUT_DIR / stl_name, _box_triangles(proxy_center, size))

        _replace_collision_with_box(collision, proxy_origin_xyz, original_rpy, size)
        _replace_collision_with_mesh(mesh_collision, stl_rel_path)

        audit_rows.append(
            {
                "link": name,
                "source_mesh": mesh_file,
                "proxy_mesh": stl_rel_path,
                "source_triangles": tri_count,
                "source_extent_x": extent[0],
                "source_extent_y": extent[1],
                "source_extent_z": extent[2],
                "proxy_center_x": proxy_center[0],
                "proxy_center_y": proxy_center[1],
                "proxy_center_z": proxy_center[2],
                "proxy_size_x": size[0],
                "proxy_size_y": size[1],
                "proxy_size_z": size[2],
                "scale_x": scale[0],
                "scale_y": scale[1],
                "scale_z": scale[2],
            }
        )

    if not audit_rows:
        raise RuntimeError(f"No mesh collisions were converted from {INPUT_URDF}")

    _write_pretty_xml(primitive_tree, OUTPUT_PRIMITIVE_URDF)
    _write_pretty_xml(mesh_tree, OUTPUT_MESH_URDF)

    with AUDIT_CSV.open("w", newline="", encoding="utf-8") as file:
        writer = csv.DictWriter(file, fieldnames=list(audit_rows[0].keys()))
        writer.writeheader()
        writer.writerows(audit_rows)

    print(f"wrote {OUTPUT_PRIMITIVE_URDF}")
    print(f"wrote {OUTPUT_MESH_URDF}")
    print(f"wrote {AUDIT_CSV}")
    print(f"wrote {len(audit_rows)} collision STL files under {OUTPUT_DIR}")


if __name__ == "__main__":
    main()