"""
prepare_meshes.py — Build a MuJoCo-friendly copy of the RX1 mesh tree.

MuJoCo's STL decoder rejects meshes with > 200,000 faces.  A few RX1 visual
meshes exceed that, so we copy every mesh into `meshes_mjcf/` and quadric-
decimate any that are too dense.  Originals under rx1_description/meshes/ are
left untouched.

Run:
    python prepare_meshes.py
"""
from __future__ import annotations

import pathlib
import struct

import trimesh

HERE = pathlib.Path(__file__).resolve().parent
SRC = HERE.parent / "humanoid_torso" / "rx1_description" / "meshes"
DST = HERE / "meshes_mjcf"
MAX_FACES = 180_000          # safety margin under MuJoCo's 200k hard limit


def face_count(path: pathlib.Path) -> int:
    return struct.unpack("<I", path.read_bytes()[80:84])[0]


def main() -> None:
    DST.mkdir(exist_ok=True)
    for src in sorted(SRC.rglob("*.stl")):
        rel = src.relative_to(SRC)
        dst = DST / rel
        dst.parent.mkdir(parents=True, exist_ok=True)

        n = face_count(src)
        if n <= MAX_FACES:
            dst.write_bytes(src.read_bytes())
            print(f"[copy ] {rel}  ({n} faces)")
            continue

        mesh = trimesh.load_mesh(str(src))
        target = int(MAX_FACES * 0.9)
        simplified = mesh.simplify_quadric_decimation(face_count=target)
        simplified.export(str(dst), file_type="stl")        # binary by default
        print(f"[decim] {rel}  {n} -> {len(simplified.faces)} faces")

    print(f"\n[prepare] mesh tree ready at {DST}")


if __name__ == "__main__":
    main()