from __future__ import annotations

import numpy as np
import cv2

# Colour palette for up to 10 classes (BGR → RGB)
DEFAULT_PALETTE = [
    [0, 0, 0],        # 0 — background (black)
    [255, 50, 50],     # 1 — building (red)
    [50, 255, 50],     # 2
    [50, 50, 255],     # 3
    [255, 255, 50],    # 4
    [255, 50, 255],    # 5
    [50, 255, 255],    # 6
    [200, 200, 200],   # 7
    [128, 0, 255],     # 8
    [0, 128, 255],     # 9
]

def overlay_mask(
    image: np.ndarray,
    mask: np.ndarray,
    alpha: float = 0.45,
    colour: tuple[int, int, int] = (255, 50, 50),
) -> np.ndarray:
    """Overlay a binary mask on an RGB image.

    Args:
        image:  HWC uint8 RGB image.
        mask:   HW boolean or int array (True / >0 = foreground).
        alpha:  Opacity of the overlay.
        colour: RGB colour for foreground pixels.

    Returns:
        Blended HWC uint8 RGB image.
    """
    overlay = image.copy()
    bool_mask = mask.astype(bool)
    overlay[bool_mask] = np.array(colour, dtype=np.uint8)
    output = image.copy()
    cv2.addWeighted(overlay, alpha, output, 1.0 - alpha, 0, output)
    return output

def colourize_labels(
    label: np.ndarray,
    palette: list[list[int]] | None = None,
) -> np.ndarray:
    """Convert an integer label map to an RGB colour image.

    Args:
        label:   HW int array of class IDs.
        palette: list of [R, G, B] per class index.

    Returns:
        HWC uint8 RGB image.
    """
    if palette is None:
        palette = DEFAULT_PALETTE

    h, w = label.shape
    rgb = np.zeros((h, w, 3), dtype=np.uint8)
    for cls_id, colour in enumerate(palette):
        rgb[label == cls_id] = colour
    return rgb

def plot_predictions(
    image: np.ndarray,
    gt_label: np.ndarray,
    pred_label: np.ndarray,
    save_path: str | None = None,
) -> np.ndarray:
    """Create a side-by-side comparison: image | GT overlay | prediction overlay.

    Returns:
        Concatenated HWC uint8 RGB image (width = 3 × original).
    """
    gt_overlay = overlay_mask(image, gt_label > 0, alpha=0.4, colour=(50, 255, 50))
    pred_overlay = overlay_mask(image, pred_label > 0, alpha=0.4, colour=(255, 50, 50))

    combined = np.concatenate([image, gt_overlay, pred_overlay], axis=1)

    if save_path is not None:
        bgr = cv2.cvtColor(combined, cv2.COLOR_RGB2BGR)
        cv2.imwrite(save_path, bgr)

    return combined