"""Display utilities for heatmaps and colormaps."""
import numpy as np
import cv2

from config.constants import COLORMAPS, COLORMAP_N_SAMPLES


def cv_colormap_to_plotly_colorscale(colormap_name, n_samples=None):
    """Build a Plotly colorscale from OpenCV colormap so UI matches download/PDF exactly."""
    n = n_samples or COLORMAP_N_SAMPLES
    cv2_cmap = COLORMAPS.get(colormap_name, cv2.COLORMAP_JET)
    gradient = np.linspace(0, 255, n, dtype=np.uint8).reshape(1, -1)
    rgb = cv2.applyColorMap(gradient, cv2_cmap)
    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
    scale = []
    for i in range(n):
        r, g, b = rgb[0, i]
        scale.append([i / (n - 1), f"rgb({r},{g},{b})"])
    return scale


def is_display_range_remapped(display_mode, clip_min, clip_max):
    """
    True when the UI uses Range mode with a clip window other than the full [0, 1]
    (colorbar tick labels map normalized 0–1 to that interval).
    """
    if display_mode != "Range":
        return False
    lo, hi = float(clip_min), float(clip_max)
    return hi > lo and not (lo == 0.0 and hi == 1.0)


def apply_display_scale(heatmap, mode, clip_min=0, clip_max=1, clamp_only=False):
    """
    Map the normalised prediction (0-1) to colormap positions. Display only; values are unchanged.

    ``Default``: the full scale, 0-1 (= 0-500 Pa single cells, 0-50 Pa spheroids).
    ``Range``: the window [clip_min, clip_max] is stretched over the whole colormap; values below the window
    are shown as the lowest colour and values above it **saturate at the highest colour**. (Earlier versions set
    values above the window to 0, which displayed the strongest stresses as zero.)
    ``clamp_only`` is kept for backward compatibility and ignored.
    """
    h = np.asarray(heatmap, dtype=np.float32)
    if mode != "Range":
        return np.clip(h, 0, 1).astype(np.float32)
    vmin, vmax = float(clip_min), float(clip_max)
    if vmax > vmin:
        return np.clip((h - vmin) / (vmax - vmin), 0, 1).astype(np.float32)
    return np.clip(h, 0, 1).astype(np.float32)