"""
GARCH-in-Mean (GARCH-M) from scratch — the risk-return tradeoff.

The conditional volatility enters the MEAN equation, so periods of higher risk carry a higher
expected return (Engle, Lilien & Robins 1987, "ARCH-M"):

    r_t = mu + lambda * m_t + sqrt(h_t) * z_t,          z_t ~ N(0,1)  or  standardized t_nu
    h_t = omega + (alpha + gamma * 1[eps_{t-1}<0]) eps_{t-1}^2 + beta h_{t-1}

where eps_t = r_t - mu - lambda * m_t and the in-mean regressor m_t is the conditional standard
deviation sqrt(h_t) ('sd', default), the variance h_t ('var'), or log h_t ('log').
lambda is the price of risk: lambda>0 is the risk-return tradeoff.

Parameter vector theta:  [mu, lambda, omega, alpha, beta] + ([gamma] if gjr) + ([nu] if student).
Because m_t depends on h_t which depends on past eps (which depend on the mean), the recursion is a
genuine forward loop -- it cannot be vectorised like a plain GARCH filter.
"""
import numpy as np
from scipy.optimize import minimize
from scipy.special import gammaln


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


def _unpack(theta, gjr):
    mu, lam, omega, alpha, beta = theta[0], theta[1], theta[2], theta[3], theta[4]
    i = 5
    gamma = theta[i] if gjr else 0.0; i += int(gjr)
    nu = theta[i] if len(theta) > i else np.inf
    return mu, lam, omega, alpha, beta, gamma, nu


def _m(h, inmean):
    if inmean == "sd":  return np.sqrt(h)
    if inmean == "var": return h
    return np.log(h)                                              # 'log'


def garchm_filter(theta, r, gjr=False, inmean="sd"):
    """Forward recursion; returns conditional variance h_t and residual eps_t = r_t - mean_t."""
    mu, lam, omega, alpha, beta, gamma, _ = _unpack(theta, gjr)
    r = np.asarray(r, float); T = len(r)
    h = np.empty(T); eps = np.empty(T)
    h[0] = np.var(r)
    for t in range(T):
        if h[t] <= 0: h[t] = 1e-10
        eps[t] = r[t] - mu - lam * _m(h[t], inmean)
        if t + 1 < T:
            lev = alpha + (gamma if eps[t] < 0 else 0.0) if gjr else alpha
            h[t + 1] = omega + lev * eps[t] ** 2 + beta * h[t]
    return h, eps


def _feasible(theta, gjr):
    _, _, omega, alpha, beta, gamma, nu = _unpack(theta, gjr)
    return (omega > 0) and (alpha >= 0) and (beta >= 0) and (alpha + gamma >= 0) \
        and (alpha + gamma / 2 + beta < 1.0) and (nu > 2.0)


def loglik(theta, r, gjr=False, inmean="sd"):
    theta = np.asarray(theta, float)
    if not _feasible(theta, gjr):
        return -np.inf
    h, eps = garchm_filter(theta, r, gjr, inmean)
    if np.any(h <= 0) or not np.all(np.isfinite(eps)):
        return -np.inf
    _, _, _, _, _, _, nu = _unpack(theta, gjr)
    if np.isfinite(nu):
        sc2 = h * (nu - 2.0) / nu                                 # scale^2 so Var(eps_t)=h_t
        return np.sum(gammaln((nu + 1) / 2) - gammaln(nu / 2) - 0.5 * np.log(nu * np.pi * sc2)
                      - (nu + 1) / 2 * np.log1p(eps ** 2 / (nu * sc2)))
    return -0.5 * np.sum(np.log(2 * np.pi * h) + eps ** 2 / h)


# ---------------------------------------------------------------- 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 mle(r, gjr=False, student=False, inmean="sd"):
    r = np.asarray(r, float); v = np.var(r)
    x0 = [np.mean(r), 0.0, 0.05 * v, 0.03 if gjr else 0.07, 0.90] \
        + ([0.06] if gjr else []) + ([8.0] if student else [])
    x0 = np.array(x0)
    def nll(th):
        ll = loglik(th, r, gjr, inmean)
        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=20000))
    se = np.full(len(res.x), np.nan)
    try:
        cov = np.linalg.inv(_num_hess(nll, res.x))
        se = np.sqrt(np.abs(np.diag(cov)))
    except Exception:
        pass
    return res.x, se


def loglik_restricted(r, gjr=False, student=False, inmean="sd"):
    """Max log-likelihood of the lambda=0 (plain-GARCH-mean) model, RE-OPTIMIZING all other parameters.
    This is the correct restricted fit for a likelihood-ratio test of the in-mean term."""
    r = np.asarray(r, float); v = np.var(r)
    x0 = np.array([np.mean(r), 0.05 * v, 0.03 if gjr else 0.07, 0.90]
                  + ([0.06] if gjr else []) + ([8.0] if student else []))
    def full(xr): return np.concatenate([[xr[0], 0.0], xr[1:]])          # insert lambda = 0
    def nll(xr):
        ll = loglik(full(xr), r, gjr, inmean)
        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=20000))
    return loglik(full(res.x), r, gjr, inmean)


# ---------------------------------------------------------------- Bayesian (random-walk Metropolis)
def metropolis(r, mle_theta, mle_se, gjr=False, student=False, inmean="sd",
               N=20000, burn=5000, scale=0.4, seed=1):
    rng = np.random.default_rng(seed)
    D = len(mle_theta)
    step = scale * np.where(np.isfinite(mle_se) & (mle_se > 0), mle_se, 0.01 * np.abs(mle_theta) + 1e-4)
    th = np.array(mle_theta, float); lp = loglik(th, r, gjr, inmean)
    out = np.empty((N, D)); acc = 0
    for i in range(N + burn):
        cand = th + step * rng.standard_normal(D)
        lpc = loglik(cand, r, gjr, inmean)
        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), names=names(gjr, student))