"""Bayesian variable selection by exact model-space enumeration -- from scratch.

The *conjugate* branch of Bayesian variable selection (George & McCulloch, 1997,
"Approaches for Bayesian Variable Selection", Statistica Sinica 7, 339-373): put a
conjugate prior on each model's coefficients so that beta and sigma^2 can be
integrated out analytically, leaving a closed-form marginal likelihood p(y | model).
When p is small the whole 2^p model space can then be enumerated exactly -- no MCMC,
no search -- giving exact posterior model probabilities and marginal inclusion
probabilities.  Reproduces Congdon (2005), Bayesian Models for Categorical Data,
Example 3.3 (cystic-fibrosis PEmax data, 9 predictors -> 2^9 = 512 models).

Prior: Zellner's g-prior on the (centered) regression coefficients,
    beta_gamma | sigma^2, gamma ~ N(0, g sigma^2 (X_gamma' X_gamma)^{-1}),
with an intercept and sigma^2 given the standard reference prior.  Marginalizing
gives the classic closed form (Liang et al. 2008) in terms of the model's ordinary
R^2:
    p(y | gamma) proportional to (1+g)^{(n-1-p_gamma)/2} / [1 + g(1-R^2_gamma)]^{(n-1)/2},
with the null (intercept-only) model as the reference (log marginal likelihood 0).
g = n (the unit-information prior) is the default, which makes model choice behave
like BIC and matches R's BAS with alpha = n.
"""
import numpy as np


def _log_ml_gprior(R2, p_gamma, n, g):
    """Log marginal likelihood of a model (up to a constant common to all models)."""
    return 0.5 * (n - 1 - p_gamma) * np.log1p(g) - 0.5 * (n - 1) * np.log1p(g * (1.0 - R2))


def enumerate_bvs(y, X, g=None, names=None):
    """Enumerate all 2^p subsets of the columns of X under a g-prior.

    Returns a dict with, per model, its 0/1 inclusion vector, log marginal likelihood
    and posterior probability, plus marginal inclusion probabilities and the
    Bayesian-model-averaged coefficient vector (g-prior shrinkage applied).
    """
    y = np.asarray(y, float)
    X = np.asarray(X, float)
    n, p = X.shape
    if g is None:
        g = float(n)                                   # unit-information prior
    if names is None:
        names = [f"X{j+1}" for j in range(p)]
    yc = y - y.mean()
    Xc = X - X.mean(0)                                  # intercept handled by centering (always included)
    TSS = yc @ yc
    shrink = g / (1.0 + g)                              # posterior mean of beta = shrink * OLS

    P = 2 ** p
    gammas = np.zeros((P, p), int)
    lml = np.zeros(P)
    size = np.zeros(P, int)
    B = np.zeros((P, p))                                # shrunk coefficient per model (0 where excluded)
    for idx in range(P):
        gv = np.array([(idx >> j) & 1 for j in range(p)])
        cols = np.where(gv == 1)[0]
        pg = cols.size
        gammas[idx] = gv
        size[idx] = pg
        if pg == 0:
            R2 = 0.0
        else:
            Xg = Xc[:, cols]
            beta, *_ = np.linalg.lstsq(Xg, yc, rcond=None)
            resid = yc - Xg @ beta
            R2 = 1.0 - (resid @ resid) / TSS
            B[idx, cols] = shrink * beta
        lml[idx] = _log_ml_gprior(R2, pg, n, g)

    w = np.exp(lml - lml.max())
    w /= w.sum()                                       # posterior model probabilities (uniform model prior)
    incl = np.array([w[gammas[:, j] == 1].sum() for j in range(p)])
    bma = (w[:, None] * B).sum(0)                       # model-averaged coefficients
    order = np.argsort(-w)
    return dict(gammas=gammas, w=w, incl=incl, size=size, lml=lml, bma=bma,
                order=order, g=g, n=n, p=p, names=list(names))


def top_models(res, k=8):
    """Return the k highest-probability models as (included-names, posterior prob)."""
    out = []
    for idx in res["order"][:k]:
        inc = [res["names"][j] for j in range(res["p"]) if res["gammas"][idx, j] == 1]
        out.append((inc if inc else ["(intercept only)"], res["w"][idx]))
    return out