"""Stochastic Search Variable Selection (SSVS) -- from scratch.

George, E. I. & McCulloch, R. E. (1993), "Variable Selection via Gibbs Sampling",
JASA 88, 881-889.  Reproduces the selection example of Congdon (2005), *Bayesian
Models for Categorical Data*, Program 3.10.

Model (linear regression, p candidate predictors, all columns standardized):

    y = X beta + eps,        eps ~ N(0, sigma^2 I)

Spike-and-slab mixture prior on each coefficient, controlled by a latent
inclusion indicator gamma_j in {0,1}:

    beta_j | gamma_j = 0 ~ N(0, tau_j^2)        (spike: pinned near zero)
    beta_j | gamma_j = 1 ~ N(0, c^2 tau_j^2)    (slab: diffuse, effect allowed)
    gamma_j ~ Bernoulli(w)                       (prior inclusion probability)
    sigma^2 ~ Inverse-Gamma(a0, b0)

Unlike a point-mass prior (Kuo-Mallick), the spike is a *narrow but proper*
normal, so every full conditional is closed form and the Gibbs sampler mixes
well.  gamma_j is drawn from the ratio of the two normal densities evaluated at
the current beta_j -- large |beta_j| pulls gamma_j toward 1 (slab), small |beta_j|
toward 0 (spike).  Sweeping gamma over 2^p models is the "stochastic search".

y and X are centered/standardized inside the sampler; the intercept is handled by
centering and reported back on the original scale.
"""
import numpy as np
from collections import Counter


def _inv_gamma(rng, shape, scale):
    """Draw from Inverse-Gamma(shape, scale) via the Gamma reciprocal."""
    return 1.0 / rng.gamma(shape, 1.0 / scale)


def ssvs_gibbs(y, X, ndraw=8000, burn=4000, tau=0.05, c=20.0, w=0.5,
               a0=1e-3, b0=1e-3, seed=0):
    """Run the George-McCulloch SSVS Gibbs sampler.

    Parameters
    ----------
    y, X   : response (n,) and predictor matrix (n, p); standardized internally.
    tau    : spike standard deviation (small -> coefficient forced near 0).
    c      : slab-to-spike ratio; slab sd = c*tau (large -> real effect allowed).
    w      : prior inclusion probability P(gamma_j = 1).
    a0, b0 : Inverse-Gamma shape/scale for the error variance sigma^2.

    Returns a dict of posterior draws plus marginal inclusion probabilities and
    the frequency table over visited 0/1 models.
    """
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float)
    X = np.asarray(X, float)
    n, p = X.shape

    # --- standardize: comparable columns so one (tau, c) governs all coefficients
    ybar, ysd = y.mean(), y.std()
    yz = (y - ybar) / ysd
    Xm, Xs = X.mean(0), X.std(0)
    Xz = (X - Xm) / Xs

    tau = np.full(p, float(tau)) if np.ndim(tau) == 0 else np.asarray(tau, float)
    XtX, Xty = Xz.T @ Xz, Xz.T @ yz

    beta = np.zeros(p)
    gamma = np.ones(p, int)
    sig2 = float(np.var(yz))
    G = np.zeros((ndraw, p)); B = np.zeros((ndraw, p)); S = np.zeros(ndraw)

    for it in range(ndraw + burn):
        # (1) beta | gamma, sigma^2, y  -- multivariate normal full conditional
        d2 = np.where(gamma == 1, (c * tau) ** 2, tau ** 2)     # prior variances
        Prec = XtX / sig2 + np.diag(1.0 / d2)                   # posterior precision
        L = np.linalg.cholesky(Prec)
        mu = np.linalg.solve(Prec, Xty / sig2)
        beta = mu + np.linalg.solve(L.T, rng.standard_normal(p))

        # (2) gamma_j | beta_j  -- Bernoulli from the spike/slab density ratio
        slab, spike = (c * tau) ** 2, tau ** 2
        l1 = np.log(w) - 0.5 * np.log(slab) - 0.5 * beta ** 2 / slab
        l0 = np.log(1 - w) - 0.5 * np.log(spike) - 0.5 * beta ** 2 / spike
        pg = 1.0 / (1.0 + np.exp(l0 - l1))
        gamma = (rng.random(p) < pg).astype(int)

        # (3) sigma^2 | beta, y  -- Inverse-Gamma full conditional
        r = yz - Xz @ beta
        sig2 = _inv_gamma(rng, a0 + n / 2.0, b0 + 0.5 * r @ r)

        if it >= burn:
            k = it - burn
            G[k], B[k], S[k] = gamma, beta, sig2

    incl = G.mean(0)
    freq = Counter(tuple(g) for g in G.astype(int))
    # coefficients rescaled back to the original y/X units (for reporting)
    B_orig = B * (ysd / Xs)
    return dict(gamma=G, beta=B, beta_orig=B_orig, sigma2=S, incl=incl, freq=freq,
                ybar=ybar, ysd=ysd, Xmean=Xm, Xstd=Xs, tau=tau, c=c, w=w)


def top_models(res, names, k=6):
    """Return the k most-visited models as (included-names, posterior probability)."""
    tot = sum(res["freq"].values())
    out = []
    for g, cnt in res["freq"].most_common(k):
        inc = [names[j] for j in range(len(g)) if g[j] == 1]
        out.append((inc if inc else ["(intercept only)"], cnt / tot))
    return out