"""Factor-Augmented VAR (Bernanke-Boivin-Eliasz 2005) — from scratch.

Summarise a large informational panel by a few principal-component factors, put them in a small VAR
with the policy rate, identify a monetary shock recursively, and recover the impulse response of EVERY
underlying series through its factor loadings.  Self-contained (numpy only); alters nothing else.
"""
import numpy as np


def extract_factors(X, K):
    """K principal-component factors from a standardised panel X (T,n).
    Returns F (T,K) unit-variance factors, loadings L (n,K), explained-variance ratios (n,)."""
    Xc = X - X.mean(0)
    U, S, Vt = np.linalg.svd(Xc, full_matrices=False)
    F = U[:, :K] * S[:K]
    F = F / F.std(0)                                              # unit-variance factors
    L = np.linalg.lstsq(F, Xc, rcond=None)[0].T                  # loadings from regressing X on F
    evr = S ** 2 / (S ** 2).sum()
    return F, L, evr


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


def _ma(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):
        acc = sum(Phi[l - 1] @ Th[h - l] for l in range(1, p + 1) if h - l >= 0)
        Th.append(acc if np.ndim(acc) else np.zeros((m, m)))
    return Th


def chol_irf(fit, H, shock):
    """Recursive structural IRF to `shock` (index), P = chol(Sigma).  Returns (H+1, m)."""
    P = np.linalg.cholesky(fit["Sigma"])
    return np.array([(Th @ P)[:, shock] for Th in _ma(fit, H)])


def favar(panel, R, K, p=2, H=20, slow_mask=None):
    """Estimate a FAVAR: factors from `panel` (T,n std.), state = [factors, policy R], monetary shock
    ordered LAST (recursive). Returns factor/policy IRFs and the reconstructed IRF of every panel series.

    panel     : (T,n) standardised informational variables
    R         : (T,) policy variable (e.g. the funds rate)
    slow_mask : optional (n,) bool of SLOW-moving variables. If given, apply the Bernanke-Boivin-Eliasz
                'slow-fast' cleaning so the factors carry no contemporaneous response to R (removes the price puzzle).
    Returns dict: 'evr', 'irf_state' (H+1, K+1), 'irf_panel' (H+1, n), 'K'
    """
    F, L, evr = extract_factors(panel, K)
    if slow_mask is not None:
        Fs, _, _ = extract_factors(panel[:, slow_mask], K)      # factors of the slow block
        Z = np.column_stack([Fs, R])
        b = np.linalg.lstsq(Z, F, rcond=None)[0]                # regress each factor on [slow factors, R]
        F = F - np.outer(R, b[-1])                              # strip the contemporaneous R component
    Y = np.column_stack([F, R])                                  # [F1..FK, R], policy last
    fit = _fit_var(Y, p)
    irf_state = chol_irf(fit, H, shock=K)                        # response of [F,R] to the policy (last) shock
    # reconstruct each panel series: regress x_i on [F, R] -> loadings, then x-IRF = loadings @ state-IRF
    Z = np.column_stack([F, R])
    Lam = np.linalg.lstsq(Z, panel, rcond=None)[0]              # (K+1, n)
    irf_panel = irf_state @ Lam                                 # (H+1, n)
    return {"evr": evr, "irf_state": irf_state, "irf_panel": irf_panel, "K": K,
            "loadings": L, "Y": Y, "Lam": Lam, "fit": fit}


def boot_irf(Y, Lam, p, H, nboot=500, seed=0):
    """Residual bootstrap of the state VAR [F,R] (policy last); returns state & panel IRF draws."""
    rng = np.random.default_rng(seed)
    T, m = Y.shape
    Xd = np.hstack([Y[p - l:T - l] for l in range(1, p + 1)] + [np.ones((T - p, 1))])
    B = np.linalg.lstsq(Xd, Y[p:], rcond=None)[0]; E = Y[p:] - Xd @ B
    S, Pn = [], []
    for _ in range(nboot):
        Yb = Y[:p].copy().tolist()
        idx = rng.integers(0, len(E), T - p)
        for t in range(p, T):
            x = np.hstack([Yb[t - l] for l in range(1, p + 1)] + [1.0])
            Yb.append(x @ B + E[idx[t - p]])
        fb = _fit_var(np.array(Yb), p)
        ir = chol_irf(fb, H, shock=m - 1); ir = ir * np.sign(ir[0, -1])
        S.append(ir); Pn.append(ir @ Lam)
    return np.array(S), np.array(Pn)