"""Blanchard-Perotti fiscal SVAR identification, layered on the shared reduced-form SVAR
toolkit (oilsvar.py, Part 8).

The reduced-form VAR, moving-average / impulse responses and Normal-inverse-Wishart posterior
draws are reused unaltered from oilsvar.py; this module adds only the fiscal identification:
recover the structural government-spending and (cyclically-adjusted) tax shocks from the
reduced-form residuals via institutional timing plus a calibrated output-elasticity of taxes,
then convert the elasticities to dollar multipliers.  NumPy only.
"""
import numpy as np
import oilsvar as O


def bp_B0(E, a1):
    """Blanchard-Perotti structural impact matrix B0 from reduced-form residuals E (columns g, t, y)
    and the calibrated output-elasticity of taxes a1.

    Identification: spending is predetermined within the quarter, so the spending shock is u_g;
    the automatic response of taxes to output is stripped with the calibrated elasticity, so the
    discretionary tax shock is u_t - a1 * u_y; output responds to both, estimated by instrumental
    variables using the structural fiscal shocks as instruments.
    """
    eg, et, ey = E[:, 0], E[:, 1], E[:, 2]
    ut = et - a1 * ey; ug = eg                                       # spending & cyclically-adjusted tax shocks
    cg, ct = np.linalg.solve(np.column_stack([ug, ut]).T @ np.column_stack([eg, et]),
                             np.column_stack([ug, ut]).T @ ey)       # output equation by IV
    den = 1 - ct * a1; B0 = np.zeros((3, 3))
    B0[:, 0] = [1.0, a1 * cg / den, cg / den]                        # response of [g, t, y] to a spending shock
    B0[:, 1] = [0.0, 1 + a1 * ct / den, ct / den]                    # ... to a tax shock
    B0[:, 2] = [0.0, a1 / den, 1 / den]                              # ... to an output shock
    return B0


def bp_B0_from_S(S, a1):
    """The same identification from a residual second-moment matrix S (used for posterior draws,
    where only the drawn reduced-form covariance is available)."""
    MtN = np.array([[S[0, 0], S[0, 1]], [S[1, 0] - a1 * S[2, 0], S[1, 1] - a1 * S[2, 1]]])
    Mty = np.array([S[0, 2], S[1, 2] - a1 * S[2, 2]])
    cg, ct = np.linalg.solve(MtN, Mty)
    den = 1 - ct * a1; B0 = np.zeros((3, 3))
    B0[:, 0] = [1.0, a1 * cg / den, cg / den]
    B0[:, 1] = [0.0, 1 + a1 * ct / den, ct / den]
    B0[:, 2] = [0.0, a1 / den, 1 / den]
    return B0


def multipliers(fit, a1, H, GY, TY):
    """Point dollar multiplier paths over horizon H: spending dY/dG and tax dY/dT.

    Converts the elasticity response of output to dollars with the spending (GY) and tax (TY)
    shares of GDP: dY/dG = (dlogY/dlogG) * (Y/G).
    """
    Th = O.ma_irf(fit, H); B0 = bp_B0(fit["resid"], a1)
    sp = np.array([(Th[h] @ B0)[2, 0] for h in range(H + 1)]) / GY
    tx = np.array([(Th[h] @ B0)[2, 1] for h in range(H + 1)]) / TY
    return sp, tx


def posterior_multipliers(fit, a1, H, GY, TY, ndraw=2000, seed=0):
    """Normal-inverse-Wishart posterior draws of the spending & tax dollar multiplier paths:
    draw (B, Sigma) from the reduced-form posterior (oilsvar._mn_iw_draw) and re-apply the
    Blanchard-Perotti identification on each draw.  Returns (spd, txd), each (ndraw, H+1)."""
    rng = np.random.default_rng(seed)
    spd = np.empty((ndraw, H + 1)); txd = np.empty((ndraw, H + 1))
    for i in range(ndraw):
        Bd, Sd = O._mn_iw_draw(fit, rng)
        Thd = O.ma_irf({"B": Bd, "Sigma": Sd, "p": fit["p"], "m": fit["m"]}, H)
        B0d = bp_B0_from_S(Sd, a1)
        spd[i] = np.array([(Thd[h] @ B0d)[2, 0] for h in range(H + 1)]) / GY
        txd[i] = np.array([(Thd[h] @ B0d)[2, 1] for h in range(H + 1)]) / TY
    return spd, txd