"""Case 02 core algorithms: SPI-3 drought and spatial clustering.

The implementation follows the usual SPI workflow: rolling precipitation
accumulation, calendar-month Gamma calibration with a zero-precipitation
probability, and the standard-normal transform.  Very short teaching arrays
cannot support a calendar-month distribution fit; in that case the function
uses a clearly documented global z-score fallback so that the example remains
runnable while real analyses still require a sufficiently long baseline.
"""
from __future__ import annotations

import numpy as np
from scipy.stats import gamma, norm


def _rolling_sum(x: np.ndarray, period: int) -> np.ndarray:
    """Return a trailing sum with the first ``period-1`` observations NaN."""
    n, p = x.shape
    out = np.full((n, p), np.nan, dtype=float)
    for i in range(period - 1, n):
        window = x[i - period + 1 : i + 1]
        valid = np.isfinite(window).sum(axis=0)
        summed = np.nansum(window, axis=0)
        out[i] = np.where(valid == period, summed, np.nan)
    return out


def _normal_fallback(values: np.ndarray) -> np.ndarray:
    """Standardize a short series; used only when Gamma calibration is underpowered."""
    finite = np.isfinite(values)
    if finite.sum() < 2:
        return np.full(values.shape, np.nan, dtype=float)
    mu = np.nanmean(values)
    sd = np.nanstd(values, ddof=1)
    if not np.isfinite(sd) or sd <= 0:
        return np.where(finite, 0.0, np.nan)
    return np.divide(values - mu, sd, out=np.full(values.shape, np.nan), where=finite)


def spi3(precip, period: int = 3, min_samples: int = 5, calibration_start: int = 0):
    """Calculate SPI for a time-by-cell monthly precipitation array.

    Parameters
    ----------
    precip : array-like, shape (time, cell)
        Monthly precipitation in a consistent unit (normally millimetres).
    period : int, default 3
        Accumulation window.  ``3`` gives SPI-3.
    min_samples : int, default 5
        Minimum finite values per calendar month and cell for Gamma fitting.
        Fewer values use a global normalised fallback for short teaching data.
    calibration_start : int, default 0
        Index of the first observation allowed in the calibration baseline.

    Notes
    -----
    For each calendar month the positive accumulated values are fitted to a
    Gamma distribution (location fixed at zero).  The probability of zero
    accumulation is mixed into the CDF before applying ``norm.ppf``.  The
    first ``period-1`` rows are NaN because no complete accumulation exists.
    """
    x = np.asarray(precip, dtype=float)
    if x.ndim != 2 or x.shape[0] < period:
        raise ValueError("precip must be time x cell with at least period rows")
    if period < 1 or int(period) != period:
        raise ValueError("period must be a positive integer")
    if np.nanmin(x) < 0:
        raise ValueError("precipitation cannot be negative")
    period = int(period)
    accum = _rolling_sum(x, period)
    out = np.full_like(accum, np.nan, dtype=float)
    n = x.shape[0]
    start = max(0, int(calibration_start))

    # In a small teaching array a seasonal fit is not identifiable.  Retain
    # an explicit fallback rather than silently pretending it is Gamma SPI.
    if n < max(12, period + min_samples):
        for j in range(x.shape[1]):
            out[:, j] = _normal_fallback(accum[:, j])
        return out

    for j in range(x.shape[1]):
        for month in range(12):
            idx = np.arange(period - 1, n)
            idx = idx[(idx >= start) & ((idx % 12) == month)]
            vals = accum[idx, j]
            vals = vals[np.isfinite(vals)]
            if vals.size < min_samples:
                # A robust, documented fallback for a sparse calendar-month
                # baseline; this is still standardisation, not a fitted SPI.
                target = accum[idx, j]
                out[idx, j] = _normal_fallback(target)
                continue
            zero_prob = float(np.mean(vals <= 0))
            positive = vals[vals > 0]
            if positive.size < 2 or np.nanstd(positive) <= 1e-12:
                target = accum[idx, j]
                out[idx, j] = _normal_fallback(target)
                continue
            try:
                shape, _, scale = gamma.fit(positive, floc=0)
                target = accum[idx, j]
                cdf = np.full(target.shape, np.nan, dtype=float)
                finite = np.isfinite(target)
                cdf[finite] = zero_prob + (1.0 - zero_prob) * gamma.cdf(
                    np.maximum(target[finite], 0.0), a=shape, loc=0, scale=scale
                )
                cdf = np.clip(cdf, 1e-8, 1 - 1e-8)
                out[idx, j] = norm.ppf(cdf)
            except (ValueError, FloatingPointError):
                out[idx, j] = _normal_fallback(accum[idx, j])
    return out


def drought_frequency(spi, threshold=-1.0):
    z = np.asarray(spi, float)
    # Comparisons with NaN return False, so mask invalid months explicitly.
    # Otherwise warm-up months from the rolling SPI window would be counted
    # as non-drought months and bias the frequency downward.
    valid = np.isfinite(z)
    drought = (z <= threshold) & valid
    denominator = valid.sum(axis=0)
    return np.divide(
        drought.sum(axis=0),
        denominator,
        out=np.full(denominator.shape, np.nan, dtype=float),
        where=denominator > 0,
    )


def moran_i(values, weights):
    x = np.asarray(values, float)
    W = np.asarray(weights, float)
    if x.ndim != 1 or W.shape != (x.size, x.size):
        raise ValueError("values and weights shape mismatch")
    z = x - x.mean()
    s0 = W.sum()
    return float(x.size / s0 * (z @ W @ z) / (z @ z)) if s0 and z @ z else float("nan")


def _bh_fdr(p_values: np.ndarray) -> np.ndarray:
    """Benjamini-Hochberg adjusted p-values, preserving original order."""
    p = np.asarray(p_values, float)
    q = np.full(p.shape, np.nan)
    finite = np.isfinite(p)
    if not finite.any():
        return q
    vals = p[finite]
    order = np.argsort(vals)
    ranked = vals[order]
    adjusted = ranked * len(ranked) / np.arange(1, len(ranked) + 1)
    adjusted = np.minimum.accumulate(adjusted[::-1])[::-1]
    adjusted = np.clip(adjusted, 0, 1)
    restored = np.empty_like(adjusted)
    restored[order] = adjusted
    q[finite] = restored
    return q


def lisa_quadrants(values, weights):
    """Return HH/LL/HL/LH labels from the centred value and spatial lag."""
    x = np.asarray(values, float)
    W = np.asarray(weights, float)
    if x.ndim != 1 or W.shape != (x.size, x.size):
        raise ValueError("values and weights shape mismatch")
    z = x - x.mean()
    lag = W @ z
    return np.where((z >= 0) & (lag >= 0), "HH", np.where((z < 0) & (lag < 0), "LL", np.where(z >= 0, "HL", "LH")))


def lisa_statistics(values, weights, permutations: int = 499, seed: int = 0, fdr_alpha: float = 0.05):
    """Local Moran's I permutation p-values with Benjamini-Hochberg FDR.

    ``weights`` may be binary or row-standardised.  The returned dictionary
    contains the local statistic, two-sided permutation p-value, FDR-adjusted
    q-value, and a boolean significance mask.  Quadrant labels remain
    available through :func:`lisa_quadrants` and should be interpreted only
    where ``q_value <= fdr_alpha``.
    """
    x = np.asarray(values, float)
    W = np.asarray(weights, float)
    if x.ndim != 1 or W.shape != (x.size, x.size):
        raise ValueError("values and weights shape mismatch")
    if np.isnan(x).any() or np.isnan(W).any():
        raise ValueError("values and weights cannot contain NaN")
    row_sum = W.sum(axis=1, keepdims=True)
    Wn = np.divide(W, row_sum, out=np.zeros_like(W), where=row_sum != 0)
    z = x - x.mean()
    local_i = z * (Wn @ z)
    p = np.full(x.shape, np.nan, dtype=float)
    rng = np.random.default_rng(seed)
    permutations = int(permutations)
    if permutations <= 0:
        raise ValueError("permutations must be positive")
    for i in range(x.size):
        neighbor_slots = np.flatnonzero(Wn[i] > 0)
        if neighbor_slots.size == 0:
            continue
        other_values = np.delete(z, i)
        active_weights = Wn[i, neighbor_slots]
        permuted = np.empty(permutations, dtype=float)
        for k in range(permutations):
            # Conditional randomisation: keep the focal value fixed and draw
            # neighbor values without replacement from all other locations.
            sampled = rng.choice(other_values, size=neighbor_slots.size, replace=False)
            permuted[k] = z[i] * np.dot(active_weights, sampled)
        exceed = np.count_nonzero(np.abs(permuted) >= abs(local_i[i]))
        p[i] = (exceed + 1) / (permutations + 1)
    q = _bh_fdr(p)
    return {"local_i": local_i, "p_value": p, "q_value": q, "significant": q <= float(fdr_alpha)}