"""Structural VAR toolkit for the global oil market (Kilian 2009; Kilian-Murphy 2012/14) — from scratch.

Self-contained (numpy only); does NOT import or modify any earlier module.  Provides:
  fit_var            - OLS reduced-form VAR
  ma_irf             - reduced-form moving-average matrices Theta_h
  chol_irf           - recursive (Cholesky) structural IRFs, shock j
  fevd               - forecast-error variance decomposition
  historical_decomp  - contribution of each structural shock to each series over time
  sign_elast_irf     - set-identification by sign restrictions + short-run elasticity bounds (rotation search)
"""
import numpy as np


def _design(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))])
    return Y[p:], X                                             # Yt (T-p,m), X (T-p, mp+1) intercept LAST


def fit_var(Y, p):
    """OLS reduced-form VAR(p) with intercept.  Returns coefficients, Sigma, residuals."""
    Yt, X = _design(Y, p)
    B = np.linalg.lstsq(X, Yt, rcond=None)[0]                  # (mp+1, m)
    E = Yt - X @ B
    Sig = (E.T @ E) / (X.shape[0] - X.shape[1])
    return {"B": B, "Sigma": Sig, "resid": E, "p": p, "m": Y.shape[1], "Yt": Yt, "X": X}


def ma_irf(fit, H):
    """Reduced-form MA matrices Theta_0=I, Theta_h  (list length H+1), from the companion recursion."""
    m, p = fit["m"], fit["p"]
    Phi = [fit["B"][l * m:(l + 1) * m].T for l in range(p)]    # lag matrices (m,m)
    Th = [np.eye(m)]
    for h in range(1, H + 1):
        acc = np.zeros((m, m))
        for l in range(1, p + 1):
            if h - l >= 0:
                acc += Phi[l - 1] @ Th[h - l]
        Th.append(acc)
    return Th


def chol_irf(fit, H, P=None):
    """Structural IRFs under a recursive (lower-triangular) impact matrix P (default chol(Sigma)).
    Returns array (H+1, m, m): [horizon, response variable, shock]."""
    if P is None:
        P = np.linalg.cholesky(fit["Sigma"])
    Th = ma_irf(fit, H)
    return np.array([Th[h] @ P for h in range(H + 1)])


def fevd(fit, H, P=None):
    """Forecast-error variance decomposition. Returns (H+1, m, m): share of var of variable i from shock j."""
    irf = chol_irf(fit, H, P)                                  # (H+1,m,m)
    m = fit["m"]
    contrib = np.cumsum(irf ** 2, axis=0)                      # cumulative MSE contribution by shock
    total = contrib.sum(axis=2, keepdims=True)
    return contrib / total


def historical_decomp(fit, P=None):
    """Historical decomposition: contribution of each structural shock to each series over the sample.
    Returns (Teff, m, m): [t, variable, shock]  (deviations from the deterministic/base projection)."""
    if P is None:
        P = np.linalg.cholesky(fit["Sigma"])
    Th = ma_irf(fit, fit["resid"].shape[0] - 1)               # MA up to sample length
    u = fit["resid"] @ np.linalg.inv(P).T                     # structural shocks (Teff, m)
    Teff, m = u.shape
    hd = np.zeros((Teff, m, m))
    for t in range(Teff):
        acc = np.zeros((m, m))
        for h in range(t + 1):
            acc += (Th[h] @ P) * u[t - h][None, :]           # response(var,shock) weighted by shock at t-h
        hd[t] = acc
    return hd


def sign_rotations(fit, H=18, ntry=50000, seed=0, e_supply_max=None):
    """Set-identify the oil-market SVAR by SIGN RESTRICTIONS (Kilian-Murphy).
    Variables assumed [oil production growth, real activity, real oil price].
    Impact sign pattern (each shock normalised to raise the real oil price):
        supply disruption : production < 0, activity < 0
        aggregate demand  : production > 0, activity > 0
        oil-specific demand: production > 0, activity < 0
    For each accepted rotation the implied short-run price elasticity of oil SUPPLY
    (production response / price response along the demand-shift) is returned, so the caller can
    impose a hard bound (Kilian-Murphy) or weight by an elasticity prior (Baumeister-Hamilton).

    Returns dict: 'irf' (n, H+1, 3, 3) ordered [supply, aggdem, oilspec], 'elast_supply' (n,), 'accept_frac'.
    """
    rng = np.random.default_rng(seed)
    C = np.linalg.cholesky(fit["Sigma"]); Th = ma_irf(fit, H); m = fit["m"]
    IRF = []; ES = []
    for _ in range(ntry):
        Q, R = np.linalg.qr(rng.standard_normal((m, m)))
        Q = Q * np.sign(np.diag(R))
        A0 = C @ Q
        A0 = A0 * np.sign(A0[2])[None, :]                    # normalise each shock so real oil price rises
        assign = {}
        okc = True
        for c in range(m):
            pp, ra = A0[0, c], A0[1, c]
            s = ("supply" if (pp < 0 and ra < 0) else "aggdem" if (pp > 0 and ra > 0)
                 else "oilspec" if (pp > 0 and ra < 0) else None)
            if s is None or s in assign:
                okc = False; break
            assign[s] = c
        if not okc or len(assign) != 3:
            continue
        es = A0[0, assign["aggdem"]] / A0[2, assign["aggdem"]]    # short-run supply elasticity (along demand shift)
        if e_supply_max is not None and abs(es) > e_supply_max:
            continue
        order = [assign["supply"], assign["aggdem"], assign["oilspec"]]
        A0o = A0[:, order]
        IRF.append(np.array([Th[h] @ A0o for h in range(H + 1)]))
        ES.append(es)
    return {"irf": np.array(IRF), "elast_supply": np.array(ES), "accept_frac": len(IRF) / ntry}


def _mn_iw_draw(fit, rng):
    """One draw from the diffuse Normal-inverse-Wishart posterior of (B, Sigma)."""
    X = fit["X"]; m = fit["m"]; k = X.shape[1]; T = X.shape[0]
    XtXi = np.linalg.inv(X.T @ X)
    S = fit["resid"].T @ fit["resid"]
    L = np.linalg.cholesky(np.linalg.inv(S)); A = L @ rng.standard_normal((m, T - k))
    Sig = np.linalg.inv(A @ A.T)
    B = fit["B"] + np.linalg.cholesky(XtXi) @ rng.standard_normal((k, m)) @ np.linalg.cholesky(Sig).T
    return B, Sig


def posterior_chol_irf(fit, H, ndraw=1000, seed=0):
    """Recursive structural IRFs over the Normal-IW posterior.  Returns (ndraw, H+1, m, m)."""
    rng = np.random.default_rng(seed)
    out = []
    for _ in range(ndraw):
        B, Sig = _mn_iw_draw(fit, rng)
        out.append(chol_irf({"B": B, "Sigma": Sig, "p": fit["p"], "m": fit["m"]}, H))
    return np.array(out)


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