#!/usr/bin/env python3
"""Contingency-table scores shared by every verification path in the project.

Keeping one implementation means the AI models of stage 2 are scored with exactly
the same definitions as the raw EC forecast of stage 1.
"""

from __future__ import annotations

import numpy as np


def contingency(pred, obs):
    """Counts (hit, false alarm, miss, correct negative) of two boolean arrays."""
    pred = pred.astype(bool); obs = obs.astype(bool)
    return (int(np.sum(pred & obs)), int(np.sum(pred & ~obs)),
            int(np.sum(~pred & obs)), int(np.sum(~pred & ~obs)))


def _safe(num, den):
    return float(num) / float(den) if den else np.nan


def edi_sedi(H, F):
    """Extremal Dependence Index and its symmetric form, from POD (H) and POFD (F).

    CSI and ETS both degenerate towards zero as the event gets rarer -- their value
    depends on the base rate as much as on the forecast -- so a system cannot be
    compared across sub-regions or across lead weeks with different climatologies
    using them alone.  EDI and SEDI are base-rate independent and non-degenerate in
    the rare-event limit; both are 0 for a random forecast and 1 for a perfect one.
    SEDI is additionally invariant when hits and correct negatives swap roles, so it
    cannot be improved by hedging the way EDI can.

        EDI  = (ln F - ln H) / (ln F + ln H)
        SEDI = [ln F - ln H - ln(1-F) + ln(1-H)]
               / [ln F + ln H + ln(1-F) + ln(1-H)]

    Undefined when H or F hits exactly 0 or 1, which is why they are returned as NaN
    there rather than clipped: a clipped value would silently invent skill.
    """
    if not (0.0 < H < 1.0) or not (0.0 < F < 1.0):
        return np.nan, np.nan
    lh, lf = np.log(H), np.log(F)
    edi = (lf - lh) / (lf + lh)
    l1h, l1f = np.log(1.0 - H), np.log(1.0 - F)
    den = lf + lh + l1f + l1h
    sedi = (lf - lh - l1f + l1h) / den if den != 0 else np.nan
    return float(edi), float(sedi)


def scores(a, b, c, d):
    """POD, POFD, FAR, FBI, CSI, ETS, EDI, SEDI and the base rate."""
    n = a + b + c + d
    ar = _safe((a + b) * (a + c), n)
    _edi, _sedi = edi_sedi(_safe(a, a + c), _safe(b, b + d))
    return {
        "EDI": _edi,
        "SEDI": _sedi,
        "POD": _safe(a, a + c),
        "POFD": _safe(b, b + d),
        "FAR": _safe(b, a + b),
        "FBI": _safe(a + b, a + c),
        "CSI": _safe(a, a + b + c),
        "ETS": _safe(a - ar, (a + b + c - ar)) if not np.isnan(ar) else np.nan,
        "base_rate": _safe(a + c, n),
        "n": n,
    }


def score_binary(pred, obs):
    return scores(*contingency(pred, obs))


def auc_score(score, obs):
    """Rank-based ROC AUC; ties get the average rank."""
    obs = obs.astype(bool)
    ny, nn = int(obs.sum()), int((~obs).sum())
    if ny == 0 or nn == 0:
        return np.nan
    order = np.argsort(score, kind="mergesort")
    ranks = np.empty(score.size, dtype="f8")
    ranks[order] = np.arange(1, score.size + 1)
    s = np.sort(score)
    i = 0
    while i < s.size:                       # average the ranks inside each tie group
        j = i
        while j + 1 < s.size and s[j + 1] == s[i]:
            j += 1
        if j > i:
            ranks[order[i:j + 1]] = 0.5 * (i + 1 + j + 1)
        i = j + 1
    return (ranks[obs].sum() - ny * (ny + 1) / 2.0) / (ny * nn)


def best_threshold(prob, obs, metric="CSI", grid=None):
    """Threshold on `prob` maximising `metric`, plus the value reached."""
    if grid is None:
        q = np.unique(np.quantile(prob, np.linspace(0.50, 0.9995, 200)))
        grid = q
    best, best_t = -np.inf, float(np.median(prob))
    for t in grid:
        s = score_binary(prob >= t, obs)[metric]
        if np.isfinite(s) and s > best:
            best, best_t = s, float(t)
    return best_t, best


def reliability(prob, obs, nbin=10):
    """Reliability-diagram points: mean forecast probability vs observed frequency."""
    edges = np.linspace(0.0, 1.0, nbin + 1)
    idx = np.clip(np.digitize(prob, edges) - 1, 0, nbin - 1)
    out = []
    for k in range(nbin):
        m = idx == k
        if m.sum():
            out.append((float(prob[m].mean()), float(obs[m].mean()), int(m.sum())))
        else:
            out.append((np.nan, np.nan, 0))
    return out


def brier_skill_score(prob, obs):
    """Brier skill score against the sample climatology."""
    obs = obs.astype("f8")
    bs = np.mean((prob - obs) ** 2)
    base = obs.mean()
    bs_ref = np.mean((base - obs) ** 2)
    return np.nan if bs_ref == 0 else 1.0 - bs / bs_ref