#!/usr/bin/env python3
"""Create lightweight visual STL meshes by deterministic vertex clustering.

The source meshes are left untouched.  This optimization is intended only for
URDF ``visual`` geometry; collision geometry and all physical properties stay
in the URDF.  Vertices are snapped to a regular grid, so the maximum movement
of any source vertex is bounded by ``sqrt(3) * grid_mm / 2``.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import numpy as np
import trimesh


def optimize_mesh(source: Path, destination: Path, grid_mm: float) -> dict[str, object]:
    mesh = trimesh.load_mesh(source, process=True)
    source_vertices = len(mesh.vertices)
    source_faces = len(mesh.faces)

    # Round each vertex to its nearest grid point.  The rounding makes the
    # geometric error bounded and gives deterministic output across runs.
    clustered_vertices = np.rint(mesh.vertices / grid_mm).astype(np.int64)
    unique_grid_vertices, inverse = np.unique(clustered_vertices, axis=0, return_inverse=True)
    faces = inverse[mesh.faces]

    # Remove triangles collapsed by clustering.
    valid = (
        (faces[:, 0] != faces[:, 1])
        & (faces[:, 1] != faces[:, 2])
        & (faces[:, 2] != faces[:, 0])
    )
    faces = faces[valid]

    # Remove duplicate triangles while preserving the winding of the first
    # occurrence.  Sorting is used only to detect duplicates.
    canonical_faces = np.sort(faces, axis=1)
    _, first_indices = np.unique(canonical_faces, axis=0, return_index=True)
    faces = faces[np.sort(first_indices)]

    vertices = unique_grid_vertices.astype(np.float64) * grid_mm
    result = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
    result.remove_unreferenced_vertices()

    destination.parent.mkdir(parents=True, exist_ok=True)
    result.export(destination, file_type="stl")

    return {
        "file": source.name,
        "source_vertices": source_vertices,
        "source_faces": source_faces,
        "optimized_vertices": len(result.vertices),
        "optimized_faces": len(result.faces),
        "source_bytes": source.stat().st_size,
        "optimized_bytes": destination.stat().st_size,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("source_dir", type=Path)
    parser.add_argument("output_dir", type=Path)
    parser.add_argument("--grid-mm", type=float, default=1.0)
    args = parser.parse_args()

    if args.grid_mm <= 0.0:
        raise ValueError("--grid-mm must be positive")

    source_dir = args.source_dir.resolve()
    output_dir = args.output_dir.resolve()
    records = [
        optimize_mesh(source, output_dir / source.name, args.grid_mm)
        for source in sorted(source_dir.glob("*.STL"))
    ]
    if not records:
        raise RuntimeError(f"No STL files found in {source_dir}")

    manifest = {
        "source_dir": str(source_dir),
        "grid_mm": args.grid_mm,
        "maximum_vertex_displacement_mm": math.sqrt(3.0) * args.grid_mm / 2.0,
        "source_faces": sum(int(record["source_faces"]) for record in records),
        "optimized_faces": sum(int(record["optimized_faces"]) for record in records),
        "source_bytes": sum(int(record["source_bytes"]) for record in records),
        "optimized_bytes": sum(int(record["optimized_bytes"]) for record in records),
        "meshes": records,
    }
    (output_dir / "optimization_manifest.json").write_text(
        json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps(manifest, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()