fluencyCAD/src/fluency/rendering/occ_to_mesh.py

"""Convert OCC BRep shapes to mesh files for render backends.

Outputs PLY files (preferred by Mitsuba) or STL files.
"""

from __future__ import annotations

import logging
import os
import tempfile
from typing import List, Optional, Tuple

import numpy as np

logger = logging.getLogger(__name__)


def occ_shape_to_ply(
    shape,
    output_path: Optional[str] = None,
    linear_deflection: float = 0.1,
    angular_deflection: float = 0.15,
) -> str:
    """Tessellate an OCC TopoDS_Shape and write as PLY.

    Returns the path to the written PLY file.
    """
    from OCP.BRepMesh import BRepMesh_IncrementalMesh
    from OCP.TopExp import TopExp_Explorer
    from OCP.TopAbs import TopAbs_FACE
    from OCP.TopoDS import TopoDS
    from OCP.BRep import BRep_Tool
    from OCP.TopLoc import TopLoc_Location

    # Tessellate
    tess = BRepMesh_IncrementalMesh(
        shape, linear_deflection, False, angular_deflection, True
    )
    tess.Perform()

    # Extract triangulation from all faces
    all_vertices: List[List[float]] = []
    all_faces: List[List[int]] = []
    vertex_offset = 0

    from OCP.TopAbs import TopAbs_FORWARD

    explorer = TopExp_Explorer(shape, TopAbs_FACE)
    while explorer.More():
        face = TopoDS.Face_s(explorer.Current())
        location = TopLoc_Location()
        triangulation = BRep_Tool.Triangulation_s(face, location)

        if triangulation is None:
            explorer.Next()
            continue

        # Transform
        trsf = location.Transformation()

        # Check face orientation: FORWARD means the surface normal points
        # outward from the solid; REVERSED means it points inward.
        is_forward = (face.Orientation() == TopAbs_FORWARD)

        # Extract vertices (apply location transform to positions)
        nb_nodes = triangulation.NbNodes()
        for i in range(1, nb_nodes + 1):
            node = triangulation.Node(i)
            pnt = node.Transformed(trsf)
            all_vertices.append([pnt.X(), pnt.Y(), pnt.Z()])

        # Extract triangles
        # For REVERSED faces, swap winding order (n1, n3, n2) so that
        # the computed normal points outward consistently.
        nb_triangles = triangulation.NbTriangles()
        for i in range(1, nb_triangles + 1):
            tri = triangulation.Triangle(i)
            n1, n2, n3 = tri.Get()
            if is_forward:
                all_faces.append([
                    n1 - 1 + vertex_offset,
                    n2 - 1 + vertex_offset,
                    n3 - 1 + vertex_offset,
                ])
            else:
                # Swap winding for REVERSED faces
                all_faces.append([
                    n1 - 1 + vertex_offset,
                    n3 - 1 + vertex_offset,
                    n2 - 1 + vertex_offset,
                ])

        vertex_offset += nb_nodes
        explorer.Next()

    if not all_vertices:
        raise ValueError("Tessellation produced no vertices")

    vertices = np.array(all_vertices, dtype=np.float32)
    faces = np.array(all_faces, dtype=np.uint32)

    logger.info(
        f"Tessellation: {len(vertices)} vertices, {len(faces)} triangles"
    )

    # Compute smooth vertex normals from face normals.
    # Winding is already corrected during tessellation using OCC face orientation.
    normals, corrected_faces = _compute_outward_normals(vertices, faces, shape)

    # Write PLY with corrected faces and normals
    if output_path is None:
        fd, output_path = tempfile.mkstemp(suffix=".ply", prefix="fluency_render_")
        os.close(fd)

    _write_ply(output_path, vertices, corrected_faces, normals)
    logger.info(f"Wrote PLY: {output_path}")
    return output_path


def _compute_outward_normals(
    vertices: np.ndarray,
    faces: np.ndarray,
    shape,
) -> Tuple[np.ndarray, np.ndarray]:
    """Compute outward-facing vertex normals and correct face winding.

    The winding is already corrected during tessellation using OCC's face
    orientation (TopAbs_FORWARD/REVERSED). This function computes smooth
    vertex normals by averaging face normals at shared vertices.

    Returns (normals, corrected_faces) for PLY export.
    """
    n_verts = len(vertices)
    v_normals = np.zeros((n_verts, 3), dtype=np.float64)

    # Ensure faces is 2D (numpy creates (3,) for single-face meshes)
    if faces.ndim == 1:
        faces = faces.reshape(1, -1)

    # Winding is already correct from tessellation (face orientation check).
    # Just compute face normals and accumulate to vertices.
    v0 = vertices[faces[:, 0]]
    v1 = vertices[faces[:, 1]]
    v2 = vertices[faces[:, 2]]

    edge1 = v1 - v0
    edge2 = v2 - v0
    fn = np.cross(edge1, edge2)

    # Normalize face normals
    lengths = np.linalg.norm(fn, axis=1, keepdims=True)
    lengths[lengths < 1e-10] = 1.0
    fn /= lengths

    # Accumulate to vertices
    for i in range(len(faces)):
        idx = faces[i]
        v_normals[idx[0]] += fn[i]
        v_normals[idx[1]] += fn[i]
        v_normals[idx[2]] += fn[i]

    # Normalize vertex normals
    v_lengths = np.linalg.norm(v_normals, axis=1, keepdims=True)
    v_lengths[v_lengths < 1e-10] = 1.0
    v_normals /= v_lengths

    return v_normals.astype(np.float32), faces.astype(np.uint32)


def occ_shape_to_stl(
    shape,
    output_path: Optional[str] = None,
    linear_deflection: float = 0.1,
) -> str:
    """Tessellate an OCC TopoDS_Shape and write as binary STL.

    Returns the path to the written STL file.
    """
    from OCP.BRepMesh import BRepMesh_IncrementalMesh
    from OCP.StlAPI import StlAPI_Writer

    # Tessellate
    tess = BRepMesh_IncrementalMesh(shape, linear_deflection, False, 0.5, True)
    tess.Perform()

    if output_path is None:
        fd, output_path = tempfile.mkstemp(suffix=".stl", prefix="fluency_render_")
        os.close(fd)

    writer = StlAPI_Writer()
    writer.SetASCIIMode(False)
    writer.Write(shape, output_path)
    logger.info(f"Wrote STL: {output_path}")
    return output_path


def occ_shape_bounds(shape) -> Tuple[Tuple[float, float, float], Tuple[float, float, float]]:
    """Return (min_xyz, max_xyz) bounding box of an OCC shape."""
    from OCP.Bnd import Bnd_Box
    from OCP.BRepBndLib import BRepBndLib

    bbox = Bnd_Box()
    BRepBndLib.Add_s(shape, bbox)
    xmin, ymin, zmin, xmax, ymax, zmax = bbox.Get()
    return (xmin, ymin, zmin), (xmax, ymax, zmax)


def _write_ply(
    path: str,
    vertices: np.ndarray,
    faces: np.ndarray,
    normals: Optional[np.ndarray] = None,
) -> None:
    """Write a binary PLY file (little-endian) with optional vertex normals."""
    import struct

    n_verts = len(vertices)
    n_faces = len(faces)
    has_normals = normals is not None and len(normals) == n_verts

    with open(path, "wb") as f:
        # Header
        header_lines = [
            "ply",
            "format binary_little_endian 1.0",
            f"element vertex {n_verts}",
            "property float x",
            "property float y",
            "property float z",
        ]
        if has_normals:
            header_lines.extend([
                "property float nx",
                "property float ny",
                "property float nz",
            ])
        header_lines.append(f"element face {n_faces}")
        header_lines.append("property list uchar int vertex_indices")
        header_lines.append("end_header")

        f.write(("\n".join(header_lines) + "\n").encode("ascii"))

        # Vertex positions (+ normals if available)
        for i in range(n_verts):
            f.write(struct.pack("<fff", vertices[i, 0], vertices[i, 1], vertices[i, 2]))
            if has_normals:
                f.write(struct.pack("<fff", normals[i, 0], normals[i, 1], normals[i, 2]))

        # Faces
        for face in faces:
            f.write(struct.pack("<B", 3))
            f.write(struct.pack("<iii", int(face[0]), int(face[1]), int(face[2])))