"""
Two more SVAR identification schemes on top of `bvar.py`, complementing the sign restrictions in svar_sign.py:

  proxy_svar()     -- external-instrument / proxy SVAR (Stock-Watson 2012/18; Mertens-Ravn 2013).
                      An observed instrument z_t, correlated with the target structural shock but not the
                      others, pins the shock's impact vector: b proportional to E[u_t z_t] (regress the
                      reduced-form residuals on z), normalized so the policy variable moves by `norm`.
  blanchard_quah() -- long-run restrictions (Blanchard-Quah 1989). The impact matrix is chosen so a shock
                      has NO permanent effect on the first variable's level: with the long-run multiplier
                      Xi = (I - sum Pi_l)^{-1}, the long-run structural matrix Xi P is lower-triangular, i.e.
                      Xi P = chol(Xi Sigma Xi'), so P = Xi^{-1} chol(Xi Sigma Xi').

Both consume a `bvar.py` fit (Minnesota) and reuse the MA-IRF machinery from svar_sign. bvar.py is unmodified.
"""
import numpy as np
import bvar
from svar_sign import _ma_irf, _pi_blocks


def proxy_svar(fit, Y, z, policy_idx, H=48, norm=1.0):
    """External-instrument SVAR. z is the instrument over the full sample (length T); it is aligned to the
    VAR residual rows p..T-1 internally. Returns dict(irf=(ndraw, H+1, m)); the shock raises variable
    `policy_idx` by `norm` on impact (a contractionary normalization if norm>0)."""
    Y = np.asarray(Y, float); m = Y.shape[1]; p = fit["p"]
    Yt, X = bvar.build_regressors(Y, p, intercept=True)          # minnesota design (with intercept)
    z = np.asarray(z, float)[p:]                                 # align instrument to residual sample
    z = z - z.mean()
    coefs = fit["B"]
    out = []
    for B in coefs:
        U = Yt - X @ B                                           # reduced-form residuals (T-p, m)
        b = (U.T @ z) / (z @ z)                                  # proxy impact vector (m,), up to scale
        b = b / b[policy_idx] * norm                             # normalize: policy variable moves `norm`
        Theta = _ma_irf(_pi_blocks(B, "minnesota", m, p), m, p, H)
        out.append(Theta @ b)                                   # structural IRF (H+1, m)
    return {"irf": np.array(out)}


def blanchard_quah(fit, H=40):
    """Long-run-restricted SVAR (2 variables). Shock 0 has a PERMANENT effect on variable 0's level,
    shock 1 has only a TRANSITORY one (no long-run effect on variable 0). Returns dict(irf=(ndraw, H+1, m, m))
    with irf[..., i, k] = response of variable i to structural shock k."""
    coefs = fit["B"] if fit["kind"] == "minnesota" else fit["Pi"]
    Sig = fit["Sigma"]; m = Sig.shape[1]; p = fit["p"]
    out = []
    for d in range(len(coefs)):
        Pis = _pi_blocks(coefs[d], fit["kind"], m, p)
        Xi = np.linalg.inv(np.eye(m) - sum(Pis))                # long-run multiplier (I - sum Pi)^{-1}
        ChL = np.linalg.cholesky(Xi @ Sig[d] @ Xi.T)            # lower-triangular long-run structural matrix
        P = np.linalg.solve(Xi, ChL)                            # structural impact matrix = Xi^{-1} ChL
        Theta = _ma_irf(Pis, m, p, H)
        out.append(np.einsum("hij,jk->hik", Theta, P))         # (H+1, m, m)
    return {"irf": np.array(out)}


def bands(res, qs=(16, 50, 84)):
    return np.percentile(res["irf"], qs, axis=0)