"""
Markov-switching GARCH (Haas, Mittnik & Paolella 2004 specification) from scratch, Normal or Student-t.

A latent 2-state Markov chain S_t (calm / turbulent) switches the volatility dynamics. In the HMP
specification each regime carries its OWN GARCH process, run on the full return series, which removes
the path-dependence that makes exact MS-GARCH likelihoods intractable:

    h_t^{(k)} = omega_k + alpha_k r_{t-1}^2 + beta_k h_{t-1}^{(k)},     k = 1,2
    r_t | S_t = k  ~  N(0, h_t^{(k)})  or  standardized t_nu with variance h_t^{(k)}
    S_t Markov with transition P[i,j] = P(S_t=j | S_{t-1}=i)

Likelihood via the Hamilton (1989) filter; smoothed regime probabilities via the Kim (1994) smoother.
theta = (o1,a1,b1, o2,a2,b2, p11, p22)  [+ nu if Student-t, one shared tail index across regimes].
"""
import numpy as np
from scipy.optimize import minimize
from scipy.stats import t as _t


def _paths(r2, th):
    o1, a1, b1, o2, a2, b2 = th[:6]
    T = len(r2); H = np.empty((T, 2))
    H[0, 0] = o1 / max(1 - a1 - b1, 1e-4); H[0, 1] = o2 / max(1 - a2 - b2, 1e-4)
    for t in range(1, T):
        H[t, 0] = o1 + a1 * r2[t - 1] + b1 * H[t - 1, 0]
        H[t, 1] = o2 + a2 * r2[t - 1] + b2 * H[t - 1, 1]
    return H


def _feasible(th, student):
    o1, a1, b1, o2, a2, b2, p11, p22 = th[:8]
    ok = (o1 > 0 and o2 > 0 and a1 >= 0 and b1 >= 0 and a2 >= 0 and b2 >= 0
          and a1 + b1 < 1 and a2 + b2 < 1 and 0 < p11 < 1 and 0 < p22 < 1)
    if student:
        ok = ok and th[8] > 2.05
    return ok


def _emissions(r, H, th, student):
    """Per-regime observation densities eta_t(k) = f(r_t | S_t=k), vectorised over t (shape T x 2)."""
    if student:
        nu = th[8]; sc = np.sqrt(H * (nu - 2.0) / nu)          # scale so Var(r_t|k)=h_t^{(k)}
        return _t.pdf(r[:, None], df=nu, scale=sc)
    return np.exp(-0.5 * (np.log(2 * np.pi * H) + r[:, None] ** 2 / H))


def filter(th, r, student=False, return_all=False):
    """Hamilton filter: log-likelihood (+ filtered probs, H paths, predicted probs if return_all)."""
    r2 = r ** 2; T = len(r); H = _paths(r2, th)
    p11, p22 = th[6], th[7]
    P = np.array([[p11, 1 - p11], [1 - p22, p22]])
    pi = np.array([1 - p22, 1 - p11]); pi = pi / pi.sum()        # stationary distribution
    eta = _emissions(r, H, th, student)
    xi = pi.copy(); ll = 0.0; filt = np.empty((T, 2)); pred = np.empty((T, 2))
    for t in range(T):
        pr = xi @ P                                             # P(S_t | data<t)
        num = pr * eta[t]; lik = num.sum()
        if lik <= 0 or not np.isfinite(lik):
            return (-np.inf, None, None, None) if return_all else -np.inf
        ll += np.log(lik); xi = num / lik
        filt[t] = xi; pred[t] = pr
    return (ll, filt, H, pred) if return_all else ll


def smooth(th, r, student=False):
    """Kim smoother: P(S_t = k | all data)."""
    _, filt, H, pred = filter(th, r, student, return_all=True)
    T = len(r); p11, p22 = th[6], th[7]
    P = np.array([[p11, 1 - p11], [1 - p22, p22]])
    sm = np.empty((T, 2)); sm[-1] = filt[-1]
    for t in range(T - 2, -1, -1):
        for j in range(2):
            sm[t, j] = filt[t, j] * np.sum(P[j, :] * sm[t + 1, :] / np.where(pred[t + 1] > 0, pred[t + 1], 1e-12))
        sm[t] /= sm[t].sum()
    return sm, H


def loglik(th, r, student=False):
    if not _feasible(th, student):
        return -np.inf
    return filter(th, r, student)


def mle(r, student=False):
    r = np.asarray(r, float); r = r - r.mean(); v = np.var(r)
    x0 = [0.02 * v, 0.05, 0.90, 0.30 * v, 0.15, 0.75, 0.98, 0.95] + ([8.0] if student else [])
    nll = lambda th: -loglik(th, r, student) if np.isfinite(loglik(th, r, student)) else 1e12
    res = minimize(nll, np.array(x0), method="Nelder-Mead", options=dict(xatol=1e-7, fatol=1e-7, maxiter=50000))
    th = res.x
    u1 = th[0] / (1 - th[1] - th[2]); u2 = th[3] / (1 - th[4] - th[5])
    if u1 > u2:                                                  # order so regime 2 is the HIGH-vol state
        tail = list(th[8:])
        th = np.array([th[3], th[4], th[5], th[0], th[1], th[2], th[7], th[6]] + tail)
    return th, r


def names(student=False):
    return ["omega1", "alpha1", "beta1", "omega2", "alpha2", "beta2", "p11", "p22"] + (["nu"] if student else [])