mesh_utils.py

"""
Mesh Processing Utilities
=========================
Advanced mesh repair, manifold enforcement, scaling, and modular
decomposition utilities for 3D-print-ready output.
"""

import re
import logging
from typing import Optional

import numpy as np
import trimesh

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Manifold Repair
# ---------------------------------------------------------------------------
def make_manifold(mesh: trimesh.Trimesh, max_iterations: int = 5) -> trimesh.Trimesh:
    """
    Aggressively repair a mesh to make it manifold (watertight).
    Runs multiple repair passes until the mesh is watertight or max iterations hit.
    """
    for i in range(max_iterations):
        if mesh.is_watertight:
            logger.info("Mesh is manifold after %d passes", i)
            return mesh

        # Fix winding and normals
        trimesh.repair.fix_normals(mesh)
        trimesh.repair.fix_inversion(mesh)
        trimesh.repair.fix_winding(mesh)

        # Remove degenerate faces
        mask = mesh.nondegenerate_faces()
        if not mask.all():
            mesh.update_faces(mask)

        # Merge close vertices (precision issues)
        mesh.merge_vertices(merge_tex=True, merge_norm=True)

        # Fill holes
        mesh.fill_holes()

    logger.info("Mesh repair done (watertight: %s) after %d passes",
                mesh.is_watertight, max_iterations)
    return mesh


def reduce_faces(mesh: trimesh.Trimesh, target_faces: int) -> trimesh.Trimesh:
    """Decimate a mesh to a target face count while preserving shape."""
    if len(mesh.faces) <= target_faces:
        return mesh

    try:
        # Use quadric decimation if available
        simplified = mesh.simplify_quadric_decimation(target_faces)
        logger.info("Decimated from %d to %d faces", len(mesh.faces), len(simplified.faces))
        return simplified
    except Exception:
        logger.warning("Quadric decimation unavailable, returning original mesh")
        return mesh


# ---------------------------------------------------------------------------
# Dimensional Scaling
# ---------------------------------------------------------------------------
UNIT_TO_MM = {
    "mm": 1.0,
    "cm": 10.0,
    "m": 1000.0,
    "in": 25.4,
    "inch": 25.4,
    "inches": 25.4,
    "ft": 304.8,
    "feet": 304.8,
}


def scale_to_real_dimensions(
    mesh: trimesh.Trimesh,
    target_size_mm: float,
    axis: str = "auto",
) -> trimesh.Trimesh:
    """
    Scale a mesh so that its extent along `axis` matches `target_size_mm`.
    
    Args:
        mesh: Input mesh
        target_size_mm: Target dimension in millimeters
        axis: 'x', 'y', 'z', or 'auto' (largest extent)
    """
    extents = mesh.extents  # [x, y, z]
    
    if axis == "auto":
        current_size = max(extents)
    else:
        idx = {"x": 0, "y": 1, "z": 2}[axis.lower()]
        current_size = extents[idx]

    if current_size <= 0:
        return mesh

    scale_factor = target_size_mm / current_size
    mesh.apply_scale(scale_factor)
    logger.info("Scaled mesh by %.4f (%.2f → %.2fmm)", scale_factor, current_size, target_size_mm)
    return mesh


def center_mesh(mesh: trimesh.Trimesh) -> trimesh.Trimesh:
    """Center mesh at origin and place on ground plane (z=0)."""
    mesh.vertices -= mesh.centroid
    # Place on ground
    min_z = mesh.vertices[:, 2].min()
    mesh.vertices[:, 2] -= min_z
    return mesh


# ---------------------------------------------------------------------------
# Modular Decomposition
# ---------------------------------------------------------------------------
def split_into_parts(mesh: trimesh.Trimesh) -> list[trimesh.Trimesh]:
    """
    Split mesh into separate connected components.
    Each component becomes an independent, manifold part.
    """
    try:
        parts = mesh.split(only_watertight=False)
        if not parts:
            return [mesh]

        # Sort parts by volume (largest first)
        parts_with_vol = []
        for p in parts:
            try:
                vol = abs(p.volume) if p.is_watertight else p.area
            except Exception:
                vol = len(p.faces)
            parts_with_vol.append((p, vol))

        parts_with_vol.sort(key=lambda x: x[1], reverse=True)
        result = [p for p, _ in parts_with_vol]

        logger.info("Split mesh into %d parts (largest: %d faces)",
                     len(result), len(result[0].faces))
        return result

    except Exception as e:
        logger.warning("Split failed: %s", e)
        return [mesh]


def label_parts(parts: list[trimesh.Trimesh]) -> dict[str, trimesh.Trimesh]:
    """
    Label parts by relative position and size.
    Returns dict like: {"base_large": mesh, "top_small": mesh, ...}
    """
    labeled = {}
    for i, part in enumerate(parts):
        centroid = part.centroid
        # Determine position description
        if centroid[2] < -0.3:
            pos = "bottom"
        elif centroid[2] > 0.3:
            pos = "top"
        else:
            pos = "middle"

        # Determine size description
        size = "large" if len(part.faces) > len(parts[0].faces) * 0.3 else "small"

        key = f"part_{i+1:02d}_{pos}_{size}"
        labeled[key] = part

    return labeled


# ---------------------------------------------------------------------------
# Quality Metrics
# ---------------------------------------------------------------------------
def mesh_quality_report(mesh: trimesh.Trimesh) -> dict:
    """Generate a quality report for a mesh."""
    report = {
        "vertices": len(mesh.vertices),
        "faces": len(mesh.faces),
        "is_watertight": mesh.is_watertight,
        "is_volume": mesh.is_volume,
        "euler_number": mesh.euler_number,
        "extents_mm": mesh.extents.tolist(),
        "bounding_box_mm": mesh.bounds.tolist(),
    }

    try:
        report["volume_mm3"] = float(mesh.volume) if mesh.is_watertight else None
    except Exception:
        report["volume_mm3"] = None

    report["surface_area_mm2"] = float(mesh.area)

    # Check for degenerate faces
    degen = ~mesh.nondegenerate_faces()
    report["degenerate_faces"] = int(degen.sum())

    return report


# ---------------------------------------------------------------------------
# CLI test
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    # Quick self-test
    box = trimesh.creation.box(extents=[2, 3, 1])
    box = make_manifold(box)
    box = scale_to_real_dimensions(box, target_size_mm=100.0)
    box = center_mesh(box)
    report = mesh_quality_report(box)
    print("Quality Report:")
    for k, v in report.items():
        print(f"  {k}: {v}")
    print(f"\nManifold: {box.is_watertight}")