# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""Forecast evaluation: point and probabilistic scores, Diebold–Mariano (HAC), Giacomini–White, Model Confidence Set, block bootstrap."""
from __future__ import annotations

import numpy as np
from scipy import stats


# ---- point scores -----------------------------------------------------------------------------
def point_scores(y: np.ndarray, f: np.ndarray, y_naive: np.ndarray | None = None) -> dict:
    e = f - y; ok = np.isfinite(e); e = e[ok]
    out = {"n": int(ok.sum()), "mae": float(np.mean(np.abs(e))), "rmse": float(np.sqrt(np.mean(e ** 2))), "medae": float(np.median(np.abs(e))), "bias": float(np.mean(e))}
    if y_naive is not None:
        en = (y_naive - y)[ok]; out["mase"] = float(np.mean(np.abs(e)) / max(np.mean(np.abs(en)), 1e-12))
    return out


# ---- probabilistic scores ---------------------------------------------------------------------
def pinball(y: np.ndarray, q: np.ndarray, tau: float) -> np.ndarray:
    d = y - q
    return np.maximum(tau * d, (tau - 1) * d)


def crps_from_quantiles(y: np.ndarray, Q: np.ndarray, taus: np.ndarray) -> np.ndarray:
    """CRPS approximated by the average pinball loss over a quantile grid (×2), Laio & Tamea (2007) relation."""
    return 2 * np.mean(np.stack([pinball(y, Q[:, j], t) for j, t in enumerate(taus)], axis=1), axis=1)


def interval_scores(y: np.ndarray, lo: np.ndarray, hi: np.ndarray, alpha: float = 0.1) -> dict:
    cov = np.mean((y >= lo) & (y <= hi)); width = np.mean(hi - lo)
    wis = np.mean((hi - lo) + 2 / alpha * (lo - y) * (y < lo) + 2 / alpha * (y - hi) * (y > hi))
    return {"coverage": float(cov), "width": float(width), "interval_score": float(wis)}


def crps_gaussian(y: np.ndarray, mu: np.ndarray, sigma: np.ndarray) -> np.ndarray:
    z = (y - mu) / sigma
    return sigma * (z * (2 * stats.norm.cdf(z) - 1) + 2 * stats.norm.pdf(z) - 1 / np.sqrt(np.pi))


# ---- tests ------------------------------------------------------------------------------------
def hac_var(d: np.ndarray, lags: int) -> float:
    d = d - d.mean(); n = len(d); v = np.sum(d * d) / n
    for L in range(1, lags + 1):
        w = 1 - L / (lags + 1); v += 2 * w * np.sum(d[L:] * d[:-L]) / n
    return max(v, 1e-18)


def diebold_mariano(loss_a: np.ndarray, loss_b: np.ndarray, h: int = 1, small_sample: bool = True) -> dict:
    """DM test of H0: E[loss_a − loss_b] = 0 with Newey–West variance (h−1 lags) and the Harvey–Leybourne–Newbold correction."""
    d = loss_a - loss_b; d = d[np.isfinite(d)]; n = len(d)
    if n < 10:
        return {"dm": np.nan, "p": np.nan, "n": n, "mean_diff": float(np.mean(d)) if n else np.nan}
    v = hac_var(d, max(h - 1, 0)); dm = d.mean() / np.sqrt(v / n)
    if small_sample:
        dm *= np.sqrt((n + 1 - 2 * h + h * (h - 1) / n) / n)
        p = 2 * (1 - stats.t.cdf(abs(dm), df=n - 1))
    else:
        p = 2 * (1 - stats.norm.cdf(abs(dm)))
    return {"dm": float(dm), "p": float(p), "n": n, "mean_diff": float(d.mean())}


def giacomini_white(loss_a: np.ndarray, loss_b: np.ndarray, Z: np.ndarray | None = None, h: int = 1) -> dict:
    """Conditional predictive ability test: regress d_t on instruments Z_{t-1} (constant + lagged d by default); Wald test with HAC."""
    d = loss_a - loss_b; ok = np.isfinite(d); d = d[ok]
    if Z is None:
        Z = np.column_stack([np.ones(len(d) - 1), d[:-1]]); dd = d[1:]
    else:
        Z = Z[ok][1:]; dd = d[1:]
    n, q = Z.shape
    if n < 10 + q:
        return {"stat": np.nan, "p": np.nan, "n": n}
    Zd = Z * dd[:, None]; m = Zd.mean(axis=0)
    S = np.zeros((q, q)); lags = max(h - 1, 0)
    Zc = Zd - m
    S = Zc.T @ Zc / n
    for L in range(1, lags + 1):
        w = 1 - L / (lags + 1); G = Zc[L:].T @ Zc[:-L] / n; S += w * (G + G.T)
    try:
        stat = float(n * m @ np.linalg.solve(S, m))
    except np.linalg.LinAlgError:
        return {"stat": np.nan, "p": np.nan, "n": n}
    return {"stat": stat, "p": float(1 - stats.chi2.cdf(stat, q)), "n": n}


def stationary_bootstrap_indices(n: int, mean_block: float, rng: np.random.Generator) -> np.ndarray:
    p = 1 / mean_block; idx = np.empty(n, int); i = rng.integers(n)
    for t in range(n):
        if t > 0 and rng.random() > p:
            i = (i + 1) % n
        else:
            i = rng.integers(n)
        idx[t] = i
    return idx


def model_confidence_set(losses: np.ndarray, names: list[str], alpha: float = 0.10, B: int = 1000, mean_block: float = 3.0, seed: int = 0) -> dict:
    """Hansen–Lunde–Nason MCS with the T_max statistic and the stationary bootstrap. losses: (T, M)."""
    rng = np.random.default_rng(seed)
    L = losses[np.all(np.isfinite(losses), axis=1)]; T, M = L.shape
    alive = list(range(M)); pvals = {}; order = []
    boots = [stationary_bootstrap_indices(T, mean_block, rng) for _ in range(B)]
    while len(alive) > 1:
        Lk = L[:, alive]; dbar = Lk.mean(axis=0) - Lk.mean()
        # bootstrap distribution of the deviations
        bd = np.array([(Lk[b].mean(axis=0) - Lk[b].mean()) - dbar for b in boots])
        var = np.maximum(bd.var(axis=0), 1e-18)
        t_stat = dbar / np.sqrt(var); tmax = t_stat.max()
        tmax_b = (bd / np.sqrt(var)).max(axis=1)
        p = float(np.mean(tmax_b >= tmax))
        worst = int(np.argmax(t_stat)); name = names[alive[worst]]
        pvals[name] = max(p, max(pvals.values()) if pvals else 0.0); order.append(name)
        if p >= alpha:
            break
        alive.pop(worst)
    for i in alive:
        pvals.setdefault(names[i], 1.0)
    return {"mcs_pvalues": pvals, "in_mcs": [n for n, p in pvals.items() if p >= alpha], "eliminated_in_order": order}


def block_bootstrap_ci(values: np.ndarray, stat=np.mean, B: int = 1000, mean_block: float = 3.0, seed: int = 0, level: float = 0.95) -> tuple[float, float]:
    rng = np.random.default_rng(seed); v = values[np.isfinite(values)]; n = len(v)
    s = np.array([stat(v[stationary_bootstrap_indices(n, mean_block, rng)]) for _ in range(B)])
    return float(np.quantile(s, (1 - level) / 2)), float(np.quantile(s, 1 - (1 - level) / 2))