"""Identification of SVARs through heteroskedasticity (Rigobon 2003) — from scratch.

If the structural impact matrix A0^{-1} is CONSTANT but the structural shock variances DIFFER across
volatility regimes, A0^{-1} is identified with no sign/timing/elasticity restrictions: solve the
simultaneous diagonalisation (generalised eigenproblem) of the two regime covariances.
Self-contained (numpy only); alters nothing else.
"""
import numpy as np


def fit_var(Y, p):
    T, m = Y.shape
    X = np.hstack([Y[p - l:T - l] for l in range(1, p + 1)] + [np.ones((T - p, 1))])
    B = np.linalg.lstsq(X, Y[p:], rcond=None)[0]
    E = Y[p:] - X @ B
    return {"B": B, "Sigma": (E.T @ E) / (X.shape[0] - X.shape[1]), "resid": E, "p": p, "m": m}


def ma_irf(fit, H):
    m, p = fit["m"], fit["p"]
    Phi = [fit["B"][l * m:(l + 1) * m].T for l in range(p)]
    Th = [np.eye(m)]
    for h in range(1, H + 1):
        Th.append(sum(Phi[l - 1] @ Th[h - l] for l in range(1, p + 1) if h - l >= 0))
    return Th


def rigobon(E, regime):
    """Structural impact matrix from two volatility regimes (E residuals, regime=bool mask).
    Returns B0 = A0^{-1} (columns = structural shocks) and lam = regime-2/regime-1 variance ratios."""
    S1 = np.cov(E[regime].T); S2 = np.cov(E[~regime].T)
    L1 = np.linalg.cholesky(S1); L1i = np.linalg.inv(L1)
    lam, W = np.linalg.eigh(L1i @ S2 @ L1i.T)                    # simultaneous diagonalisation
    return L1 @ W, lam                                          # B0 columns diagonalise BOTH regimes


def label(B0, order_to=None):
    """Fix the sign/ordering indeterminacy: assign each structural shock to the variable it loads on
    most (permute), and sign so the diagonal is positive."""
    m = B0.shape[0]
    perm = []
    avail = list(range(m))
    for i in range(m):                                          # variable i -> its dominant shock column
        j = max(avail, key=lambda c: abs(B0[i, c]))
        perm.append(j); avail.remove(j)
    Bp = B0[:, perm]
    Bp = Bp * np.sign(np.diag(Bp))[None, :]
    return Bp


def struct_irf(fit, B0, H):
    return np.array([Th @ B0 for Th in ma_irf(fit, H)])         # (H+1, m, m) [h, response, shock]


def _iw(S, nu, rng):
    m = S.shape[0]; L = np.linalg.cholesky(np.linalg.inv(S))
    A = L @ rng.standard_normal((m, int(nu)))
    return np.linalg.inv(A @ A.T)


def _match(Bd, Bref):
    """Fix label switching: permute & sign Bd's columns to best match the reference Bref."""
    from itertools import permutations, product
    m = Bd.shape[1]; best, bv = Bd, np.inf
    for perm in permutations(range(m)):
        for sg in product((1, -1), repeat=m):
            C = Bd[:, list(perm)] * np.array(sg)
            v = np.sum((C - Bref) ** 2)
            if v < bv: bv, best = v, C
    return best


def posterior_irf(fit, regime, H, ndraw=1000, seed=0):
    """Posterior het-ID and recursive-Cholesky structural IRFs (Sigma drawn from inverse-Wishart per regime).
    Returns het (ndraw,H+1,m,m) and chol (ndraw,H+1,m,m), columns consistently labelled."""
    rng = np.random.default_rng(seed)
    E = fit["resid"]; E1, E2 = E[regime], E[~regime]
    S1, S2 = E1.T @ E1, E2.T @ E2; n1, n2 = len(E1), len(E2)
    Bref = label(rigobon(E, regime)[0])
    het, cho = [], []
    for _ in range(ndraw):
        Sig1, Sig2 = _iw(S1, n1, rng), _iw(S2, n2, rng)
        L1 = np.linalg.cholesky(Sig1); L1i = np.linalg.inv(L1)
        lam, W = np.linalg.eigh(L1i @ Sig2 @ L1i.T)
        het.append(struct_irf(fit, _match(L1 @ W, Bref), H))
        cho.append(struct_irf(fit, np.linalg.cholesky(_iw(S1 + S2, n1 + n2, rng)), H))
    return np.array(het), np.array(cho)