"""Proper scoring rules and secondary metrics. Conventions: y = realised cumulative log return at horizon H (or the
per-minute return), forecasts are either parametric Student-t (closed form) or sample arrays [n, N]."""
from __future__ import annotations

import numpy as np
from scipy import special, stats


# ----------------------------------------------------------------------------- closed-form Student-t CRPS
def crps_student_t(y: np.ndarray, mu: np.ndarray, scale: np.ndarray, nu: np.ndarray) -> np.ndarray:
    """CRPS of a location-scale Student-t (Jordan, Krüger & Lerch 2019, scoringRules), valid for nu > 1."""
    y, mu, scale, nu = (np.asarray(a, dtype=np.float64) for a in (y, mu, scale, nu))
    z = (y - mu) / scale
    F = stats.t.cdf(z, nu); f = stats.t.pdf(z, nu)
    b1 = special.beta(0.5, nu - 0.5); b2 = special.beta(0.5, nu / 2) ** 2
    term = z * (2 * F - 1) + 2 * f * (nu + z * z) / (nu - 1) - 2 * np.sqrt(nu) * b1 / ((nu - 1) * b2)
    return scale * term


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


# ----------------------------------------------------------------------------- sample based
def crps_samples(y: np.ndarray, s: np.ndarray) -> np.ndarray:
    """Energy-form estimator from samples s [n, N] for observations y [N] (unbiased pairwise term)."""
    n = s.shape[0]
    t1 = np.abs(s - y[None, :]).mean(0)
    ss = np.sort(s, axis=0)
    # E|X - X'| via order statistics: (2/(n(n-1))) sum_i (2i - n - 1) x_(i)
    i = np.arange(1, n + 1)[:, None]
    t2 = (2.0 / (n * (n - 1))) * ((2 * i - n - 1) * ss).sum(0)
    return t1 - 0.5 * t2


def quantiles_from_samples(s: np.ndarray, levels: np.ndarray) -> np.ndarray:
    return np.quantile(s, levels, axis=0)   # [Q, N]


def quantile_loss(y: np.ndarray, q: np.ndarray, levels: np.ndarray) -> np.ndarray:
    """Mean pinball loss over levels; q [Q, N]."""
    d = y[None, :] - q
    return np.maximum(levels[:, None] * d, (levels[:, None] - 1) * d).mean(0)


def interval_score(y: np.ndarray, lo: np.ndarray, hi: np.ndarray, alpha: float) -> np.ndarray:
    return (hi - lo) + 2 / alpha * (lo - y) * (y < lo) + 2 / alpha * (y - hi) * (y > hi)


def weighted_interval_score(y: np.ndarray, s: np.ndarray, alphas=(0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)) -> np.ndarray:
    """WIS (Bracher et al. 2021) with K central intervals + median."""
    med = np.median(s, axis=0)
    K = len(alphas)
    tot = 0.5 * np.abs(y - med)
    for a in alphas:
        lo, hi = np.quantile(s, [a / 2, 1 - a / 2], axis=0)
        tot = tot + a / 2 * interval_score(y, lo, hi, a)
    return tot / (K + 0.5)


def coverage(y: np.ndarray, s: np.ndarray, level: float) -> tuple[float, float]:
    """Empirical coverage and mean width of the central `level` interval."""
    a = 1 - level
    lo, hi = np.quantile(s, [a / 2, 1 - a / 2], axis=0)
    return float(((y >= lo) & (y <= hi)).mean()), float((hi - lo).mean())


def pit_samples(y: np.ndarray, s: np.ndarray) -> np.ndarray:
    """Randomised PIT from samples (ties broken uniformly)."""
    n = s.shape[0]
    below = (s < y[None, :]).sum(0); equal = (s == y[None, :]).sum(0)
    u = np.random.default_rng(0).random(len(y))
    return (below + u * (equal + 1)) / (n + 1)


def log_score_student_t(y, mu, scale, nu) -> np.ndarray:
    return -stats.t.logpdf((y - mu) / scale, nu) + np.log(scale)


def log_score_kde(y: np.ndarray, s: np.ndarray) -> np.ndarray:
    """Negative log density from a Gaussian KDE of the samples (comparable only across sample-based forecasts)."""
    n = s.shape[0]
    sd = s.std(0, ddof=1) + 1e-12
    h = 1.06 * sd * n ** (-0.2)
    dens = stats.norm.pdf((y[None, :] - s) / h[None, :]).mean(0) / h
    return -np.log(dens + 1e-300)


# ----------------------------------------------------------------------------- secondary
def qlike(rv: np.ndarray, var_hat: np.ndarray) -> np.ndarray:
    v = np.maximum(var_hat, 1e-14)
    return np.log(v) + rv / v


def brier(event: np.ndarray, p: np.ndarray) -> np.ndarray:
    return (p - event.astype(float)) ** 2


def directional_accuracy(y: np.ndarray, mu: np.ndarray) -> float:
    nz = y != 0
    return float((np.sign(mu[nz]) == np.sign(y[nz])).mean()) if nz.any() else float("nan")


# ----------------------------------------------------------------------------- multivariate (paths)
def energy_score(Y: np.ndarray, S: np.ndarray) -> np.ndarray:
    """Y [N, H] observations, S [n, N, H] samples. ES = E||S - Y|| - 0.5 E||S - S'|| (pairs via a random permutation)."""
    n = S.shape[0]
    t1 = np.linalg.norm(S - Y[None], axis=-1).mean(0)
    perm = np.random.default_rng(1).permutation(n)
    t2 = np.linalg.norm(S - S[perm], axis=-1).mean(0)
    return t1 - 0.5 * t2


def variogram_score(Y: np.ndarray, S: np.ndarray, p: float = 0.5, max_pairs: int = 400) -> np.ndarray:
    """Variogram score of order p (Scheuerer & Hamill 2015) on a random subset of index pairs, equal weights."""
    N, H = Y.shape
    rng = np.random.default_rng(2)
    i = rng.integers(0, H, max_pairs); j = rng.integers(0, H, max_pairs)
    keep = i != j; i, j = i[keep], j[keep]
    obs = np.abs(Y[:, i] - Y[:, j]) ** p
    sim = (np.abs(S[:, :, i] - S[:, :, j]) ** p).mean(0)
    return ((obs - sim) ** 2).mean(-1)


def path_stats(S: np.ndarray) -> dict[str, np.ndarray]:
    """From sample paths of per-minute returns S [n, N, H]: cumulative path, max drawdown, realised vol."""
    cum = np.cumsum(S, axis=-1)
    peak = np.maximum.accumulate(np.concatenate([np.zeros_like(cum[..., :1]), cum], -1), axis=-1)[..., 1:]
    mdd = (peak - cum).max(-1)
    rv = (S ** 2).sum(-1)
    return {"cum": cum, "mdd": mdd, "rv": rv}