"""
Sign-restricted structural VAR (Uhlig 2005; Rubio-Ramirez, Waggoner & Zha 2010), on top of `bvar.py`.

The reduced-form BVAR posterior (bvar.gibbs_minnesota / gibbs_steady_state) gives draws of the VAR
coefficients and Sigma. Any structural decomposition is eps_t = P u_t with P = chol(Sigma) Q for some
orthogonal Q; the data do not pin Q, so it is SET-identified. Instead of a recursive ordering we impose
only the *signs* we believe in (e.g. a monetary tightening has rate up, inflation down) and keep every
rotation consistent with them -- the identified set, reported as an IRF band.

This module does NOT modify bvar.py: it consumes a bvar fit dict and adds the rotation search.
"""
import numpy as np


def _pi_blocks(coefs_draw, kind, m, p):
    """Extract the list [Pi_1, ..., Pi_p] (each m x m) from one posterior coefficient draw."""
    off = 1 if kind == "minnesota" else 0
    return [coefs_draw[off + (l - 1) * m: off + l * m].T for l in range(1, p + 1)]


def _ma_irf(Pis, m, p, H):
    """Reduced-form moving-average IRFs Theta_h (Theta_0 = I), shape (H+1, m, m), via the companion matrix."""
    F = np.zeros((m * p, m * p)); F[:m] = np.hstack(Pis)
    if p > 1:
        F[m:, :m * (p - 1)] = np.eye(m * (p - 1))
    J = np.zeros((m, m * p)); J[:, :m] = np.eye(m)
    Theta = np.empty((H + 1, m, m)); Fh = np.eye(m * p)
    for h in range(H + 1):
        Theta[h] = J @ Fh @ J.T; Fh = Fh @ F
    return Theta


def _satisfies(g, restr):
    """g is a candidate structural IRF (H+1, m); restr = list of (var_idx, sign, hmax)."""
    for vi, sgn, hmax in restr:
        if not (sgn * g[:hmax + 1, vi] >= 0).all():
            return False
    return True


def sign_irf(fit, restr, H=20, max_tries=200, seed=0):
    """One structural shock identified by sign restrictions.

    fit   : a bvar fit dict (kind 'minnesota' or 'steady_state').
    restr : list of (variable_index, +1/-1, max_horizon) the shock's IRF must obey.
    Returns dict(irf=(n_accepted, H+1, m), accept_frac). For each posterior draw we draw random
    Haar rotations until one column's IRF satisfies the signs (trying both sign flips), then keep it.
    """
    rng = np.random.default_rng(seed)
    coefs = fit["B"] if fit["kind"] == "minnesota" else fit["Pi"]
    Sig = fit["Sigma"]; m = Sig.shape[1]; p = fit["p"]
    kept = []; tries_hit = 0
    for d in range(len(coefs)):
        Theta = _ma_irf(_pi_blocks(coefs[d], fit["kind"], m, p), m, p, H)
        P0 = np.linalg.cholesky(Sig[d])
        found = None
        for _ in range(max_tries):
            q = rng.standard_normal(m); q /= np.linalg.norm(q)     # Haar-uniform impulse direction
            g = Theta @ (P0 @ q)                                   # candidate structural IRF (H+1, m)
            if _satisfies(g, restr):
                found = g; break
            if _satisfies(-g, restr):
                found = -g; break
        if found is not None:
            kept.append(found)
        else:
            tries_hit += 1
    return {"irf": np.array(kept), "accept_frac": len(kept) / len(coefs)}


def bands(res, qs=(16, 50, 84)):
    """Pointwise percentile bands of the accepted IRFs, shape (len(qs), H+1, m)."""
    return np.percentile(res["irf"], qs, axis=0)