"""
Dynamic Conditional Correlation GARCH (Engle 2002) from scratch.

DCC relaxes CCC's constant correlation into a time-varying R_t driven by its own GARCH-like recursion.
Same volatility/correlation split H_t = D_t R_t D_t, but now, on the standardized residuals z_t = D_t^{-1} eps_t,

    Q_t = (1 - a - b) Qbar + a z_{t-1} z_{t-1}' + b Q_{t-1},        Qbar = uncond. corr of z,
    R_t = diag(Q_t)^{-1/2} Q_t diag(Q_t)^{-1/2}      (rescale Q_t to a proper correlation matrix).

Estimated in two steps (Engle 2002): stage 1 = a univariate GARCH per series (reused from ccc.garch11);
stage 2 = maximize the correlation part of the Gaussian log-likelihood over the two scalars (a, b),
    L_c(a,b) = -1/2 sum_t [ log|R_t| + z_t' R_t^{-1} z_t ].
CCC is the special case a = b = 0 (R_t == Qbar), so DCC strictly nests it.
"""
import numpy as np
from scipy.optimize import minimize
from ccc import garch11


def _Rt_path(Z, Qbar, a, b):
    """The full sequence of correlation matrices R_t (T x N x N)."""
    T, N = Z.shape; Q = Qbar.copy(); Rt = np.empty((T, N, N))
    for t in range(T):
        if t > 0:
            Q = (1 - a - b) * Qbar + a * np.outer(Z[t - 1], Z[t - 1]) + b * Q
        d = 1.0 / np.sqrt(np.diag(Q))
        Rt[t] = Q * np.outer(d, d)
    return Rt


def _corr_loglik(Z, Qbar, a, b):
    """Stage-2 correlation log-likelihood (the part that depends on a, b)."""
    T, N = Z.shape; Q = Qbar.copy(); ll = 0.0
    for t in range(T):
        if t > 0:
            Q = (1 - a - b) * Qbar + a * np.outer(Z[t - 1], Z[t - 1]) + b * Q
        d = 1.0 / np.sqrt(np.diag(Q)); Rt = Q * np.outer(d, d)
        _, ld = np.linalg.slogdet(Rt)
        ll += -0.5 * (ld + Z[t] @ np.linalg.solve(Rt, Z[t]))
    return ll


def _full_loglik(Rret, H, Rt, mu):
    """Full Gaussian log-likelihood with H_t = D_t R_t D_t (for AIC vs CCC)."""
    T, N = Rret.shape; e = Rret - np.asarray(mu); D = np.sqrt(H); Z = e / D; ll = 0.0
    for t in range(T):
        _, ld = np.linalg.slogdet(Rt[t])
        logdetH = 2.0 * np.sum(np.log(D[t])) + ld
        ll += -0.5 * (N * np.log(2 * np.pi) + logdetH + Z[t] @ np.linalg.solve(Rt[t], Z[t]))
    return float(ll)


def dcc_fit(Rret):
    """Two-step DCC. Rret is (T x N). Returns dict(fits, H, Z, Qbar, a, b, Rt, loglik)."""
    Rret = np.asarray(Rret, float); T, N = Rret.shape
    fits = [garch11(Rret[:, i]) for i in range(N)]           # stage 1: univariate GARCH per series
    H = np.column_stack([f["h"] for f in fits])
    Z = np.column_stack([f["z"] for f in fits])
    Qbar = (Z.T @ Z) / T                                     # unconditional corr of standardized resids

    def neg(x):
        a, b = x
        if a < 0 or b < 0 or a + b > 0.9999:
            return 1e12
        return -_corr_loglik(Z, Qbar, a, b)

    best = None
    for x0 in ([0.02, 0.97], [0.05, 0.90]):                  # stage 2: optimize (a, b)
        res = minimize(neg, x0, method="Nelder-Mead",
                       options=dict(xatol=1e-7, fatol=1e-7, maxiter=4000))
        if best is None or res.fun < best.fun:
            best = res
    a, b = best.x
    Rt = _Rt_path(Z, Qbar, a, b)
    ll = _full_loglik(Rret, H, Rt, mu=np.array([f["mu"] for f in fits]))
    return dict(fits=fits, H=H, Z=Z, Qbar=Qbar, a=a, b=b, Rt=Rt, loglik=ll)


def pair_corr(Rt, i, j):
    """Extract the conditional-correlation time series between assets i and j."""
    return Rt[:, i, j]