visualizer/wired3d_viewer/stl_viewer.py

"""Interactive 3D viewer for raw STL meshes (pre-slicing).

This is the mesh-side companion to ``viewer.py`` (which renders sliced NC
toolpaths). It loads an STL with ``trimesh`` and builds a single self-contained
Plotly HTML document showing the solid surface, so you can rotate/zoom the part
in 3D *before* slicing.

A per-face **direction-vector** overlay (the surface normals, via
:func:`build_normal_trace`) is kept in this module as the seed of the
projectile-direction feature, but is **not drawn yet** — once slicing is wired
in, the direction shown will be the tool-head orientation at every toolpath
point rather than raw facet normals.

The styling deliberately mirrors ``viewer.py``: dark theme, perspective camera,
axis ranges pinned to the part bounding box, the wired3d logo embedded top-left,
and a stats table. It reuses :func:`viewer.logo_data_uri` and
:data:`viewer.DEFAULT_CAMERA` so the two viewers stay visually consistent.

Run (from ``visualizer/``):
  python -m wired3d_viewer view-stl <file.stl>   # writes <stem>_stl_viewer.html

Dependencies: trimesh, plotly, numpy.
"""

from pathlib import Path

import numpy as np
import plotly.graph_objects as go
import plotly.io as pio
import trimesh

from .viewer import DEFAULT_CAMERA, FIG_BG, SCENE_BG, logo_data_uri

# Surface + direction-vector colours, in keeping with the NC viewer palette.
MESH_COLOUR = "#4a90d9"       # cool steel blue for the solid surface
NORMAL_COLOUR = "#ffcc00"     # amber direction vectors — pop against the blue
EDGE_COLOUR = "#1b2a44"

# Cap on how many normal vectors we draw. A dense mesh can have 10k+ faces;
# drawing one cone per face is both unreadable and slow, so we evenly sample
# down to this many. The full mesh is always drawn — only the vector overlay is
# decimated (and the count of dropped vectors is reported in the stats table).
MAX_NORMALS = 1500


def load_stl(filepath: str) -> trimesh.Trimesh:
    """Load an STL file and return a single concrete ``Trimesh``.

    Inputs:
        filepath -- path to a binary or ASCII ``.stl`` file.
    Output:
        a ``trimesh.Trimesh``. If the file is a scene (multiple bodies) it is
        concatenated into one mesh so the rest of the pipeline sees one object.
    Raises:
        ValueError if the file contains no triangulated geometry.
    """
    loaded = trimesh.load(filepath, force="mesh")
    if isinstance(loaded, trimesh.Scene):
        loaded = trimesh.util.concatenate(loaded.dump())
    if not isinstance(loaded, trimesh.Trimesh) or len(loaded.faces) == 0:
        raise ValueError(f"No triangle mesh found in {filepath!r}")
    return loaded


def _axis_ranges(mesh: trimesh.Trimesh) -> dict:
    """Return padded, equal-aspect X/Y/Z scene ranges for the mesh bounds.

    Inputs:
        mesh -- the loaded ``Trimesh``.
    Output:
        dict with ``xaxis``/``yaxis``/``zaxis`` range entries. A common padded
        cube edge is used for all three axes so the part is never distorted and
        the view does not zoom-jump (same approach as the NC viewer).
    """
    lo, hi = mesh.bounds
    centre = (lo + hi) / 2.0
    half = float((hi - lo).max()) / 2.0
    half = half * 1.1 if half > 0 else 1.0  # 10% padding; guard a flat part
    return dict(
        xaxis=dict(range=[centre[0] - half, centre[0] + half]),
        yaxis=dict(range=[centre[1] - half, centre[1] + half]),
        zaxis=dict(range=[centre[2] - half, centre[2] + half]),
    )


def build_mesh_trace(mesh: trimesh.Trimesh) -> go.Mesh3d:
    """Build the solid-surface ``Mesh3d`` trace.

    Inputs:
        mesh -- the loaded ``Trimesh``.
    Output:
        a lit ``go.Mesh3d`` using the mesh vertices/faces.
    """
    v = mesh.vertices
    f = mesh.faces
    return go.Mesh3d(
        x=v[:, 0], y=v[:, 1], z=v[:, 2],
        i=f[:, 0], j=f[:, 1], k=f[:, 2],
        color=MESH_COLOUR,
        opacity=1.0,
        flatshading=True,
        name="STL surface",
        showlegend=True,
        lighting=dict(ambient=0.45, diffuse=0.85, specular=0.25, roughness=0.55),
        lightposition=dict(x=1000, y=1000, z=2000),
        hoverinfo="skip",
    )


def _sample_faces(n_faces: int, cap: int) -> np.ndarray:
    """Return indices of faces to draw a normal for (evenly sampled, capped).

    Inputs:
        n_faces -- total number of faces.
        cap     -- maximum number of vectors to draw.
    Output:
        a 1-D int array of face indices (all of them when ``n_faces <= cap``,
        otherwise an even stride across the range).
    """
    if n_faces <= cap:
        return np.arange(n_faces)
    return np.linspace(0, n_faces - 1, cap).astype(int)


def build_normal_trace(mesh: trimesh.Trimesh) -> tuple[go.Cone, int]:
    """Build the per-face direction-vector overlay (surface normals as cones).

    Inputs:
        mesh -- the loaded ``Trimesh``.
    Output:
        ``(cone_trace, shown)`` where ``cone_trace`` is a ``go.Cone`` quiver
        placed at each sampled face centroid pointing along its outward normal,
        and ``shown`` is how many vectors were actually drawn (after sampling).
    The cones are length-normalised relative to the part size so they read as
    little arrows regardless of model scale.
    """
    idx = _sample_faces(len(mesh.faces), MAX_NORMALS)
    origins = mesh.triangles_center[idx]
    normals = mesh.face_normals[idx]

    # Scale the cone vectors to ~3% of the bounding-box diagonal so they are
    # visible but don't swamp the surface. go.Cone sizes cones by vector norm.
    diag = float(np.linalg.norm(mesh.bounds[1] - mesh.bounds[0])) or 1.0
    vec = normals * (diag * 0.03)

    cone = go.Cone(
        x=origins[:, 0], y=origins[:, 1], z=origins[:, 2],
        u=vec[:, 0], v=vec[:, 1], w=vec[:, 2],
        colorscale=[[0, NORMAL_COLOUR], [1, NORMAL_COLOUR]],
        showscale=False,
        sizemode="absolute",
        sizeref=diag * 0.03,
        anchor="tail",
        name="Direction (face normals)",
        showlegend=True,
        hoverinfo="x+y+z",
    )
    return cone, len(idx)


def compute_stats(mesh: trimesh.Trimesh) -> dict:
    """Compute a small dictionary of human-readable mesh statistics.

    Inputs:
        mesh -- the loaded ``Trimesh``.
    Output:
        ordered dict of ``label -> value`` strings for the stats table.
    """
    lo, hi = mesh.bounds
    size = hi - lo
    stats = {
        "Vertices": f"{len(mesh.vertices):,}",
        "Faces": f"{len(mesh.faces):,}",
        "Size X×Y×Z (mm)": f"{size[0]:.1f} × {size[1]:.1f} × {size[2]:.1f}",
        "Bounds min (mm)": f"({lo[0]:.1f}, {lo[1]:.1f}, {lo[2]:.1f})",
        "Bounds max (mm)": f"({hi[0]:.1f}, {hi[1]:.1f}, {hi[2]:.1f})",
        "Watertight": "yes" if mesh.is_watertight else "no",
        "Volume (mm³)": f"{mesh.volume:,.1f}" if mesh.is_watertight else "n/a",
        "Surface area (mm²)": f"{mesh.area:,.1f}",
    }
    return stats


def build_stats_table(stats: dict) -> go.Table:
    """Build the bottom-right stats ``go.Table`` trace.

    Inputs:
        stats -- the dict from :func:`compute_stats`.
    Output:
        a styled ``go.Table`` (dark theme, two columns).
    """
    return go.Table(
        columnwidth=[1.1, 1.4],
        header=dict(
            values=["<b>Property</b>", "<b>Value</b>"],
            fill_color="#16213e", font=dict(color="#e8e8f0", size=13),
            align="left", height=26,
        ),
        cells=dict(
            values=[list(stats.keys()), list(stats.values())],
            fill_color="#0f1626", font=dict(color="#cfd3e6", size=12),
            align="left", height=24,
        ),
    )


def build_figure(mesh: trimesh.Trimesh, part_name: str) -> go.Figure:
    """Assemble the full STL viewer figure (3D mesh + stats table).

    Inputs:
        mesh      -- the loaded ``Trimesh``.
        part_name -- title shown in the header.
    Output:
        a ``go.Figure`` with a large left 3D scene and a bottom-right stats
        table, dark themed, perspective camera, axis ranges pinned to bounds.
    """
    from plotly.subplots import make_subplots

    fig = make_subplots(
        rows=2, cols=2,
        column_widths=[0.74, 0.26], row_heights=[0.55, 0.45],
        specs=[[{"type": "scene", "rowspan": 2}, {"type": "table"}],
               [None, {"type": "table"}]],
        horizontal_spacing=0.02, vertical_spacing=0.04,
    )

    # Direction vectors (build_normal_trace) are intentionally not drawn yet —
    # see the module docstring. Only the solid surface is shown for now.
    fig.add_trace(build_mesh_trace(mesh), row=1, col=1)

    table = build_stats_table(compute_stats(mesh))
    fig.add_trace(table, row=1, col=2)

    ranges = _axis_ranges(mesh)
    fig.update_layout(
        title=dict(text=f"STL Viewer — {part_name}", x=0.5,
                   font=dict(color="#e8e8f0", size=18)),
        paper_bgcolor=FIG_BG,
        font=dict(color="#e8e8f0"),
        showlegend=True,
        legend=dict(x=0.0, y=1.0, bgcolor="rgba(22,33,62,0.6)",
                    font=dict(size=12), traceorder="normal"),
        autosize=True,
        margin=dict(l=0, r=0, t=44, b=0),
        scene=dict(
            bgcolor=SCENE_BG,
            xaxis=dict(title="X", color="#9aa0b5", **ranges["xaxis"]),
            yaxis=dict(title="Y", color="#9aa0b5", **ranges["yaxis"]),
            zaxis=dict(title="Z", color="#9aa0b5", **ranges["zaxis"]),
            aspectmode="cube",
            camera=DEFAULT_CAMERA,
            uirevision="keep",
        ),
    )

    logo = logo_data_uri()
    if logo:
        fig.add_layout_image(dict(
            source=logo, xref="paper", yref="paper",
            x=0.01, y=0.99, sizex=0.12, sizey=0.12,
            xanchor="left", yanchor="top", layer="above",
        ))
    return fig


def write_html(mesh: trimesh.Trimesh, part_name: str, out_path: str) -> str:
    """Build the figure and write a self-contained HTML file.

    Inputs:
        mesh      -- the loaded ``Trimesh``.
        part_name -- title / part name.
        out_path  -- destination ``.html`` path.
    Output:
        ``out_path`` (also written to disk). plotly.js is pulled from the CDN
        so the file is small but needs a network connection on first open.
    """
    fig = build_figure(mesh, part_name)
    pio.write_html(
        fig, out_path, include_plotlyjs="cdn", full_html=True,
        default_width="100%", default_height="100vh",
        config={"responsive": True},
    )
    return out_path


def main(stl_file: str) -> None:
    """CLI entry: render one STL file to ``<stem>_stl_viewer.html``.

    Inputs:
        stl_file -- path to the ``.stl`` file.
    Output:
        None. Prints the saved HTML path.
    """
    path = Path(stl_file)
    mesh = load_stl(str(path))
    out = path.with_name(f"{path.stem}_stl_viewer.html")
    write_html(mesh, path.stem, str(out))
    print(f"Saved: {out}")


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        print("usage: python -m wired3d_viewer.stl_viewer <file.stl>")
        sys.exit(1)
    main(sys.argv[1])