"""
Constant Conditional Correlation GARCH (Bollerslev 1990) from scratch.

The CCC decomposition writes the conditional covariance matrix as
    H_t = D_t R D_t,     D_t = diag(sqrt(h_{1t}), ..., sqrt(h_{Nt})),   R = constant correlation,
so each asset gets its OWN univariate GARCH for the conditional variance h_{it}, and a SINGLE constant
correlation matrix R ties them together. Estimated in two steps (Bollerslev 1990; Engle 2002 stage-1):
    1. fit a univariate GARCH(1,1) to each series -> conditional variances h_{it}, standardized resids z_{it}
    2. R = sample correlation of the standardized residuals z.
The constancy of R is exactly the assumption the DCC model relaxes; rolling_corr() exposes it.

Self-contained (own univariate GARCH), independent of the univariate-arc scripts.
"""
import numpy as np
from scipy.optimize import minimize


def garch11(r):
    """Univariate GARCH(1,1) with constant mean, Gaussian QMLE. Returns dict(mu,omega,alpha,beta,h,z)."""
    r = np.asarray(r, float); mu = r.mean(); e = r - mu; v = e.var()

    def filt(o, a, b):
        h = np.empty(len(e)); h[0] = v
        for t in range(1, len(e)):
            h[t] = o + a * e[t - 1] ** 2 + b * h[t - 1]
        return h

    def nll(th):
        o, a, b = th
        if o <= 0 or a < 0 or b < 0 or a + b >= 0.9999:
            return 1e12
        h = filt(o, a, b)
        return 0.5 * np.sum(np.log(2 * np.pi * h) + e ** 2 / h)

    best = None
    for x0 in ([0.05 * v, 0.05, 0.90], [0.10 * v, 0.10, 0.85]):
        res = minimize(nll, x0, method="Nelder-Mead",
                       options=dict(xatol=1e-9, fatol=1e-9, maxiter=20000))
        if best is None or res.fun < best.fun:
            best = res
    o, a, b = best.x; h = filt(o, a, b)
    return dict(mu=mu, omega=o, alpha=a, beta=b, persistence=a + b, h=h, z=e / np.sqrt(h))


def ccc_fit(R):
    """Two-step CCC. R is a (T x N) return matrix. Returns dict(fits, H, Z, Rbar, loglik)."""
    R = np.asarray(R, float); T, N = R.shape
    fits = [garch11(R[:, i]) for i in range(N)]
    H = np.column_stack([f["h"] for f in fits])         # conditional variances h_{it}
    Z = np.column_stack([f["z"] for f in fits])         # standardized residuals
    Rbar = np.corrcoef(Z, rowvar=False)                 # the constant conditional correlation matrix
    ll = ccc_loglik(R, H, Rbar, mu=np.array([f["mu"] for f in fits]))
    return dict(fits=fits, H=H, Z=Z, Rbar=Rbar, loglik=ll)


def ccc_loglik(R, H, Rbar, mu=None):
    """Gaussian log-likelihood of the CCC model, H_t = D_t Rbar D_t."""
    R = np.asarray(R, float); T, N = R.shape
    if mu is not None:
        R = R - mu
    Ri = np.linalg.inv(Rbar); ldR = np.linalg.slogdet(Rbar)[1]
    D = np.sqrt(H); z = R / D                            # z_t = D_t^{-1} (r_t - mu)
    quad = np.einsum("ti,ij,tj->t", z, Ri, z)           # z_t' Rbar^{-1} z_t
    logdetH = 2.0 * np.sum(np.log(D), axis=1) + ldR     # log|D_t Rbar D_t|
    return float(-0.5 * np.sum(N * np.log(2 * np.pi) + logdetH + quad))


def rolling_corr(x, y, window=126):
    """Rolling Pearson correlation of two series (window trading days). NaN for the first window-1."""
    x = np.asarray(x, float); y = np.asarray(y, float); T = len(x)
    out = np.full(T, np.nan)
    for t in range(window - 1, T):
        xs = x[t - window + 1:t + 1]; ys = y[t - window + 1:t + 1]
        out[t] = np.corrcoef(xs, ys)[0, 1]
    return out