"""
Diagonal BEKK-GARCH (Engle & Kroner 1995) from scratch, bivariate.

BEKK models the conditional covariance matrix H_t DIRECTLY (not via a volatility/correlation split), in a
quadratic form that keeps H_t positive-definite for free at every t:

    H_t = C C' + A' eps_{t-1} eps_{t-1}' A + B' H_{t-1} B,

with C lower-triangular (so C C' is p.d.). In the DIAGONAL BEKK, A = diag(a1,..,aN) and B = diag(b1,..,bN),
so element-by-element
    H_t[i,j] = (C C')[i,j] + a_i a_j eps_{t-1,i} eps_{t-1,j} + b_i b_j H_{t-1}[i,j].
Every H_t is a sum of a p.d. matrix and outer-product / congruence terms, hence p.d. by construction -- BEKK's
signature guarantee. The cost is parameters: full BEKK needs 2N^2 + N(N+1)/2; diagonal trims to 2N + N(N+1)/2;
scalar (A=aI, B=bI) to 2 + N(N+1)/2. This "N^2 curse" is why DCC, not BEKK, scales to many assets.

theta = (c11, c21, c22, a1, a2, b1, b2)  for N = 2.
"""
import numpy as np
from scipy.optimize import minimize


def _unpack(theta):
    C = np.array([[theta[0], 0.0], [theta[1], theta[2]]])
    a = np.array([theta[3], theta[4]]); b = np.array([theta[5], theta[6]])
    return C, a, b


def _filter(theta, R):
    C, a, b = _unpack(theta)
    CC = C @ C.T; aa = np.outer(a, a); bb = np.outer(b, b)
    T = len(R); H = np.empty((T, 2, 2)); H[0] = np.cov(R, rowvar=False)
    for t in range(1, T):
        e = R[t - 1]
        H[t] = CC + aa * np.outer(e, e) + bb * H[t - 1]
    return H


def _feasible(theta):
    C, a, b = _unpack(theta)
    return (theta[0] > 0 and theta[2] > 0 and                      # positive diagonal of C (identification)
            a[0] ** 2 + b[0] ** 2 < 1 and a[1] ** 2 + b[1] ** 2 < 1)  # covariance stationarity (diagonal case)


def loglik(theta, R):
    if not _feasible(theta):
        return -np.inf
    H = _filter(theta, R)
    h11, h12, h22 = H[:, 0, 0], H[:, 0, 1], H[:, 1, 1]
    det = h11 * h22 - h12 ** 2
    if np.any(det <= 0):
        return -np.inf
    quad = (h22 * R[:, 0] ** 2 - 2 * h12 * R[:, 0] * R[:, 1] + h11 * R[:, 1] ** 2) / det
    return float(-0.5 * np.sum(2 * np.log(2 * np.pi) + np.log(det) + quad))


def bekk_fit(R):
    """Fit diagonal BEKK by Gaussian MLE. Returns dict(theta, C, a, b, H, corr, loglik)."""
    R = np.asarray(R, float); R = R - R.mean(0)
    S = np.cov(R, rowvar=False)
    C0 = np.linalg.cholesky(0.05 * S)
    x0 = np.array([C0[0, 0], C0[1, 0], C0[1, 1], 0.20, 0.20, 0.96, 0.96])
    nll = lambda th: -loglik(th, R)
    best = None
    for jit in (0.0, 0.03):
        xx = x0.copy(); xx[3:5] += jit
        res = minimize(nll, xx, method="Nelder-Mead",
                       options=dict(xatol=1e-7, fatol=1e-7, maxiter=40000, maxfev=40000))
        if best is None or res.fun < best.fun:
            best = res
    th = best.x; C, a, b = _unpack(th); H = _filter(th, R)
    corr = H[:, 0, 1] / np.sqrt(H[:, 0, 0] * H[:, 1, 1])
    return dict(theta=th, C=C, a=a, b=b, H=H, corr=corr, loglik=-best.fun)