from __future__ import annotations
import numpy as np

THRESHOLDS = {'mild': -0.5, 'mord': -1.0, 'seve': -1.5, 'extr': -2.0}
LEVEL_TO_THRESHOLD = {'D1': -0.5, 'D2': -1.0, 'D3': -1.5, 'D4': -2.0}

def drought_probabilities(spei_members: np.ndarray, min_members: int = 1) -> dict[str, np.ndarray]:
    x = np.asarray(spei_members, dtype=float)  # member,time,lat,lon
    nvalid = np.sum(np.isfinite(x), axis=0).astype(float)
    out = {}
    for name, thr in THRESHOLDS.items():
        n = np.sum(np.isfinite(x) & (x <= thr), axis=0).astype(float)
        p = np.full(nvalid.shape, np.nan, dtype=np.float32)
        ok = nvalid >= int(min_members)
        p[ok] = 100.0 * n[ok] / nvalid[ok]
        out[name] = p.astype(np.float32)
    return out

def event_bool(spei: np.ndarray, level: str) -> np.ndarray:
    return np.asarray(spei) <= LEVEL_TO_THRESHOLD[level]