"""
SV vs GARCH out-of-sample comparison — 1-step-ahead predictive scoring + VaR backtest.

Fixed-parameter rolling filter: parameters are estimated once on a training window, then the
volatility filter is rolled forward over the test set producing genuine 1-step-ahead predictive
densities p(r_t | r_{1:t-1}). GARCH's predictive is closed form (a scaled Student-t / Normal);
SV's is NOT — because SV is a non-linear non-Gaussian state-space model, its predictive is obtained
with a bootstrap PARTICLE FILTER, and comes out as a scale-mixture of normals (fat-tailed for free).

Scoring: mean log predictive density (LPS, higher is better) and VaR/ES backtests
(Kupiec 1995 unconditional coverage, Christoffersen 1998 independence/conditional coverage).
"""
import numpy as np
from scipy.stats import t as student_t, norm, chi2
from scipy.optimize import minimize
import garch_scratch as G


# ---------------------------------------------------------------- GARCH predictive (closed form)
def garch_predict(r, theta, s0, gjr=False, alphas=(0.01, 0.05)):
    """1-step predictive log-density and VaR for every t, GARCH/GJR with Normal or Student-t innovations."""
    r = np.asarray(r, float); r2 = r ** 2
    neg = (r < 0).astype(float)
    s2 = G.garch_var(theta, r2, s0, neg, gjr)                 # conditional variance, s2[t] uses info < t
    o, a, b, g, nu = G._split(theta, gjr)
    lps = np.empty(len(r)); var = {al: np.empty(len(r)) for al in alphas}
    if np.isfinite(nu):
        scale = np.sqrt(s2 * (nu - 2.0) / nu)                 # so Var(r_t) = s2
        lps = student_t.logpdf(r, df=nu, scale=scale)
        for al in alphas:
            var[al] = scale * student_t.ppf(al, df=nu)
    else:
        sd = np.sqrt(s2)
        lps = norm.logpdf(r, scale=sd)
        for al in alphas:
            var[al] = sd * norm.ppf(al)
    return lps, var, np.sqrt(s2)


# ---------------------------------------------------------------- SV predictive (bootstrap particle filter)
def _systematic_resample(w, rng):
    N = len(w); pos = (rng.random() + np.arange(N)) / N
    cum = np.cumsum(w); cum[-1] = 1.0
    return np.searchsorted(cum, pos)


def _mix_quantile(scale, alpha, nu):
    """alpha-quantile (VaR, <0) of the zero-mean mixture (1/N) sum F(.; scale_i), by bisection.
    F is Student-t_nu (if nu finite) or Normal, each with the given per-particle scale."""
    cdf = (lambda x: student_t.cdf(x, df=nu)) if np.isfinite(nu) else norm.cdf
    lo, hi = -40.0, 0.0
    for _ in range(50):
        mid = 0.5 * (lo + hi)
        if np.mean(cdf(mid / scale)) < alpha:
            lo = mid
        else:
            hi = mid
    return 0.5 * (lo + hi)


def sv_particle_filter(r, mu, phi, sig, nu=np.inf, rho=0.0, N=3000, seed=0,
                       alphas=(0.01, 0.05), var_mask=None, rng=None):
    """Bootstrap particle filter for the SV model, optionally with Student-t innovations and leverage.

        r_t = e^{h_t/2} * eps_t,  eps_t ~ N(0,1) or standardized t_nu (unit variance)
        h_t = mu + phi (h_{t-1}-mu) + sigma * ( rho * u_{t-1} + sqrt(1-rho^2) * zeta_t )
    where u_{t-1} = r_{t-1} e^{-h_{t-1}/2} * sqrt((nu-2)/nu) is the unit-variance return shock, so rho<0
    is the leverage effect (a down-move raises tomorrow's volatility). Returns 1-step LPS, VaR, vol band.

    var_mask: VaR (the expensive bisection) is only computed where True; the filter runs over all t.
    """
    r = np.asarray(r, float); T = len(r)
    if rng is None:
        rng = np.random.default_rng(seed)
    student = np.isfinite(nu)
    snu = np.sqrt((nu - 2.0) / nu) if student else 1.0               # scale so Var(r_t|h_t)=e^{h_t}
    h = mu + rng.standard_normal(N) * sig / np.sqrt(1 - phi ** 2)     # stationary init (posterior at t-1)
    lps = np.empty(T); var = {al: np.full(T, np.nan) for al in alphas}
    vol = np.empty(T); vlo = np.empty(T); vhi = np.empty(T)
    if var_mask is None:
        var_mask = np.ones(T, bool)
    r_prev = None
    for t in range(T):
        if rho != 0.0 and r_prev is not None:                        # leverage: correlate with last shock
            u = r_prev * np.exp(-h / 2) * snu
            hp = mu + phi * (h - mu) + sig * (rho * u + np.sqrt(1 - rho ** 2) * rng.standard_normal(N))
        else:
            hp = mu + phi * (h - mu) + sig * rng.standard_normal(N)   # predict h_t | past
        scale = np.exp(hp / 2) * snu
        vol[t] = np.exp(hp / 2).mean(); vlo[t], vhi[t] = np.quantile(np.exp(hp / 2), [0.05, 0.95])
        if student:
            logp = student_t.logpdf(r[t], df=nu, scale=scale)
        else:
            logp = norm.logpdf(r[t], scale=np.exp(hp / 2))
        m = logp.max(); lse = m + np.log(np.sum(np.exp(logp - m)))
        lps[t] = lse - np.log(N)                                      # predictive = mean of components
        if var_mask[t]:
            sc = scale if student else np.exp(hp / 2)
            for al in alphas:
                var[al][t] = _mix_quantile(sc, al, nu)
        w = np.exp(logp - lse)
        h = hp[_systematic_resample(w, rng)]                         # resample -> posterior at t
        r_prev = r[t]
    return lps, var, dict(vol=vol, lo=vlo, hi=vhi)


def sv_fit_pf(r, N=2000, seed=0, student=True, leverage=True):
    """Estimate SV (optionally Student-t / leverage) on a training series by maximizing the
    common-random-number particle-filter likelihood (simulated ML). Returns (mu, phi, sigma, nu, rho)."""
    r = np.asarray(r, float)

    def unpack(x):
        mu = x[0]; phi = np.tanh(x[1]); sig = np.exp(x[2])
        nu = 2.0 + np.exp(x[3]) if student else np.inf
        rho = np.tanh(x[4]) if leverage else 0.0
        return mu, phi, sig, nu, rho

    def nll(x):
        mu, phi, sig, nu, rho = unpack(x)
        rng = np.random.default_rng(seed)                            # CRN: identical noise every evaluation
        lps, _, _ = sv_particle_filter(r, mu, phi, sig, nu, rho, N=N, alphas=(), rng=rng)
        return -np.sum(lps)

    x0 = np.array([np.log(np.var(r)), np.arctanh(0.97), np.log(0.15), np.log(8.0), np.arctanh(-0.4)])
    res = minimize(nll, x0, method="Nelder-Mead",
                   options=dict(xatol=1e-3, fatol=1e-2, maxiter=2000))
    return unpack(res.x)


# ---------------------------------------------------------------- VaR backtests
def kupiec(viol, alpha):
    n = len(viol); x = int(viol.sum()); pihat = x / n
    if x == 0:
        lr = -2 * n * np.log(1 - alpha)
    else:
        lr = -2 * ((n - x) * np.log(1 - alpha) + x * np.log(alpha)
                   - (n - x) * np.log(1 - pihat) - x * np.log(pihat))
    return lr, 1 - chi2.cdf(lr, 1), pihat


def christoffersen(viol, alpha):
    v = viol.astype(int); n = len(v)
    n00 = n01 = n10 = n11 = 0
    for i in range(1, n):
        a, b = v[i - 1], v[i]
        n00 += (a == 0 and b == 0); n01 += (a == 0 and b == 1)
        n10 += (a == 1 and b == 0); n11 += (a == 1 and b == 1)
    p01 = n01 / max(n00 + n01, 1); p11 = n11 / max(n10 + n11, 1)
    p = (n01 + n11) / max(n00 + n01 + n10 + n11, 1)
    def ll(pi, k0, k1): return (k0 * np.log(1 - pi) if pi < 1 and k0 else 0) + (k1 * np.log(pi) if pi > 0 and k1 else 0)
    lr_ind = -2 * (ll(p, n00 + n10, n01 + n11) - ll(p01, n00, n01) - ll(p11, n10, n11))
    lr_uc = kupiec(viol, alpha)[0]
    lr_cc = lr_uc + lr_ind
    return lr_cc, 1 - chi2.cdf(lr_cc, 2)


def backtest(r_test, var_test, alpha):
    """Return dict: violation rate, Kupiec p, Christoffersen (cc) p, realized ES (mean breach loss)."""
    viol = r_test < var_test
    _, p_kup, pihat = kupiec(viol, alpha)
    _, p_cc = christoffersen(viol, alpha)
    es = r_test[viol].mean() if viol.any() else np.nan
    return dict(rate=pihat, n_viol=int(viol.sum()), p_kupiec=p_kup, p_cc=p_cc, es=es)