"""SSVS for VARs -- George, Sun & Ni (2008) -- from scratch.

George, E. I., Sun, D. & Ni, S. (2008), "Bayesian stochastic search for VAR model
restrictions", Journal of Econometrics 142, 553-580.

Extends spike-and-slab variable selection (Part 1) to a vector autoregression, and
does it on TWO sets of parameters at once:

  (a) the VAR coefficients  B (k x m, k = 1 + m*p), and
  (b) the off-diagonal elements of the upper-triangular Cholesky factor Psi of the
      error PRECISION,  Sigma^{-1} = Psi Psi'  --  i.e. which contemporaneous links
      between equations are present.

Each free element carries a latent inclusion indicator with a spike-and-slab prior:
    theta | gamma=0 ~ N(0, kappa0^2)   (spike),   theta | gamma=1 ~ N(0, kappa1^2)  (slab).
For the coefficients the spike/slab scales are "semiautomatic" (George-Sun-Ni): each
kappa is tied to that coefficient's OLS standard error, kappa0 = c0*se, kappa1 = c1*se.
Intercepts are always retained (diffuse prior, no indicator).

Gibbs sampler (one sweep):
  1. beta = vec(B) | Sigma, gamma   -- multivariate normal (GLS + spike/slab prior).
  2. gamma_i | beta_i               -- Bernoulli from the spike/slab density ratio.
  3. Psi | B, omega                 -- column by column: diagonal psi_jj^2 ~ Gamma,
                                       off-diagonals eta_j | psi_jj ~ Normal
                                       (the exact Cholesky-of-precision update).
  4. omega_ij | eta_ij              -- Bernoulli from the spike/slab density ratio.
"""
import numpy as np


def _lags(Y, p):
    """Build VAR design: response Yt (T x m) and regressors X = [1, y_{t-1},...,y_{t-p}]."""
    T, m = Y.shape
    Yt = Y[p:]
    X = np.column_stack([np.ones(T - p)] + [Y[p - l: T - l] for l in range(1, p + 1)])
    return Yt, X


def ssvs_var(Y, p=2, ndraw=6000, burn=3000, c0=0.1, c1=10.0,
             kap0_psi=0.1, kap1_psi=1.0, w=0.5, seed=0):
    """Run the George-Sun-Ni SSVS-VAR sampler.

    Returns posterior inclusion probabilities for the coefficients (k x m) and for the
    Psi off-diagonals, plus posterior mean coefficient matrix and error covariance.
    """
    rng = np.random.default_rng(seed)
    Y = np.asarray(Y, float)
    Yt, X = _lags(Y, p)
    T, m = Yt.shape
    k = X.shape[1]                                        # regressors per equation
    XtX = X.T @ X
    XtY = X.T @ Yt

    # --- OLS for semiautomatic spike/slab scales ---
    B_ols = np.linalg.solve(XtX, XtY)
    E_ols = Yt - X @ B_ols
    Sig_ols = (E_ols.T @ E_ols) / T
    XtX_inv = np.linalg.inv(XtX)
    se = np.sqrt(np.outer(np.diag(XtX_inv), np.diag(Sig_ols)))    # k x m OLS std errors
    intercept = np.zeros((k, m), bool); intercept[0, :] = True    # row 0 = intercept, always in

    kap0 = c0 * se; kap1 = c1 * se                                 # spike / slab sds (coefficients)
    v_int = 100.0                                                  # diffuse variance for intercepts

    # off-diagonal (i<j) index list for Psi
    offdiag = [(i, j) for j in range(m) for i in range(j)]
    n_off = len(offdiag)

    B = B_ols.copy()
    Sig = Sig_ols.copy()
    Sinv = np.linalg.inv(Sig)
    gam = np.ones((k, m), int)
    om = {ij: 1 for ij in offdiag}

    GAM = np.zeros((k, m)); OM = np.zeros(n_off)
    Bsum = np.zeros((k, m)); Ssum = np.zeros((m, m)); nkeep = 0

    for it in range(ndraw + burn):
        # ---- (1) beta = vec(B) | Sigma, gamma  (vectorized GLS + diagonal prior) ----
        d2 = np.where(gam == 1, kap1 ** 2, kap0 ** 2)             # prior variances (k x m)
        d2[intercept] = v_int
        Dinv = np.diag(1.0 / d2.flatten("F"))                    # column-major vec(B)
        Prec = np.kron(Sinv, XtX) + Dinv
        rhs = (XtY @ Sinv).flatten("F")                          # vec(X'Y Sigma^{-1})
        L = np.linalg.cholesky(Prec)
        mean = np.linalg.solve(Prec, rhs)
        beta = mean + np.linalg.solve(L.T, rng.standard_normal(k * m))
        B = beta.reshape((k, m), order="F")

        # ---- (2) gamma_i | beta_i  (Bernoulli spike/slab), intercepts fixed in ----
        l1 = np.log(w) - np.log(kap1) - 0.5 * B ** 2 / kap1 ** 2
        l0 = np.log(1 - w) - np.log(kap0) - 0.5 * B ** 2 / kap0 ** 2
        pg = 1.0 / (1.0 + np.exp(l0 - l1))
        gam = (rng.random((k, m)) < pg).astype(int)
        gam[intercept] = 1

        # ---- (3) Psi | B, omega  (Cholesky-of-precision, column by column) ----
        E = Yt - X @ B
        S = E.T @ E
        Psi = np.zeros((m, m))
        a0, b0 = 0.01, 0.01                                      # Gamma prior on psi_jj^2
        for j in range(m):
            if j == 0:
                psi2 = rng.gamma(a0 + T / 2.0, 1.0 / (b0 + 0.5 * S[0, 0]))
                Psi[0, 0] = np.sqrt(psi2)
            else:
                dη = np.array([kap1_psi ** 2 if om[(i, j)] else kap0_psi ** 2 for i in range(j)])
                A = S[:j, :j] + np.diag(1.0 / dη)
                Ainv = np.linalg.inv(A)
                s = S[:j, j]
                rate = b0 + 0.5 * (S[j, j] - s @ Ainv @ s)
                psi2 = rng.gamma(a0 + T / 2.0, 1.0 / max(rate, 1e-10))
                psi = np.sqrt(psi2)
                mη = -psi * (Ainv @ s)
                eta = mη + np.linalg.cholesky(Ainv) @ rng.standard_normal(j)
                Psi[j, j] = psi
                Psi[:j, j] = eta
                # ---- (4) omega_ij | eta_ij ----
                for i in range(j):
                    e = eta[i]
                    q1 = w * np.exp(-0.5 * e ** 2 / kap1_psi ** 2) / kap1_psi
                    q0 = (1 - w) * np.exp(-0.5 * e ** 2 / kap0_psi ** 2) / kap0_psi
                    om[(i, j)] = int(rng.random() < q1 / (q1 + q0))
        Sinv = Psi @ Psi.T
        Sig = np.linalg.inv(Sinv)

        if it >= burn:
            GAM += gam; OM += np.array([om[ij] for ij in offdiag])
            Bsum += B; Ssum += Sig; nkeep += 1

    return dict(incl_coef=GAM / nkeep, incl_psi=OM / nkeep, offdiag=offdiag,
                B_mean=Bsum / nkeep, Sigma_mean=Ssum / nkeep, se=se, m=m, p=p, k=k)


def coef_labels(m, p, varnames):
    """Row labels [const, var.L1, ..., var.Lp] for the k x m coefficient/inclusion matrix."""
    labs = ["const"]
    for l in range(1, p + 1):
        labs += [f"{v}.L{l}" for v in varnames]
    return labs