"""
GARCH(1,1) from scratch — five ways to estimate the same model.
Symmetric OR asymmetric (GJR leverage), Normal OR Student-t innovations.

Variance recursion:
    symmetric (GARCH):   sigma2_t = omega + alpha * r2_{t-1}                 + beta * sigma2_{t-1}
    asymmetric (GJR):    sigma2_t = omega + (alpha + gamma*1[r_{t-1}<0]) r2_{t-1} + beta * sigma2_{t-1}
Innovations: eps_t ~ N(0,1)  or  standardized t_nu (Var 1).

Parameter vector theta (order):  [omega, alpha, beta] + ([gamma] if gjr) + ([nu] if Student-t]
    e.g. GARCH-N (3), GARCH-t (4), GJR-N (4), GJR-t (5).
Admissible: omega>0, alpha>=0, beta>=0, alpha+gamma>=0, alpha + gamma/2 + beta < 1, (nu>2).

The five estimators read the layout from the `gjr` flag + len(theta):
    garch_mle · gh_quadrature · importance_sampling · metropolis_hastings · griddy_gibbs
Prior: flat on the admissible region, so log-posterior = log-likelihood.  Recursion via lfilter -> O(T).
Backward compatible: with gjr=False (default) and no `neg`, the symmetric API is unchanged.
"""
import numpy as np
from scipy.signal import lfilter
from scipy.optimize import minimize
from scipy.special import gammaln
from scipy.stats import multivariate_t

NAMES_N = ["omega", "alpha", "beta"]
NAMES_T = ["omega", "alpha", "beta", "nu"]


def names(gjr=False, student=False):
    return ["omega", "alpha", "beta"] + (["gamma"] if gjr else []) + (["nu"] if student else [])


def _split(theta, gjr):
    """Return omega, alpha, beta, gamma, nu (nu=inf if Gaussian) from theta given the gjr flag."""
    o, a, b = theta[0], theta[1], theta[2]; i = 3
    g = theta[i] if gjr else 0.0; i += int(gjr)
    nu = theta[i] if len(theta) > i else np.inf
    return o, a, b, g, nu


# ---------------------------------------------------------------- core likelihood
def garch_var(theta, r2, s0, neg=None, gjr=False):
    o, a, b, g, _ = _split(theta, gjr)
    lev = (a + g * neg[:-1]) if gjr else a
    v = o + lev * r2[:-1]
    rest, _ = lfilter([1.0], [1.0, -b], v, zi=[b * s0])
    return np.concatenate([[s0], rest])


def _feasible(theta, gjr):
    o, a, b, g, nu = _split(theta, gjr)
    return (o > 0) and (a >= 0) and (b >= 0) and (a + g >= 0) and (a + g / 2 + b < 1.0) and (nu > 2.0)


def garch_loglik(theta, r2, s0, neg=None, gjr=False):
    """GARCH/GJR (x) Normal/Student-t log-likelihood; -inf outside the admissible region."""
    theta = np.asarray(theta, float)
    if not _feasible(theta, gjr):
        return -np.inf
    s2 = garch_var(theta, r2, s0, neg, gjr)
    if np.any(s2 <= 0):
        return -np.inf
    _, _, _, _, nu = _split(theta, gjr)
    if np.isfinite(nu):
        sc2 = s2 * (nu - 2.0) / nu                       # scale^2 so that Var(r_t) = s2
        return np.sum(gammaln((nu + 1) / 2) - gammaln(nu / 2)
                      - 0.5 * np.log(nu * np.pi * sc2)
                      - (nu + 1) / 2 * np.log1p(r2 / (nu * sc2)))
    return -0.5 * np.sum(np.log(2 * np.pi * s2) + r2 / s2)


def prepare(r):
    r = np.asarray(r, float)
    return r ** 2, float(np.var(r))


# ---------------------------------------------------------------- 1. MLE
def _num_hess(f, x, eps=1e-4):
    n = len(x); H = np.zeros((n, n)); fx = f(x)
    for i in range(n):
        for j in range(i, n):
            xi = x.copy(); xj = x.copy(); xij = x.copy()
            xi[i] += eps; xj[j] += eps; xij[i] += eps; xij[j] += eps
            H[i, j] = H[j, i] = (f(xij) - f(xi) - f(xj) + fx) / eps ** 2
    return H


def garch_mle(r, gjr=False, student=False):
    r2, s0 = prepare(r); neg = (r < 0).astype(float)
    x0 = np.array([0.05 * s0, 0.03 if gjr else 0.08, 0.90]
                  + ([0.06] if gjr else []) + ([8.0] if student else []))
    def nll(th):
        ll = garch_loglik(th, r2, s0, neg, gjr)
        return 1e10 if not np.isfinite(ll) else -ll
    res = minimize(nll, x0, method="Nelder-Mead",
                   options=dict(xatol=1e-8, fatol=1e-8, maxiter=8000))
    mle = res.x
    cov = np.linalg.inv(_num_hess(nll, mle))
    return mle, np.sqrt(np.diag(cov)), cov


# ---------------------------------------------------------------- 2. Gauss-Hermite quadrature
def gh_quadrature(r, mle, cov, K=10, gjr=False):
    """Adaptive (Naylor-Smith) Gauss-Hermite over D = len(mle) dims: posterior mean/cov + log marginal lik."""
    r2, s0 = prepare(r); neg = (r < 0).astype(float); D = len(mle)
    x, w = np.polynomial.hermite.hermgauss(K)
    Gd = np.meshgrid(*([x] * D), indexing="ij")
    Z = np.stack([g.ravel() for g in Gd], 1)                                  # (K^D, D)
    Wg = np.stack([g.ravel() for g in np.meshgrid(*([w] * D), indexing="ij")], 1).prod(1)
    C = np.linalg.cholesky(cov)
    theta = mle + np.sqrt(2.0) * (Z @ C.T)
    logh = np.array([garch_loglik(t, r2, s0, neg, gjr) for t in theta])
    logwt = np.log(Wg) + np.sum(Z ** 2, 1) + logh
    m = np.max(logwt); wt = np.exp(logwt - m)
    logml = m + np.log(wt.sum()) + np.log(np.abs(np.linalg.det(np.sqrt(2.0) * C)))
    wt /= wt.sum()
    mean = wt @ theta; d = theta - mean
    return dict(mean=mean, cov=(d * wt[:, None]).T @ d, logml=logml)


# ---------------------------------------------------------------- 3. Importance sampling
def importance_sampling(r, mle, cov, N=30000, df=5, inflate=2.0, seed=0, gjr=False):
    r2, s0 = prepare(r); neg = (r < 0).astype(float)
    rng = np.random.default_rng(seed)
    prop = multivariate_t(loc=mle, shape=inflate * cov, df=df, allow_singular=True)
    theta = np.atleast_2d(prop.rvs(size=N, random_state=rng))
    logh = np.array([garch_loglik(t, r2, s0, neg, gjr) for t in theta])
    logw = logh - prop.logpdf(theta)
    w = np.exp(logw - logw.max()); w /= w.sum()
    return dict(mean=w @ theta, theta=theta, w=w, ess=1.0 / np.sum(w ** 2))


# ---------------------------------------------------------------- 4. Metropolis-Hastings
def metropolis_hastings(r, mle, cov, N=20000, burn=5000, scale=0.5, seed=1, gjr=False):
    r2, s0 = prepare(r); neg = (r < 0).astype(float); D = len(mle)
    rng = np.random.default_rng(seed)
    L = np.linalg.cholesky(scale ** 2 * cov)
    th = np.array(mle, float); lp = garch_loglik(th, r2, s0, neg, gjr)
    out = np.empty((N, D)); acc = 0
    for i in range(N + burn):
        cand = th + L @ rng.standard_normal(D)
        lpc = garch_loglik(cand, r2, s0, neg, gjr)
        if np.log(rng.random()) < lpc - lp:
            th, lp = cand, lpc; acc += 1
        if i >= burn:
            out[i - burn] = th
    return dict(theta=out, mean=out.mean(0), accept=acc / (N + burn))


# ---------------------------------------------------------------- 5. Griddy Gibbs
def _grid_range(name, j, th, mle, se, span, gjr):
    lo = mle[j] - span * se[j]; hi = mle[j] + span * se[j]
    a, b = th[1], th[2]; g = th[3] if gjr else 0.0
    if name == "omega":   lo = max(1e-10, lo)
    elif name == "alpha": lo = max(0.0, lo);   hi = min(1.0 - g / 2 - b - 1e-4, hi)
    elif name == "beta":  lo = max(0.0, lo);   hi = min(1.0 - a - g / 2 - 1e-4, hi)
    elif name == "gamma": lo = max(-a, lo);    hi = min(2.0 * (1.0 - a - b) - 1e-4, hi)
    elif name == "nu":    lo = max(2.05, lo);  hi = min(60.0, hi)
    return lo, hi


def griddy_gibbs(r, mle, se, N=5000, burn=1000, G=45, span=8.0, seed=2, gjr=False):
    r2, s0 = prepare(r); neg = (r < 0).astype(float); D = len(mle)
    nm = names(gjr, len(mle) > 3 + int(gjr))
    rng = np.random.default_rng(seed)
    th = np.array(mle, float); out = np.empty((N, D))
    for it in range(N + burn):
        for j in range(D):
            lo, hi = _grid_range(nm[j], j, th, mle, se, span, gjr)
            grid = np.linspace(lo, hi, G); logf = np.empty(G)
            for gg in range(G):
                th[j] = grid[gg]; logf[gg] = garch_loglik(th, r2, s0, neg, gjr)
            f = np.exp(logf - np.nanmax(logf)); f[~np.isfinite(f)] = 0.0
            cdf = np.concatenate([[0.0], np.cumsum((f[1:] + f[:-1]) / 2 * np.diff(grid))])
            th[j] = mle[j] if cdf[-1] <= 0 else float(np.interp(rng.random() * cdf[-1], cdf, grid))
        if it >= burn:
            out[it - burn] = th
    return dict(theta=out, mean=out.mean(0))