"""OpenCV visualization helpers for example scripts."""

from __future__ import annotations

import cv2
import numpy as np


def depth_colormap(depth_map, colormap=cv2.COLORMAP_JET) -> np.ndarray:
    """Normalize a scalar map to uint8 and apply an OpenCV colormap."""
    values = np.asarray(depth_map)
    if values.size == 0:
        return np.zeros((1, 1, 3), dtype=np.uint8)
    if values.max() > values.min():
        normalized = ((values - values.min()) / (values.max() - values.min()) * 255).astype(np.uint8)
    else:
        normalized = np.zeros_like(values, dtype=np.uint8)
    return cv2.applyColorMap(normalized, colormap)


def draw_magnitude_map(
    vfield,
    threshold=3.0,
    colormap=cv2.COLORMAP_JET,
    output_size=(400, 400),
) -> np.ndarray:
    """Render a smoothed optical-flow magnitude map as a BGR pseudo-color image."""
    magnitude = np.sqrt(vfield[:, :, 0] ** 2 + vfield[:, :, 1] ** 2)

    downsample_factor = 10
    h, w = magnitude.shape
    magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor]
    magnitude_smoothed = cv2.resize(magnitude_downsampled, (w, h), interpolation=cv2.INTER_CUBIC)

    magnitude_thresholded = magnitude_smoothed.copy()
    magnitude_thresholded[magnitude_smoothed < threshold] = 0

    mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max()
    if mag_max > mag_min:
        magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8)
    else:
        magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8)

    magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap)
    out_w, out_h = output_size
    return cv2.resize(magnitude_colored, (out_w, out_h), interpolation=cv2.INTER_CUBIC)


def frame_to_bgr(frame: np.ndarray) -> np.ndarray:
    """Convert a grayscale or single-channel frame to BGR for display/recording."""
    if frame is None:
        raise ValueError("frame is None")
    if len(frame.shape) == 2:
        return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
    if len(frame.shape) == 3 and frame.shape[2] == 1:
        return cv2.cvtColor(frame.squeeze(2), cv2.COLOR_GRAY2BGR)
    return frame