"""Shared flatmap plotting helpers."""

from __future__ import annotations

import numpy as np
import matplotlib
from PIL import Image, ImageDraw, ImageFont

matplotlib.use("Agg")
from matplotlib import colormaps  # noqa: E402


def load_font(size: int, bold: bool = False):
    candidates = [
        "DejaVuSans-Bold.ttf" if bold else "DejaVuSans.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    ]
    for candidate in candidates:
        try:
            return ImageFont.truetype(candidate, size)
        except OSError:
            continue
    return ImageFont.load_default()


def text_size(draw: ImageDraw.ImageDraw, text: str, font) -> tuple[int, int]:
    box = draw.textbbox((0, 0), text, font=font)
    return box[2] - box[0], box[3] - box[1]


def split_rows(mask: np.ndarray) -> tuple[int, int]:
    row_counts = mask.sum(axis=1)
    height = mask.shape[0]
    lo = int(height * 0.35)
    hi = int(height * 0.65)
    split = lo + int(np.argmin(row_counts[lo:hi]))
    return split, int(row_counts[split])


def half_mask(mask: np.ndarray, split: int, half: str) -> np.ndarray:
    out = np.zeros_like(mask, dtype=bool)
    if half == "left":
        out[: split + 1] = mask[: split + 1]
    elif half == "right":
        out[split + 1 :] = mask[split + 1 :]
    else:
        raise ValueError(f"Unknown hemisphere {half!r}")
    return out


def crop_box_for_half(mask: np.ndarray, split: int, half: str, margin: int) -> tuple[slice, slice]:
    hemi = half_mask(mask, split, half)
    return crop_box_for_mask(hemi, margin)


def crop_box_for_mask(mask: np.ndarray, margin: int) -> tuple[slice, slice]:
    rows, cols = np.where(mask)
    if rows.size == 0:
        raise ValueError("No pixels found for mask.")
    r0 = max(int(rows.min()) - margin, 0)
    r1 = min(int(rows.max()) + margin + 1, mask.shape[0])
    c0 = max(int(cols.min()) - margin, 0)
    c1 = min(int(cols.max()) + margin + 1, mask.shape[1])
    return slice(r0, r1), slice(c0, c1)


def colorize_scalar(values: np.ndarray, mask: np.ndarray, limit: float, cmap_name: str = "coolwarm") -> Image.Image:
    cmap = colormaps[cmap_name]
    normed = np.clip((values / max(limit, 1e-8) + 1.0) / 2.0, 0.0, 1.0)
    rgba = (cmap(normed) * 255).astype(np.uint8)
    rgba[~mask, 3] = 0
    rgba[mask & ~np.isfinite(values), 3] = 0
    return Image.fromarray(rgba, mode="RGBA")


def render_colorbar(width: int, height: int, cmap_name: str = "coolwarm") -> Image.Image:
    cmap = colormaps[cmap_name]
    gradient = np.linspace(-1.0, 1.0, width, dtype=np.float32)
    rgba = (cmap((gradient + 1.0) / 2.0) * 255).astype(np.uint8)
    bar = np.repeat(rgba[None, :, :], height, axis=0)
    return Image.fromarray(bar, mode="RGBA")


def colorize_positive(values: np.ndarray, mask: np.ndarray, limit: float, cmap_name: str = "magma") -> Image.Image:
    cmap = colormaps[cmap_name]
    normed = np.clip(values / max(limit, 1e-8), 0.0, 1.0)
    rgba = (cmap(normed) * 255).astype(np.uint8)
    rgba[~mask, 3] = 0
    rgba[mask & ~np.isfinite(values), 3] = 0
    return Image.fromarray(rgba, mode="RGBA")


def render_positive_colorbar(width: int, height: int, cmap_name: str = "magma") -> Image.Image:
    cmap = colormaps[cmap_name]
    gradient = np.linspace(0.0, 1.0, width, dtype=np.float32)
    rgba = (cmap(gradient) * 255).astype(np.uint8)
    bar = np.repeat(rgba[None, :, :], height, axis=0)
    return Image.fromarray(bar, mode="RGBA")