Source code for lsurf.visualization.common

# The Clear BSD License
#
# Copyright (c) 2026 Tobias Heibges
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted (subject to the limitations in the disclaimer
# below) provided that the following conditions are met:
#
#      * Redistributions of source code must retain the above copyright notice,
#      this list of conditions and the following disclaimer.
#
#      * Redistributions in binary form must reproduce the above copyright
#      notice, this list of conditions and the following disclaimer in the
#      documentation and/or other materials provided with the distribution.
#
#      * Neither the name of the copyright holder nor the names of its
#      contributors may be used to endorse or promote products derived from this
#      software without specific prior written permission.
#
# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

"""
Visualization Common Utilities

Shared styling, color mapping, and helper functions for all visualization modules.
"""

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize

# =============================================================================
# Default Style Configuration
# =============================================================================

DEFAULT_FIGSIZE = (10, 6)
DEFAULT_DPI = 150

# Color schemes
RAY_CMAP = "viridis"
INTENSITY_CMAP = "plasma"
WAVELENGTH_CMAP = "rainbow"
TIME_CMAP = "coolwarm"
GENERATION_CMAP = "tab10"

# Plot styling
GRID_ALPHA = 0.3
LINE_ALPHA = 0.6
SCATTER_ALPHA = 0.5
DEFAULT_LINEWIDTH = 0.5
DEFAULT_MARKERSIZE = 5


# =============================================================================
# Color Mapping Utilities
# =============================================================================




def get_color_mapping(
    values: np.ndarray,
    cmap_name: str = "viridis",
    vmin: float | None = None,
    vmax: float | None = None,
) -> tuple[np.ndarray, Normalize, ScalarMappable]:
    """
    Create color mapping for values.

    Parameters
    ----------
    values : ndarray
        Values to map to colors.
    cmap_name : str
        Matplotlib colormap name.
    vmin, vmax : float, optional
        Value range. If None, uses data min/max.

    Returns
    -------
    colors : ndarray
        RGBA colors for each value.
    norm : Normalize
        Normalization object.
    sm : ScalarMappable
        ScalarMappable for colorbar.
    """
    cmap = plt.cm.get_cmap(cmap_name)

    if vmin is None:
        vmin = np.min(values)
    if vmax is None:
        vmax = np.max(values)

    norm = Normalize(vmin=vmin, vmax=vmax)
    colors = cmap(norm(values))
    sm = ScalarMappable(norm=norm, cmap=cmap)
    sm.set_array([])

    return colors, norm, sm





def wavelength_to_color(
    wavelength_m: float | np.ndarray,
) -> tuple[float, float, float] | np.ndarray:
    """
    Convert wavelength (in meters) to approximate visible color.

    Parameters
    ----------
    wavelength_m : float or ndarray
        Wavelength in meters.

    Returns
    -------
    color : tuple or ndarray
        RGB color tuple(s) in [0, 1] range.

    Notes
    -----
    Uses simplified visible spectrum approximation.
    Wavelengths outside visible range return gray.
    """
    wavelength_nm = np.asarray(wavelength_m) * 1e9

    # Simple approximation of visible spectrum
    r = np.zeros_like(wavelength_nm)
    g = np.zeros_like(wavelength_nm)
    b = np.zeros_like(wavelength_nm)

    # Violet to blue (380-440nm)
    mask = (wavelength_nm >= 380) & (wavelength_nm < 440)
    r[mask] = -(wavelength_nm[mask] - 440) / (440 - 380)
    b[mask] = 1.0

    # Blue to cyan (440-490nm)
    mask = (wavelength_nm >= 440) & (wavelength_nm < 490)
    g[mask] = (wavelength_nm[mask] - 440) / (490 - 440)
    b[mask] = 1.0

    # Cyan to green (490-510nm)
    mask = (wavelength_nm >= 490) & (wavelength_nm < 510)
    g[mask] = 1.0
    b[mask] = -(wavelength_nm[mask] - 510) / (510 - 490)

    # Green to yellow (510-580nm)
    mask = (wavelength_nm >= 510) & (wavelength_nm < 580)
    r[mask] = (wavelength_nm[mask] - 510) / (580 - 510)
    g[mask] = 1.0

    # Yellow to red (580-645nm)
    mask = (wavelength_nm >= 580) & (wavelength_nm < 645)
    r[mask] = 1.0
    g[mask] = -(wavelength_nm[mask] - 645) / (645 - 580)

    # Red (645-750nm)
    mask = (wavelength_nm >= 645) & (wavelength_nm <= 750)
    r[mask] = 1.0

    # Gray for out-of-range
    mask = (wavelength_nm < 380) | (wavelength_nm > 750)
    r[mask] = g[mask] = b[mask] = 0.5

    if np.isscalar(wavelength_m):
        return (float(r), float(g), float(b))

    return np.stack([r, g, b], axis=-1)



# =============================================================================
# Axis Setup Utilities
# =============================================================================




def setup_axis_grid(
    ax: plt.Axes,
    xlabel: str = "",
    ylabel: str = "",
    title: str = "",
    grid: bool = True,
) -> None:
    """
    Apply standard axis formatting.

    Parameters
    ----------
    ax : Axes
        Matplotlib axes.
    xlabel, ylabel : str
        Axis labels.
    title : str
        Axis title.
    grid : bool
        Whether to show grid.
    """
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    if title:
        ax.set_title(title)
    if grid:
        ax.grid(True, alpha=GRID_ALPHA)





def add_colorbar(
    ax: plt.Axes,
    mappable: ScalarMappable,
    label: str = "",
    orientation: str = "vertical",
) -> plt.colorbar:
    """
    Add colorbar to axes.

    Parameters
    ----------
    ax : Axes
        Matplotlib axes.
    mappable : ScalarMappable
        Object with color mapping.
    label : str
        Colorbar label.
    orientation : str
        'vertical' or 'horizontal'.

    Returns
    -------
    colorbar
        Matplotlib colorbar object.
    """
    cbar = plt.colorbar(mappable, ax=ax, orientation=orientation, pad=0.02)
    if label:
        cbar.set_label(label)
    return cbar



# =============================================================================
# Figure Creation Utilities
# =============================================================================




def create_figure(
    nrows: int = 1,
    ncols: int = 1,
    figsize: tuple[float, float] | None = None,
    constrained_layout: bool = True,
    **kwargs,
) -> tuple[plt.Figure, plt.Axes | np.ndarray]:
    """
    Create figure with standard settings.

    Parameters
    ----------
    nrows, ncols : int
        Number of subplot rows and columns.
    figsize : tuple, optional
        Figure size. Defaults to scaled DEFAULT_FIGSIZE.
    constrained_layout : bool
        Use constrained layout.
    **kwargs
        Additional arguments to plt.subplots.

    Returns
    -------
    fig : Figure
        Matplotlib figure.
    axes : Axes or ndarray
        Axes object(s).
    """
    if figsize is None:
        width = DEFAULT_FIGSIZE[0] * ncols / 2
        height = DEFAULT_FIGSIZE[1] * nrows / 2
        figsize = (width, height)

    return plt.subplots(
        nrows, ncols, figsize=figsize, constrained_layout=constrained_layout, **kwargs
    )





def save_figure(
    fig: plt.Figure,
    path: str,
    dpi: int = DEFAULT_DPI,
    bbox_inches: str = "tight",
    **kwargs,
) -> None:
    """
    Save figure with standard settings.

    Parameters
    ----------
    fig : Figure
        Matplotlib figure.
    path : str
        Output path.
    dpi : int
        Resolution.
    bbox_inches : str
        Bounding box setting.
    **kwargs
        Additional arguments to savefig.
    """
    fig.savefig(path, dpi=dpi, bbox_inches=bbox_inches, **kwargs)



# =============================================================================
# Coordinate Projection Utilities
# =============================================================================




def get_projection_config(projection: str) -> tuple[int, int, str, str]:
    """Get axis indices and labels for xy/xz/yz projections.

    Parameters
    ----------
    projection : str
        Projection plane: 'xy', 'xz', or 'yz'.

    Returns
    -------
    idx1 : int
        First coordinate index (0=x, 1=y, 2=z).
    idx2 : int
        Second coordinate index.
    xlabel : str
        Label for first axis.
    ylabel : str
        Label for second axis.
    """
    proj_map = {"xy": (0, 1), "xz": (0, 2), "yz": (1, 2)}
    label_map = {
        "xy": ("X (m)", "Y (m)"),
        "xz": ("X (m)", "Z (m)"),
        "yz": ("Y (m)", "Z (m)"),
    }
    projection = projection.lower()
    idx1, idx2 = proj_map.get(projection, (0, 2))
    xlabel, ylabel = label_map.get(projection, ("X (m)", "Z (m)"))
    return idx1, idx2, xlabel, ylabel