"""Bayesian selection of log-linear models for contingency tables -- from scratch.

Albert, J. H. (1996), "The Bayesian selection of log-linear models", Canadian Journal
of Statistics 24, 327-347.  Reproduces the spirit of Congdon (2005) Chapter-7 log-linear
model comparison.

A contingency table is modelled as independent Poisson counts with a log-linear mean,
    log E[y_cell] = intercept + main effects + (some) interaction terms.
Model selection = deciding WHICH interaction terms (associations) to include.  For a
3-way table with factors A, G, D the candidate associations are the pairwise terms
AG, AD, GD and the three-way term AGD; keeping the model *hierarchical* (a higher-order
term requires all its lower-order relatives) leaves nine models, from mutual
independence [A][G][D] to the saturated [AGD].

Each model's marginal likelihood is obtained by a Laplace approximation of the Poisson
GLM evidence,
    log p(y | model) ~= loglik(theta_hat) + log prior(theta_hat) + (d/2) log 2pi - 1/2 log|H|,
with H = X' diag(mu_hat) X the Fisher information (the mode ~= MLE for well-populated
tables).  Following Albert, the *interaction* coefficients carry a heavy-tailed
**Cauchy(0, s)** prior -- robust, so the selection is not driven by the prior's tails --
while intercept and main effects get a diffuse normal.  Posterior model probabilities
and association inclusion probabilities follow from the marginal likelihoods under a
uniform prior over models.
"""
import numpy as np
import pandas as pd
import patsy
import statsmodels.api as sm


# the four candidate association terms for a 3-way table, and the hierarchy
def _hierarchical_models(f2):
    """All hierarchical models: any subset of the pairwise terms in f2, plus the
    saturated model (all pairwise + three-way) when all three pairwise terms are in."""
    from itertools import combinations
    models = []
    for r in range(len(f2) + 1):
        for combo in combinations(f2, r):
            models.append(list(combo))
    # saturated: all pairwise + the three-way term
    return models


def loglin_select(df, response, factors, s_cauchy=1.0, tau_main=10.0):
    """Enumerate hierarchical log-linear models for a 3-factor table and score each by
    its Laplace marginal likelihood (Albert's robust Cauchy prior on interactions).

    Returns per-model log marginal likelihoods and posterior probabilities, and the
    posterior inclusion probability of each association term.
    """
    A, G, D = factors
    pair = [f"{A}:{G}", f"{A}:{D}", f"{G}:{D}"]
    three = f"{A}:{G}:{D}"
    mains = [A, G, D]
    y = df[response].values.astype(float)

    # pre-build every term's design columns once, from the saturated formula
    rhs = " * ".join([f"C({v})" for v in factors])
    full = patsy.dmatrix(rhs, df, return_type="dataframe")
    slc = full.design_info.term_slices
    def cols(termname):
        # map "A:G" -> patsy term "C(A):C(G)"
        key = ":".join(f"C({t})" for t in termname.split(":"))
        for term, s in slc.items():
            if term.name() == key:
                return full.values[:, s]
        raise KeyError(key)
    inter = np.ones((len(df), 1))
    Xmain = np.column_stack([cols(v) for v in mains])
    n_main = Xmain.shape[1]

    def build(terms):
        blocks = [inter, Xmain]
        n_int = 0
        for t in terms:
            c = cols(t); blocks.append(c); n_int += c.shape[1]
        X = np.column_stack(blocks)
        return X, n_int

    def cauchy_logpdf(x, s):
        return -np.log(np.pi * s) - np.log1p((x / s) ** 2)

    def normal_logpdf(x, sd):
        return -0.5 * np.log(2 * np.pi * sd ** 2) - 0.5 * (x / sd) ** 2

    # candidate models: subsets of pairwise terms + saturated
    subsets = _hierarchical_models(pair)
    models = []
    for terms in subsets:
        label = "+".join(t.replace(":", "") for t in terms) if terms else "indep"
        models.append((label, list(terms)))
    models.append(("saturated", pair + [three]))         # all pairwise + three-way

    rows = []
    for label, terms in models:
        X, n_int = build(terms)
        d = X.shape[1]
        res = sm.GLM(y, X, family=sm.families.Poisson()).fit()
        theta = res.params
        mu = res.mu
        ll = np.sum(y * np.log(mu) - mu)                 # Poisson log-likelihood (up to const)
        # log prior: diffuse normal on intercept+mains, robust Cauchy on interaction coefs
        lp = normal_logpdf(theta[0], 1000.0) + normal_logpdf(theta[1:1 + n_main], tau_main).sum()
        if n_int:
            lp += cauchy_logpdf(theta[1 + n_main:], s_cauchy).sum()
        H = X.T @ (mu[:, None] * X)                      # Fisher information X' W X
        sign, logdet = np.linalg.slogdet(H)
        logml = ll + lp + 0.5 * d * np.log(2 * np.pi) - 0.5 * logdet
        rows.append(dict(model=label, terms=terms, k=d, logml=logml))

    R = pd.DataFrame(rows)
    R["postprob"] = np.exp(R["logml"] - R["logml"].max())
    R["postprob"] /= R["postprob"].sum()
    R = R.sort_values("postprob", ascending=False).reset_index(drop=True)

    # association inclusion probabilities
    assoc = pair + [three]
    incl = {a: R.loc[R["terms"].apply(lambda ts: a in ts), "postprob"].sum() for a in assoc}
    return R, incl