"""Color-scaling helpers for B-scan visualization."""

from __future__ import annotations

import numpy as np
import pyqtgraph as pg

def bscan_lookup_table(axis_mode: str) -> np.ndarray:
    """Build B-scan colormap table for selected axis mode."""
    if axis_mode == "abs":
        return build_lut(["#440154", "#31688e", "#35b779", "#fde725"])
    return build_lut(["#2166ac", "#67a9cf", "#f7f7f7", "#ef8a62", "#b2182b"])


def build_lut(stops: list[str], *, size: int = 256) -> np.ndarray:
    """Interpolate hex color stops into 8-bit RGB LUT array."""
    stop_positions = np.linspace(0.0, 1.0, num=len(stops), dtype=np.float32)
    sample_positions = np.linspace(0.0, 1.0, num=size, dtype=np.float32)
    stop_colors = np.asarray([pg.mkColor(value).getRgb()[:3] for value in stops], dtype=np.float32)

    lut = np.empty((size, 3), dtype=np.uint8)
    for channel in range(3):
        lut[:, channel] = np.interp(sample_positions, stop_positions, stop_colors[:, channel]).astype(np.uint8)
    return lut


def bscan_levels(sweeps: np.ndarray, axis_mode: str) -> tuple[float, float]:
    """Compute image levels for B-scan data based on axis mode."""
    min_value = float(np.min(sweeps))
    max_value = float(np.max(sweeps))
    if axis_mode == "abs":
        if max_value <= min_value:
            return min_value, min_value + 1e-6
        return min_value, max_value

    max_abs = max(abs(min_value), abs(max_value), 1e-6)
    return -max_abs, max_abs