"""Bayesian variable selection for binary logistic regression -- from scratch.

Two routes for a binary outcome, both filling the "logistic selection" gap:

1. `pg_ssvs` -- spike-and-slab SSVS (as in Part 1), made to work for the non-conjugate
   logistic likelihood by **Polya-Gamma data augmentation** (Polson, Scott & Windle 2013).
   Introducing latent weights omega_i ~ PG(1, x_i'beta) renders the logistic likelihood
   conditionally Gaussian in beta, so the exact same conjugate spike-and-slab Gibbs sweep
   of the Gaussian case applies:  omega | beta  (Polya-Gamma), beta | omega, gamma (Normal),
   gamma_j | beta_j (Bernoulli).  The Polya-Gamma draws use the Devroye/PSW exact sampler.

2. `enumerate_logit` -- with only p=9 predictors the whole 2^9 = 512 model space can be
   scored exactly (as in Parts 2/4): each logistic model's marginal likelihood by a Laplace
   approximation of the GLM evidence, giving exact posterior model and inclusion probabilities.

Predictors are standardized; the intercept is always included.
"""
import numpy as np
from scipy.stats import norm
import statsmodels.api as sm

# ----------------------------------------------------------------------------- Polya-Gamma
_T = 0.64

def _a(n, x):
    if x > _T:
        return np.pi * (n + 0.5) * np.exp(-(n + 0.5) ** 2 * np.pi ** 2 * x / 2.0)
    return (2.0 / (np.pi * x)) ** 1.5 * np.pi * (n + 0.5) * np.exp(-2.0 * (n + 0.5) ** 2 / x)

def _mass_texpon(z):
    fz = np.pi ** 2 / 8 + z ** 2 / 2
    b = np.sqrt(1.0 / _T) * (_T * z - 1); a = -np.sqrt(1.0 / _T) * (_T * z + 1)
    x0 = np.log(fz) + fz * _T
    qdivp = 4 / np.pi * (np.exp(x0 - z + norm.logcdf(b)) + np.exp(x0 + z + norm.logcdf(a)))
    return 1.0 / (1.0 + qdivp)

def _rtigauss(z, rng):
    z = max(abs(z), 1e-9); mu = 1.0 / z; x = _T + 1.0
    if mu > _T:
        alpha = 0.0
        while rng.random() > alpha:
            E1 = rng.exponential(); E2 = rng.exponential()
            while E1 * E1 > 2 * E2 / _T:
                E1 = rng.exponential(); E2 = rng.exponential()
            x = _T / (1 + _T * E1) ** 2; alpha = np.exp(-0.5 * z * z * x)
    else:
        while x > _T:
            y = rng.normal() ** 2
            x = mu + 0.5 * mu * mu * y - 0.5 * mu * np.sqrt(4 * mu * y + (mu * y) ** 2)
            if rng.random() > mu / (mu + x): x = mu * mu / x
    return x

def _pg1(z, rng):
    """One draw from PG(1, z) (Devroye / Polson-Scott-Windle exact sampler)."""
    z = abs(z) * 0.5; fz = np.pi ** 2 / 8 + z * z / 2
    while True:
        x = _T + rng.exponential() / fz if rng.random() < _mass_texpon(z) else _rtigauss(z, rng)
        s = _a(0, x); y = rng.random() * s; n = 0
        while True:
            n += 1
            if n % 2 == 1:
                s -= _a(n, x)
                if y <= s: return 0.25 * x
            else:
                s += _a(n, x)
                if y > s: break


# ----------------------------------------------------------------------------- PG-SSVS
def pg_ssvs(y, X, tau=0.1, c=20.0, w=0.5, ndraw=4000, burn=2000, seed=0):
    """Spike-and-slab SSVS for logistic regression via Polya-Gamma augmentation."""
    rng = np.random.default_rng(seed)
    y = np.asarray(y, float); X = np.asarray(X, float)
    n, p = X.shape
    Xz = (X - X.mean(0)) / X.std(0)
    D = np.column_stack([np.ones(n), Xz])                     # [intercept, standardized predictors]
    kappa = y - 0.5
    beta = np.zeros(p + 1); gamma = np.ones(p, int)
    v_int = 100.0
    G = np.zeros((ndraw, p)); B = np.zeros((ndraw, p + 1))
    for it in range(ndraw + burn):
        psi = D @ beta
        omega = np.array([_pg1(z, rng) for z in psi])         # (1) omega | beta  ~ Polya-Gamma
        d2 = np.empty(p + 1); d2[0] = v_int                   # prior variances
        d2[1:] = np.where(gamma == 1, (c * tau) ** 2, tau ** 2)
        Prec = (D * omega[:, None]).T @ D + np.diag(1.0 / d2) # (2) beta | omega, gamma ~ Normal
        L = np.linalg.cholesky(Prec)
        m = np.linalg.solve(Prec, D.T @ kappa)
        beta = m + np.linalg.solve(L.T, rng.standard_normal(p + 1))
        bs = beta[1:]                                          # (3) gamma_j | beta_j  ~ Bernoulli
        l1 = np.log(w) - np.log(c * tau) - 0.5 * bs ** 2 / (c * tau) ** 2
        l0 = np.log(1 - w) - np.log(tau) - 0.5 * bs ** 2 / tau ** 2
        gamma = (rng.random(p) < 1.0 / (1.0 + np.exp(l0 - l1))).astype(int)
        if it >= burn:
            G[it - burn] = gamma; B[it - burn] = beta
    from collections import Counter
    freq = Counter(tuple(g) for g in G.astype(int))
    return dict(incl=G.mean(0), beta=B.mean(0), freq=freq, n=n, p=p)


# ----------------------------------------------------------------------------- exact enumeration
def enumerate_logit(y, X, names, tau_slope=2.0):
    """Laplace marginal likelihood over all 2^p logistic models (intercept always in)."""
    y = np.asarray(y, float); X = np.asarray(X, float); n, p = X.shape
    Xz = (X - X.mean(0)) / X.std(0)
    def logml(cols):
        d = 1 + len(cols)
        Xd = np.column_stack([np.ones(n)] + [Xz[:, j] for j in cols])
        res = sm.GLM(y, Xd, family=sm.families.Binomial()).fit()
        mu = res.mu; W = mu * (1 - mu)
        ll = np.sum(y * np.log(mu.clip(1e-12, 1)) + (1 - y) * np.log((1 - mu).clip(1e-12, 1)))
        th = res.params
        lp = norm.logpdf(th[0], 0, 10) + norm.logpdf(th[1:], 0, tau_slope).sum()
        H = (Xd * W[:, None]).T @ Xd
        _, logdet = np.linalg.slogdet(H)
        return ll + lp + 0.5 * d * np.log(2 * np.pi) - 0.5 * logdet
    P = 2 ** p; lml = np.zeros(P); gammas = np.zeros((P, p), int)
    for idx in range(P):
        gv = np.array([(idx >> j) & 1 for j in range(p)])
        gammas[idx] = gv; lml[idx] = logml(list(np.where(gv == 1)[0]))
    wgt = np.exp(lml - lml.max()); wgt /= wgt.sum()
    incl = np.array([wgt[gammas[:, j] == 1].sum() for j in range(p)])
    order = np.argsort(-wgt)
    top = [([names[j] for j in range(p) if gammas[i, j]] or ["(intercept only)"], wgt[i]) for i in order[:6]]
    return dict(incl=incl, top=top, names=list(names))