"""Verification metrics for probabilistic and binary EHD forecasts.

All functions drop NaN pairs. Probabilities are P(event); binaries use a decision threshold.
Skill scores are (score - ref)/(perfect - ref) with the sign set so higher = better.
"""
from __future__ import annotations
import numpy as np


def _clean(y, p):
    y = np.asarray(y, float); p = np.asarray(p, float)
    m = np.isfinite(y) & np.isfinite(p)
    return y[m], p[m]


def brier(y, p):
    y, p = _clean(y, p)
    return float(np.mean((p - y) ** 2)) if y.size else np.nan


def brier_skill_score(y, p, p_ref):
    bs = brier(y, p)
    y2, pr = _clean(y, p_ref)
    bs_ref = float(np.mean((pr - y2) ** 2)) if y2.size else np.nan
    return float(1 - bs / bs_ref) if bs_ref and bs_ref > 0 else np.nan


def reliability_curve(y, p, nbins=10):
    y, p = _clean(y, p)
    edges = np.linspace(0, 1, nbins + 1)
    idx = np.clip(np.digitize(p, edges[1:-1]), 0, nbins - 1)
    conf, obs, cnt = [], [], []
    for b in range(nbins):
        m = idx == b
        cnt.append(int(m.sum()))
        conf.append(float(p[m].mean()) if m.any() else np.nan)
        obs.append(float(y[m].mean()) if m.any() else np.nan)
    return np.array(conf), np.array(obs), np.array(cnt)


def auc(y, p):
    """ROC AUC via the rank (Mann-Whitney) identity; ties handled by average rank."""
    y, p = _clean(y, p)
    n1 = int(y.sum()); n0 = int((y == 0).sum())
    if n1 == 0 or n0 == 0:
        return np.nan
    order = np.argsort(p, kind="mergesort")
    ranks = np.empty_like(order, float)
    ranks[order] = np.arange(1, len(p) + 1)
    # average ranks for ties
    sp = p[order]
    i = 0
    while i < len(sp):
        j = i
        while j + 1 < len(sp) and sp[j + 1] == sp[i]:
            j += 1
        if j > i:
            ranks[order[i:j + 1]] = (i + 1 + j + 1) / 2.0
        i = j + 1
    sum_pos = ranks[y == 1].sum()
    return float((sum_pos - n1 * (n1 + 1) / 2.0) / (n1 * n0))


def contingency(y, p, thr):
    y, p = _clean(y, p)
    f = p >= thr
    a = int(np.sum(f & (y == 1)))    # hits
    b = int(np.sum(f & (y == 0)))    # false alarms
    c = int(np.sum(~f & (y == 1)))   # misses
    d = int(np.sum(~f & (y == 0)))   # correct negatives
    return a, b, c, d


def binary_scores_from_counts(a, b, c, d):
    """All binary contingency-table skill scores from raw counts (hits a, false alarms b,
    misses c, correct negatives d). Used to pool per-fold / per-region decision thresholds:
    accumulate a,b,c,d across samples that each used their own threshold, then score once."""
    n = a + b + c + d
    eps = 1e-12
    pod = a / (a + c + eps)
    pofd = b / (b + d + eps)
    far = b / (a + b + eps)
    csi = a / (a + b + c + eps)
    bias = (a + b) / (a + c + eps)
    peirce = pod - pofd
    ar = (a + b) * (a + c) / (n + eps)          # random hits for ETS
    ets = (a - ar) / (a + b + c - ar + eps)
    hr = ((a + d) / (n + eps))
    he = ((a + b) * (a + c) + (c + d) * (b + d)) / (n * n + eps)
    hss = (hr - he) / (1 - he + eps)
    # extremal dependence indices
    F = b / (b + d + eps); H = pod
    edi = ((np.log(F + eps) - np.log(H + eps)) /
           (np.log(F + eps) + np.log(H + eps) + eps))
    sedi = ((np.log(F + eps) - np.log(H + eps) - np.log(1 - F + eps) + np.log(1 - H + eps)) /
            (np.log(F + eps) + np.log(H + eps) + np.log(1 - F + eps) + np.log(1 - H + eps) + eps))
    return dict(POD=pod, POFD=pofd, FAR=far, CSI=csi, freq_bias=bias, Peirce=peirce,
                ETS=ets, HSS=hss, EDI=edi, SEDI=sedi, hits=a, false_alarms=b,
                misses=c, correct_neg=d, base_rate=(a + c) / (n + eps))


def binary_scores(y, p, thr):
    a, b, c, d = contingency(y, p, thr)
    return binary_scores_from_counts(a, b, c, d)


def best_threshold(y, p, metric="Peirce", grid=None):
    """Pick the probability threshold maximizing a binary metric on (y, p)."""
    if grid is None:
        grid = np.linspace(0.02, 0.98, 49)
    best, best_thr = -np.inf, 0.5
    for t in grid:
        s = binary_scores(y, p, t)[metric]
        if np.isfinite(s) and s > best:
            best, best_thr = s, t
    return float(best_thr), float(best)


# ---- continuous metrics (for the Tx MOS side) ----
def mae(o, f):
    o, f = _clean(o, f); return float(np.mean(np.abs(f - o))) if o.size else np.nan


def rmse(o, f):
    o, f = _clean(o, f); return float(np.sqrt(np.mean((f - o) ** 2))) if o.size else np.nan


def bias(o, f):
    o, f = _clean(o, f); return float(np.mean(f - o)) if o.size else np.nan


def crps_gaussian(o, mu, sigma):
    """CRPS for a Gaussian predictive distribution (Gneiting & Raftery 2007)."""
    from scipy.stats import norm
    o = np.asarray(o, float); mu = np.asarray(mu, float); sigma = np.asarray(sigma, float)
    m = np.isfinite(o) & np.isfinite(mu) & np.isfinite(sigma) & (sigma > 0)
    o, mu, sigma = o[m], mu[m], sigma[m]
    z = (o - mu) / sigma
    return float(np.mean(sigma * (z * (2 * norm.cdf(z) - 1) + 2 * norm.pdf(z) - 1 / np.sqrt(np.pi))))