"""Small, locally trainable colony-candidate detector for the PetriCount demo.

It deliberately uses only the annotated AGAR images that ship with this project.
The model ranks multiscale bright/dark blob candidates; it is a baseline for UI
integration, not a validated microbiology measurement method.
"""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter, maximum_filter

MAX_SIDE = 720


@dataclass
class PreparedImage:
    rgb: np.ndarray
    gray: np.ndarray
    scale: float
    original_size: tuple[int, int]


def prepare_image(path_or_file) -> PreparedImage:
    image = Image.open(path_or_file).convert("RGB")
    original_size = image.size
    scale = min(1.0, MAX_SIDE / max(image.size))
    if scale < 1:
        image = image.resize((round(image.width * scale), round(image.height * scale)), Image.Resampling.LANCZOS)
    rgb = np.asarray(image, dtype=np.float32) / 255.0
    gray = rgb[..., 0] * 0.2126 + rgb[..., 1] * 0.7152 + rgb[..., 2] * 0.0722
    return PreparedImage(rgb, gray, scale, original_size)


def _near_edge(x: int, y: int, width: int, height: int, margin: int = 13) -> bool:
    return x < margin or y < margin or x >= width - margin or y >= height - margin


def candidate_points(prepared: PreparedImage) -> list[tuple[int, int]]:
    """Find potential round colonies as local extrema at several blob scales."""
    gray = prepared.gray
    height, width = gray.shape
    candidates: list[tuple[float, int, int]] = []
    # Colonies can be lighter or darker than the surrounding agar.
    for sigma in (1.5, 2.5, 4.0, 6.0):
        local = gaussian_filter(gray, sigma) - gaussian_filter(gray, sigma * 3.2)
        radius = max(4, round(sigma * 2.5))
        for response in (local, -local):
            maxima = response == maximum_filter(response, size=radius * 2 + 1)
            threshold = np.percentile(response, 97.2)
            ys, xs = np.where(maxima & (response >= threshold))
            for y, x in zip(ys, xs):
                if _near_edge(int(x), int(y), width, height):
                    continue
                # Suppress the metal/background outside the central dish region.
                if np.hypot(x - width / 2, y - height / 2) > min(width, height) * 0.48:
                    continue
                candidates.append((float(response[y, x]), int(x), int(y)))
    candidates.sort(reverse=True)
    selected: list[tuple[int, int]] = []
    for _, x, y in candidates:
        if all((x - px) ** 2 + (y - py) ** 2 > 13 ** 2 for px, py in selected):
            selected.append((x, y))
        if len(selected) >= 450:
            break
    # Diffuse colonies often have no sharp local maximum. Add a coarse candidate
    # grid over visually changing parts of the dish and let the learned ranker
    # decide whether each patch is a colony or background.
    local_change = np.abs(gaussian_filter(gray, 2) - gaussian_filter(gray, 15))
    change_threshold = np.percentile(local_change, 58)
    for y in range(14, height - 14, 16):
        for x in range(14, width - 14, 16):
            if np.hypot(x - width / 2, y - height / 2) > min(width, height) * 0.46:
                continue
            if local_change[y, x] < change_threshold:
                continue
            if all((x - px) ** 2 + (y - py) ** 2 > 8 ** 2 for px, py in selected):
                selected.append((x, y))
    return selected


def _patch(array: np.ndarray, x: int, y: int, radius: int) -> np.ndarray:
    return array[max(0, y - radius):y + radius + 1, max(0, x - radius):x + radius + 1]


def features_for_points(prepared: PreparedImage, points: Iterable[tuple[int, int]]) -> np.ndarray:
    rows = []
    smooth_small = gaussian_filter(prepared.gray, 2)
    smooth_large = gaussian_filter(prepared.gray, 12)
    contrast = smooth_small - smooth_large
    gradient_y, gradient_x = np.gradient(prepared.gray)
    gradient = np.hypot(gradient_x, gradient_y)
    for x, y in points:
        inner_gray = _patch(prepared.gray, x, y, 4)
        outer_gray = _patch(prepared.gray, x, y, 12)
        inner_rgb = _patch(prepared.rgb, x, y, 4).reshape(-1, 3)
        outer_rgb = _patch(prepared.rgb, x, y, 12)
        inner_gradient = _patch(gradient, x, y, 8)
        rows.append([
            *inner_rgb.mean(axis=0), *inner_rgb.std(axis=0),
            float(inner_gray.mean()), float(inner_gray.std()),
            float(outer_gray.mean()), float(outer_gray.std()),
            float(contrast[y, x]), float(abs(contrast[y, x])),
            float(inner_gradient.mean()), float(inner_gradient.std()),
            x / prepared.gray.shape[1], y / prepared.gray.shape[0],
        ])
    return np.asarray(rows, dtype=np.float32)


def label_centers(labels: list[dict], scale: float) -> list[tuple[int, int, float]]:
    return [((item["x"] + item["width"] / 2) * scale, (item["y"] + item["height"] / 2) * scale,
             max(item["width"], item["height"]) * scale / 2) for item in labels]


def label_candidates(points: list[tuple[int, int]], labels: list[dict], scale: float) -> np.ndarray:
    centers = label_centers(labels, scale)
    labels_for_points = []
    for x, y in points:
        labels_for_points.append(any((x - cx) ** 2 + (y - cy) ** 2 <= max(7, radius * .78) ** 2 for cx, cy, radius in centers))
    return np.asarray(labels_for_points, dtype=np.int8)


def add_annotation_centers(points: list[tuple[int, int]], labels: list[dict], scale: float) -> list[tuple[int, int]]:
    result = list(points)
    for center_x, center_y, _ in label_centers(labels, scale):
        point = (round(center_x), round(center_y))
        if all((point[0] - px) ** 2 + (point[1] - py) ** 2 > 5 ** 2 for px, py in result):
            result.append(point)
    return result