from __future__ import annotations

from typing import Iterable

import numpy as np
from PySide6 import QtGui


COLOR_PALETTE = [
    (128, 0, 0),      # maroon
    (0, 0, 139),      # darkblue
    (0, 100, 0),      # darkgreen
    (0, 255, 255),    # cyan
    (0, 139, 139),    # darkcyan
    (255, 0, 255),    # magenta
    (75, 0, 130),     # indigo
    (128, 128, 128),  # grey
    (205, 133, 63),   # peru
    (106, 90, 205),   # slateblue
    (0, 250, 154),    # mediumspringgreen
    (255, 69, 0),     # orangered
]


def ensure_uint8(img: np.ndarray) -> np.ndarray:
    if img.dtype == np.uint8:
        return img
    img = np.clip(img, 0, 255)
    return img.astype("uint8")


def numpy_to_qimage(img: np.ndarray) -> QtGui.QImage:
    if img.ndim == 2:
        gray = ensure_uint8(img)
        h, w = gray.shape
        return QtGui.QImage(gray.data, w, h, w, QtGui.QImage.Format_Grayscale8).copy()
    if img.ndim == 3 and img.shape[2] == 3:
        rgb = ensure_uint8(img)
        h, w, _ = rgb.shape
        bytes_per_line = 3 * w
        return QtGui.QImage(rgb.data, w, h, bytes_per_line, QtGui.QImage.Format_RGB888).copy()
    raise ValueError("Unsupported image shape for display.")


def colorize_labels(labels: np.ndarray, palette: Iterable[tuple[int, int, int]] | None = None) -> np.ndarray:
    if palette is None:
        palette = COLOR_PALETTE
    palette = list(palette)
    h, w = labels.shape
    img = np.zeros((h, w, 3), dtype="uint8")
    unique_labels = np.unique(labels[np.isfinite(labels)])
    for idx, label in enumerate(unique_labels):
        if label < 0:
            continue
        color = palette[int(idx) % len(palette)]
        img[labels == label] = color
    return img


def colorize_probability(prob_map: np.ndarray) -> np.ndarray:
    prob = np.nan_to_num(prob_map, nan=0.0)
    prob = np.clip(prob, 0, 1)
    r = (prob * 255).astype("uint8")
    g = ((1 - prob) * 140).astype("uint8")
    b = (255 - r).astype("uint8")
    return np.stack([r, g, b], axis=-1)


def colorize_index(index_map: np.ndarray) -> np.ndarray:
    idx = np.nan_to_num(index_map, nan=0.0)
    idx = np.clip(idx, -1, 1)
    norm = (idx + 1) / 2  # 0..1
    r = (norm * 255).astype("uint8")
    b = ((1 - norm) * 255).astype("uint8")
    g = (0.5 * 255 - np.abs(norm - 0.5) * 255).astype("uint8")
    return np.stack([r, g, b], axis=-1)
