"""
Heston-Nandi (2000) closed-form GARCH option valuation, from scratch.

Physical dynamics (log return R_t = ln S_t/S_{t-1}, per period):
    R_t = r + lambda * h_t + sqrt(h_t) * z_t,   z_t ~ N(0,1)
    h_t = omega + beta*h_{t-1} + alpha*(z_{t-1} - gamma*sqrt(h_{t-1}))^2
Variance persistence  beta + alpha*gamma^2 (< 1);  unconditional var = (omega+alpha)/(1 - beta - alpha*gamma^2).

Risk-neutral (LRNVR): lambda* = -1/2,  gamma* = gamma + lambda + 1/2, same (omega,alpha,beta).
The generating function E*_t[S_T^phi] = S_t^phi * exp(A_t + B_t h_{t+1}) has a closed recursion, so European
options price by Fourier inversion (Heston-Nandi 2000, eqs 5-6, 14).
"""
import numpy as np
from scipy.optimize import minimize
from scipy.stats import norm
from scipy.optimize import brentq

_trapz = getattr(np, "trapezoid", getattr(np, "trapz", None))    # numpy 2.x renamed trapz -> trapezoid


# ---------------------------------------------------------------- estimation (physical measure)
def hn_filter(theta, R, r):
    omega, alpha, beta, gamma, lam = theta
    T = len(R); h = np.empty(T); z = np.empty(T)
    h[0] = (omega + alpha) / max(1 - beta - alpha * gamma**2, 1e-6)        # unconditional start
    for t in range(T):
        if h[t] <= 0: h[t] = 1e-12
        z[t] = (R[t] - r - lam * h[t]) / np.sqrt(h[t])
        if t + 1 < T:
            h[t + 1] = omega + beta * h[t] + alpha * (z[t] - gamma * np.sqrt(h[t]))**2
    return h, z


def hn_loglik(theta, R, r):
    omega, alpha, beta, gamma, lam = theta
    if omega <= 0 or alpha < 0 or beta < 0 or beta + alpha * gamma**2 >= 1:
        return -np.inf
    h, z = hn_filter(theta, R, r)
    if np.any(h <= 0) or not np.all(np.isfinite(z)):
        return -np.inf
    return -0.5 * np.sum(np.log(2 * np.pi * h) + z**2)


def hn_mle(R, r=0.0):
    v = np.var(R)
    x0 = np.array([v * 0.1, v * 0.2, 0.6, 2.0, 0.5])                       # omega, alpha, beta, gamma, lambda
    nll = lambda th: -hn_loglik(th, R, r) if np.isfinite(hn_loglik(th, R, r)) else 1e10
    res = minimize(nll, x0, method="Nelder-Mead", options=dict(xatol=1e-10, fatol=1e-6, maxiter=20000))
    return res.x


# ---------------------------------------------------------------- option pricing (risk-neutral)
def hn_gf(cphi, S, h_next, r, omega, alpha, beta, gamma, lam, T):
    """Risk-neutral generating function E*[S_T^phi] for complex phi, T periods ahead."""
    gs = gamma + lam + 0.5                                                 # risk-neutral gamma*
    ls = -0.5
    A = np.zeros_like(cphi); B = np.zeros_like(cphi)
    for _ in range(int(T)):
        Bn = cphi * (ls + gs) - 0.5 * gs**2 + beta * B + 0.5 * (cphi - gs)**2 / (1 - 2 * alpha * B)
        A = A + cphi * r + B * omega - 0.5 * np.log(1 - 2 * alpha * B)
        B = Bn
    return np.exp(cphi * np.log(S) + A + B * h_next)


def hn_call(S, K, r, T, theta, h_next, phimax=120.0, nphi=400):
    """European call price(s); K may be scalar or array. h_next = h_{t+1} (next-period variance)."""
    omega, alpha, beta, gamma, lam = theta
    phi = np.linspace(1e-8, phimax, nphi)
    f1 = hn_gf(1j * phi + 1, S, h_next, r, omega, alpha, beta, gamma, lam, T)
    f0 = hn_gf(1j * phi, S, h_next, r, omega, alpha, beta, gamma, lam, T)
    disc = np.exp(-r * T)
    K = np.atleast_1d(np.asarray(K, float)); out = np.empty(len(K))
    for i, k in enumerate(K):
        i1 = (k ** (-1j * phi) * f1 / (1j * phi)).real
        i2 = (k ** (-1j * phi) * f0 / (1j * phi)).real
        P1 = 0.5 * S + disc / np.pi * _trapz(i1, phi)
        P2 = 0.5 + 1.0 / np.pi * _trapz(i2, phi)
        out[i] = P1 - k * disc * P2
    return out if out.size > 1 else float(out[0])


# ---------------------------------------------------------------- Black-Scholes + implied vol
def bs_call(S, K, r, T, sig):
    if sig <= 0 or T <= 0: return max(S - K * np.exp(-r * T), 0.0)
    d1 = (np.log(S / K) + (r + 0.5 * sig**2) * T) / (sig * np.sqrt(T))
    d2 = d1 - sig * np.sqrt(T)
    return S * norm.cdf(d1) - K * np.exp(-r * T) * norm.cdf(d2)


def bs_iv(price, S, K, r, T):
    intrinsic = max(S - K * np.exp(-r * T), 0.0)
    if price <= intrinsic + 1e-10: return np.nan
    try:
        return brentq(lambda s: bs_call(S, K, r, T, s) - price, 1e-6, 5.0, maxiter=200)
    except Exception:
        return np.nan


# ---------------------------------------------------------------- Duan (1995) Monte-Carlo (risk-neutral simulation)
def hn_mc_call(S, K, r, T, theta, h_next, M=200000, seed=0, antithetic=True):
    """Monte-Carlo European call under the risk-neutral GARCH (Duan 1995 LRNVR); returns (price, std-err)."""
    omega, alpha, beta, gamma, lam = theta
    gs = gamma + lam + 0.5                                        # risk-neutral gamma*
    rng = np.random.default_rng(seed)
    if antithetic:
        Zh = rng.standard_normal((int(T), M // 2)); Z = np.concatenate([Zh, -Zh], axis=1)
    else:
        Z = rng.standard_normal((int(T), M))
    Mn = Z.shape[1]
    h = np.full(Mn, h_next); logS = np.full(Mn, np.log(S))
    for t in range(int(T)):
        z = Z[t]
        logS += r - 0.5 * h + np.sqrt(h) * z                     # R_t = r - h/2 + sqrt(h) z  (risk-neutral)
        h = omega + beta * h + alpha * (z - gs * np.sqrt(h))**2   # next-period variance
    ST = np.exp(logS); disc = np.exp(-r * T)
    K = np.atleast_1d(np.asarray(K, float))
    pay = np.maximum(ST[None, :] - K[:, None], 0.0)
    price = disc * pay.mean(1); se = disc * pay.std(1) / np.sqrt(Mn)
    return (price, se) if K.size > 1 else (float(price[0]), float(se[0]))