"""Horseshoe (global-local shrinkage) prior for a Bayesian VAR — from scratch.

Each reduced-form equation is a horseshoe regression: coefficient j has prior N(0, lambda_j^2 tau^2 sigma^2),
with local scales lambda_j ~ C+(0,1) and a global scale tau ~ C+(0,1). The global tau shrinks everything;
the heavy-tailed local lambda_j let a few important coefficients escape shrinkage -- adaptive, unlike the
fixed Minnesota structure. Sampled by the Makalic-Schmidt (2016) auxiliary-variable Gibbs (all inverse-gamma).
Self-contained (numpy/scipy); alters nothing else.
"""
import numpy as np
from scipy.stats import invgamma


def _build(Y, p):
    T, m = Y.shape
    X = np.hstack([np.ones((T - p, 1))] + [Y[p - l:T - l] for l in range(1, p + 1)])
    return Y[p:], X


def gibbs_hs_var(Y, p, ndraw=2000, burn=1000, seed=1):
    """Equation-by-equation horseshoe VAR.  Returns coefficient draws B (ndraw,k,m), Sigma draws,
    and kappa (ndraw,k,m): per-coefficient shrinkage weight 1/(1+lambda^2 tau^2) (1=fully shrunk)."""
    rng = np.random.default_rng(seed)
    Yt, X = _build(Y, p); n, m = Yt.shape; k = X.shape[1]
    XtX = X.T @ X
    ig = lambda a, scale: invgamma.rvs(a, scale=scale, random_state=rng)
    # per-equation state
    beta = np.zeros((k, m)); lam2 = np.ones((k, m)); nu = np.ones((k, m))
    tau2 = np.ones(m); xi = np.ones(m); sig2 = np.var(Yt, 0)
    keepB, keepK, keepS = [], [], []
    for it in range(ndraw + burn):
        for i in range(m):
            y = Yt[:, i]
            d = lam2[:, i] * tau2[i]; d[0] = 1e6                         # diffuse prior on the intercept
            Prec = XtX / sig2[i] + np.diag(1.0 / (d * sig2[i]))
            L = np.linalg.cholesky(Prec)
            mu = np.linalg.solve(L.T, np.linalg.solve(L, X.T @ y / sig2[i]))
            beta[:, i] = mu + np.linalg.solve(L.T, rng.standard_normal(k))
            b2 = beta[:, i] ** 2
            # local scales (Makalic-Schmidt)
            for j in range(1, k):
                lam2[j, i] = ig(1.0, 1.0 / nu[j, i] + b2[j] / (2 * tau2[i] * sig2[i]))
                nu[j, i] = ig(1.0, 1.0 + 1.0 / lam2[j, i])
            # global scale
            tau2[i] = ig((k) / 2.0, 1.0 / xi[i] + np.sum(b2[1:] / lam2[1:, i]) / (2 * sig2[i]))
            xi[i] = ig(1.0, 1.0 + 1.0 / tau2[i])
            # residual variance
            r = y - X @ beta[:, i]
            sig2[i] = ig((n + k - 1) / 2.0, (r @ r + np.sum(b2[1:] / (lam2[1:, i] * tau2[i]))) / 2.0)
        if it >= burn:
            E = Yt - X @ beta
            keepB.append(beta.copy()); keepS.append((E.T @ E) / n)
            kap = 1.0 / (1.0 + lam2 * tau2[None, :]); kap[0] = 0.0
            keepK.append(kap.copy())
    return {"B": np.array(keepB), "Sigma": np.array(keepS), "kappa": np.array(keepK),
            "p": p, "m": m, "k": k}