from __future__ import annotations

from typing import Dict, Iterable, Tuple

import numpy as np

CATEGORY_CODES = {
    'D4': -4,
    'D3': -3,
    'D2': -2,
    'D1': -1,
    'D0': 0,
    'W1': 1,
    'W2': 2,
    'W3': 3,
    'W4': 4,
}


def spei_category(spei: np.ndarray) -> np.ndarray:
    """Map SPEI to drought/wetness category codes.

    D1: -1.0 < SPEI <= -0.5
    D2: -1.5 < SPEI <= -1.0
    D3: -2.0 < SPEI <= -1.5
    D4: SPEI <= -2.0
    D0: -0.5 < SPEI < 0.5
    W1-W4 analogous positive thresholds.
    """
    x = np.asarray(spei, dtype=float)
    out = np.full(x.shape, -99, dtype=np.int16)
    valid = np.isfinite(x)
    out[valid & (x <= -2.0)] = CATEGORY_CODES['D4']
    out[valid & (x > -2.0) & (x <= -1.5)] = CATEGORY_CODES['D3']
    out[valid & (x > -1.5) & (x <= -1.0)] = CATEGORY_CODES['D2']
    out[valid & (x > -1.0) & (x <= -0.5)] = CATEGORY_CODES['D1']
    out[valid & (x > -0.5) & (x < 0.5)] = CATEGORY_CODES['D0']
    out[valid & (x >= 0.5) & (x < 1.0)] = CATEGORY_CODES['W1']
    out[valid & (x >= 1.0) & (x < 1.5)] = CATEGORY_CODES['W2']
    out[valid & (x >= 1.5) & (x < 2.0)] = CATEGORY_CODES['W3']
    out[valid & (x >= 2.0)] = CATEGORY_CODES['W4']
    return out


def drought_exceedance_probability(
    spei_members: np.ndarray,
    thresholds: Iterable[float] = (-0.5, -1.0, -1.5, -2.0),
    min_members: int = 1,
) -> Dict[str, np.ndarray]:
    """Compute P(SPEI <= threshold) from member dimension.

    spei_members shape must be (member, time, lat, lon) or (member, lat, lon).
    Missing members are ignored.  If the number of valid members is below
    min_members, probability is missing, never interpreted as no drought.
    """
    x = np.asarray(spei_members, dtype=float)
    if x.ndim not in (3, 4):
        raise ValueError('spei_members must have member dimension first and 3 or 4 dimensions')
    valid_count = np.sum(np.isfinite(x), axis=0).astype(float)
    probs: Dict[str, np.ndarray] = {}
    for thr in thresholds:
        n = np.sum(np.isfinite(x) & (x <= float(thr)), axis=0).astype(float)
        p = np.full(valid_count.shape, np.nan, dtype=np.float32)
        ok = valid_count >= int(min_members)
        p[ok] = (100.0 * n[ok] / valid_count[ok]).astype(np.float32)
        name = f'p_spei_le_{str(thr).replace("-", "m").replace(".", "p")}'
        probs[name] = p
    return probs