from __future__ import annotations

from collections import deque
from dataclasses import asdict, dataclass

import numpy as np


def normalised_difference(high: np.ndarray, low: np.ndarray) -> np.ndarray:
    high_value = np.asarray(high, dtype=np.float32)
    low_value = np.asarray(low, dtype=np.float32)
    denominator = high_value + low_value
    return np.divide(
        high_value - low_value,
        denominator,
        out=np.full_like(high_value, np.nan),
        where=np.abs(denominator) > 1e-8,
    )


def compute_indices(bands: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    required = {"green", "red", "red_edge", "nir"}
    if set(bands) != required:
        raise ValueError("bands must contain exactly green, red, red_edge, and nir")
    shapes = {np.asarray(value).shape for value in bands.values()}
    if len(shapes) != 1 or len(next(iter(shapes))) != 2:
        raise ValueError("all bands must be aligned two-dimensional arrays")
    return {
        "ndvi": normalised_difference(bands["nir"], bands["red"]),
        "ndre": normalised_difference(bands["nir"], bands["red_edge"]),
        "gndvi": normalised_difference(bands["nir"], bands["green"]),
    }


def robust_low_score(indices: dict[str, np.ndarray], valid: np.ndarray) -> np.ndarray:
    component_scores = []
    for name in ("ndvi", "ndre", "gndvi"):
        value = np.asarray(indices[name], dtype=np.float32)
        sample = value[valid & np.isfinite(value)]
        if sample.size < 100:
            raise ValueError("at least one hundred valid pixels are required")
        median = float(np.median(sample))
        mad = max(1.4826 * float(np.median(np.abs(sample - median))), 1e-4)
        component_scores.append(np.clip((median - value) / mad, 0, 6) / 6)
    score = np.mean(np.stack(component_scores), axis=0).astype(np.float32)
    score[~valid] = np.nan
    return score


@dataclass(frozen=True)
class ReviewZone:
    zone_id: str
    area_m2: int
    mean_score: float
    peak_score: float
    centroid_x_m: float
    centroid_y_m: float


def filter_connected_zones(
    mask: np.ndarray,
    score: np.ndarray,
    minimum_area_m2: int = 100,
) -> tuple[np.ndarray, list[ReviewZone]]:
    active = np.asarray(mask, dtype=bool)
    if active.shape != np.asarray(score).shape:
        raise ValueError("mask and score must have the same shape")
    if minimum_area_m2 < 1:
        raise ValueError("minimum_area_m2 must be positive")
    visited = np.zeros(active.shape, dtype=bool)
    retained = np.zeros(active.shape, dtype=bool)
    candidates: list[ReviewZone] = []
    for row, col in zip(*np.where(active)):
        if visited[row, col]:
            continue
        queue: deque[tuple[int, int]] = deque([(int(row), int(col))])
        visited[row, col] = True
        pixels: list[tuple[int, int]] = []
        while queue:
            r_value, c_value = queue.popleft()
            pixels.append((r_value, c_value))
            for r_next, c_next in ((r_value - 1, c_value), (r_value + 1, c_value), (r_value, c_value - 1), (r_value, c_value + 1)):
                if 0 <= r_next < active.shape[0] and 0 <= c_next < active.shape[1] and active[r_next, c_next] and not visited[r_next, c_next]:
                    visited[r_next, c_next] = True
                    queue.append((r_next, c_next))
        if len(pixels) < minimum_area_m2:
            continue
        for retained_row, retained_col in pixels:
            retained[retained_row, retained_col] = True
        values = np.array([score[r, c] for r, c in pixels], dtype=float)
        candidates.append(ReviewZone(
            zone_id="",
            area_m2=len(pixels),
            mean_score=round(float(values.mean()), 6),
            peak_score=round(float(values.max()), 6),
            centroid_x_m=round(float(np.mean([c + 0.5 for _, c in pixels])), 3),
            centroid_y_m=round(float(np.mean([r + 0.5 for r, _ in pixels])), 3),
        ))
    candidates.sort(key=lambda item: (-item.mean_score, -item.area_m2))
    zones = [
        ReviewZone(zone_id=f"RZ-{index:02d}", **{k: v for k, v in asdict(zone).items() if k != "zone_id"})
        for index, zone in enumerate(candidates, 1)
    ]
    return retained, zones


def connected_zones(mask: np.ndarray, score: np.ndarray, minimum_area_m2: int = 100) -> list[ReviewZone]:
    _, zones = filter_connected_zones(mask, score, minimum_area_m2)
    return zones
