"""
lcasel.py -- From-scratch LATENT CLASS ANALYSIS with the MODEL-SELECTION toolkit.

Backs the notebooks in  "Latent Class Analysis -- Choosing the Number of Classes"  (Project 2 of the Latent
Class Analysis arc): choosing the number of classes.

The core LCA (EM, data-augmentation Gibbs, likelihood, simulation) is copied here
so the folder is self-contained, and EXTENDED with the tools for deciding how many
latent classes a dataset supports:

    info_criteria      -- AIC, BIC, sample-size-adjusted BIC (SABIC), CAIC
    pointwise_loglik   -- per-subject log-likelihood over posterior draws (for WAIC/LOO)
    waic               -- Watanabe-Akaike information criterion, from scratch
    blrt               -- bootstrap likelihood-ratio test for C vs C+1 classes
    bivariate_residuals-- limited-information fit: observed vs expected pairwise association
    entropy_R2         -- classification separation

With J items there are 2^J response patterns, so for more than a handful of items the
full-table G^2 / X^2 becomes unusable (sparse); the BLRT and the bivariate residuals are
the practical alternatives.
"""

import numpy as np
from itertools import product
from scipy.special import gammaln, betaln, logsumexp


# --------------------------------------------------------------------------- #
#  Core LCA (self-contained copy)                                               #
# --------------------------------------------------------------------------- #

def _log_comp(X, lam, delta):
    X = np.asarray(X, float)
    ld = np.log(np.clip(delta, 1e-12, 1 - 1e-12)); l1 = np.log(np.clip(1 - delta, 1e-12, 1 - 1e-12))
    ll = (X[:, None, :] * ld[None] + (1 - X[:, None, :]) * l1[None]).sum(-1)
    return np.log(np.clip(lam, 1e-300, None))[None, :] + ll

def responsibilities(X, lam, delta):
    lc = _log_comp(X, lam, delta); lc -= lc.max(1, keepdims=True)
    r = np.exp(lc); return r / r.sum(1, keepdims=True)

def loglik(X, lam, delta):
    lc = _log_comp(X, lam, delta); m = lc.max(1, keepdims=True)
    return float((m[:, 0] + np.log(np.exp(lc - m).sum(1))).sum())

def _em_once(X, C, rng, max_iter=500, tol=1e-8):
    N, J = X.shape
    delta = rng.uniform(0.1, 0.9, (C, J)); lam = rng.dirichlet(np.ones(C)); ll_old = -np.inf
    for _ in range(max_iter):
        r = responsibilities(X, lam, delta); Nc = r.sum(0)
        lam = Nc / N; delta = np.clip((r.T @ X) / Nc[:, None], 1e-6, 1 - 1e-6)
        ll = loglik(X, lam, delta)
        if ll - ll_old < tol: break
        ll_old = ll
    return lam, delta, loglik(X, lam, delta)

def em_lca(X, C, rng, n_starts=20, max_iter=500):
    X = np.asarray(X, float); best = None
    for _ in range(n_starts):
        lam, delta, ll = _em_once(X, C, rng, max_iter)
        if best is None or ll > best[2]: best = (lam, delta, ll)
    lam, delta, ll = best; lam, delta = order_by_prevalence(lam, delta)
    return dict(lam=lam, delta=delta, loglik=ll, resp=responsibilities(X, lam, delta),
                npar=C - 1 + C * X.shape[1], C=C)

def gibbs_lca(X, C, rng, draws=4000, burn=1000, a_lam=1.0, a_delta=1.0, b_delta=1.0, relabel=True):
    X = np.asarray(X, float); N, J = X.shape
    delta = rng.uniform(0.2, 0.8, (C, J)); lam = rng.dirichlet(np.ones(C))
    LAM = np.empty((draws, C)); DEL = np.empty((draws, C, J))
    for t in range(draws + burn):
        r = responsibilities(X, lam, delta); u = rng.random((N, 1))
        T = (np.cumsum(r, 1) > u).argmax(1); counts = np.bincount(T, minlength=C)
        lam = rng.dirichlet(a_lam + counts)
        for c in range(C):
            xc = X[T == c]; s = xc.sum(0) if len(xc) else np.zeros(J); f = (len(xc) - s) if len(xc) else np.zeros(J)
            delta[c] = rng.beta(a_delta + s, b_delta + f)
        if t >= burn:
            i = t - burn
            if relabel: lam, delta = order_by_prevalence(lam, delta)
            LAM[i] = lam; DEL[i] = delta
    return dict(lam=LAM, delta=DEL)

def order_by_prevalence(lam, delta):
    o = np.argsort(lam)[::-1]; return np.asarray(lam)[o], np.asarray(delta)[o]

def entropy_R2(resp):
    r = np.clip(resp, 1e-12, 1); N, C = r.shape
    if C < 2: return np.nan
    return 1 - (-(r * np.log(r)).sum()) / (N * np.log(C))

def simulate_lca(N, lam, delta, rng):
    lam = np.asarray(lam, float); delta = np.asarray(delta, float); C, J = delta.shape
    T = rng.choice(C, size=N, p=lam); X = (rng.random((N, J)) < delta[T]).astype(int)
    return X, T


# --------------------------------------------------------------------------- #
#  Model-selection toolkit                                                      #
# --------------------------------------------------------------------------- #

def info_criteria(ll, npar, N):
    """AIC, BIC, sample-size-adjusted BIC (SABIC, n*=(N+2)/24), and CAIC."""
    return dict(AIC=-2 * ll + 2 * npar,
                BIC=-2 * ll + npar * np.log(N),
                SABIC=-2 * ll + npar * np.log((N + 2) / 24),
                CAIC=-2 * ll + npar * (np.log(N) + 1))

def pointwise_loglik(X, lam_draws, delta_draws):
    """Per-subject log-likelihood log p(x_i | theta_s) at every posterior draw s.
    Returns (S, N) -- the input to WAIC and LOO."""
    X = np.asarray(X, float); S = len(lam_draws); N = X.shape[0]
    out = np.empty((S, N))
    for s in range(S):
        lc = _log_comp(X, lam_draws[s], delta_draws[s]); m = lc.max(1, keepdims=True)
        out[s] = m[:, 0] + np.log(np.exp(lc - m).sum(1))
    return out

def waic(llmat):
    """Watanabe-Akaike IC from the (S,N) pointwise log-likelihood.
    lppd_i = logmeanexp_s ll[s,i] ; p_waic_i = var_s ll[s,i] ; WAIC = -2 sum_i(lppd_i - p_waic_i)."""
    llmat = np.asarray(llmat, float); S, N = llmat.shape
    m = llmat.max(0)
    lppd = m + np.log(np.exp(llmat - m).mean(0))            # (N,)
    p_waic = llmat.var(0, ddof=1)                            # (N,)
    elpd_i = lppd - p_waic
    waic = -2 * elpd_i.sum()
    se = np.sqrt(N * elpd_i.var(ddof=1)) * 2
    return dict(WAIC=float(waic), lppd=float(lppd.sum()), p_waic=float(p_waic.sum()), se=float(se),
                elpd_i=elpd_i)

def blrt(X, C, rng, B=100, n_starts_obs=25, n_starts_boot=8):
    """Bootstrap likelihood-ratio test of C classes (null) vs C+1 (alt).
    LR = 2( loglik_{C+1} - loglik_C ); the null distribution of LR is obtained by parametric
    bootstrap -- simulate B datasets from the fitted C-class model and refit both. p = P(LR >= LR_obs)."""
    X = np.asarray(X, float); N = X.shape[0]
    f0 = em_lca(X, C, rng, n_starts=n_starts_obs)
    f1 = em_lca(X, C + 1, rng, n_starts=n_starts_obs)
    LR_obs = 2 * (f1["loglik"] - f0["loglik"])
    LR_null = np.empty(B)
    for b in range(B):
        Xb, _ = simulate_lca(N, f0["lam"], f0["delta"], rng)
        g0 = em_lca(Xb, C, rng, n_starts=n_starts_boot)
        g1 = em_lca(Xb, C + 1, rng, n_starts=n_starts_boot)
        LR_null[b] = 2 * (g1["loglik"] - g0["loglik"])
    pval = (1 + np.sum(LR_null >= LR_obs)) / (B + 1)
    return dict(C=C, LR_obs=float(LR_obs), LR_null=LR_null, pval=float(pval))

def bivariate_residuals(X, lam, delta):
    """Limited-information fit: for each item pair, the Pearson X^2 of the observed 2x2 table
    against its model-expected counts (~ chi^2_1; values > ~3.8 flag an under-fitted association)."""
    X = np.asarray(X, int); N, J = X.shape
    lam = np.asarray(lam, float); delta = np.asarray(delta, float)
    BVR = np.zeros((J, J))
    for j in range(J):
        for k in range(j + 1, J):
            tot = 0.0
            for a in (0, 1):
                for b in (0, 1):
                    O = np.sum((X[:, j] == a) & (X[:, k] == b))
                    pj = delta[:, j] if a == 1 else 1 - delta[:, j]
                    pk = delta[:, k] if b == 1 else 1 - delta[:, k]
                    E = N * np.sum(lam * pj * pk)
                    tot += (O - E) ** 2 / max(E, 1e-9)
            BVR[j, k] = BVR[k, j] = tot
    return BVR


# --------------------------------------------------------------------------- #
#  Marginal likelihood -- Chib's method (with the label-switching correction)   #
# --------------------------------------------------------------------------- #

def _dirichlet_logpdf(x, alpha):
    x = np.clip(np.asarray(x, float), 1e-300, None)
    return gammaln(alpha.sum()) - gammaln(alpha).sum() + np.sum((alpha - 1) * np.log(x))

def _beta_logpdf(x, a, b):
    x = np.clip(np.asarray(x, float), 1e-300, 1 - 1e-12)
    return (a - 1) * np.log(x) + (b - 1) * np.log1p(-x) - betaln(a, b)

def exact_logml_c1(X):
    """Exact log marginal likelihood of the ONE-class model (a product of independent
    Beta-Bernoulli items), used to validate the Chib estimator (which needs no correction at C=1)."""
    X = np.asarray(X, int); N, J = X.shape; s = X.sum(0)
    return float(np.sum(betaln(1 + s, 1 + N - s)))         # prod_j B(1+s_j, 1+f_j) / B(1,1)

def chib_logml(X, C, rng, draws=6000, burn=1500):
    """log marginal likelihood log p(x | C) by CHIB'S METHOD from the conjugate Gibbs sampler.
        log p(x) = log f(x|theta*) + log prior(theta*) - log posterior(theta*|x) + log C!
    The posterior ordinate is Rao-Blackwellised over the imputed class labels; the +log C!
    corrects for the C! symmetric relabellings the (relabelled) sampler does not visit."""
    X = np.asarray(X, float); N, J = X.shape
    delta = rng.uniform(0.2, 0.8, (C, J)); lam = rng.dirichlet(np.ones(C))
    LAM, DEL, NC, SS, FF = [], [], [], [], []
    for t in range(draws + burn):
        r = responsibilities(X, lam, delta); u = rng.random((N, 1))
        T = (np.cumsum(r, 1) > u).argmax(1); counts = np.bincount(T, minlength=C)
        s = np.array([X[T == c].sum(0) if (T == c).any() else np.zeros(J) for c in range(C)])
        f = counts[:, None] - s
        lam = rng.dirichlet(1 + counts)
        for c in range(C):
            delta[c] = rng.beta(1 + s[c], 1 + f[c])
        if t >= burn:
            o = np.argsort(lam)[::-1]                        # relabel by prevalence (one mode)
            LAM.append(lam[o]); DEL.append(delta[o]); NC.append(counts[o]); SS.append(s[o]); FF.append(f[o])
    LAM = np.array(LAM); DEL = np.array(DEL); NC = np.array(NC); SS = np.array(SS); FF = np.array(FF)
    lam_s = LAM.mean(0); del_s = DEL.mean(0)                 # theta* = posterior mean (an interior point)
    log_f = loglik(X, lam_s, del_s)
    log_prior = _dirichlet_logpdf(lam_s, np.ones(C))         # Beta(1,1) log-density = 0
    M = len(LAM); logdens = np.empty(M)
    for m in range(M):
        logdens[m] = _dirichlet_logpdf(lam_s, 1 + NC[m]) + _beta_logpdf(del_s, 1 + SS[m], 1 + FF[m]).sum()
    log_post = logsumexp(logdens) - np.log(M)               # Rao-Blackwellised posterior ordinate
    return dict(logml=float(log_f + log_prior - log_post + gammaln(C + 1)),
                theta_star=(lam_s, del_s), log_f=float(log_f), log_prior=float(log_prior),
                log_post=float(log_post), correction=float(gammaln(C + 1)))


# --------------------------------------------------------------------------- #
#  Full-table goodness of fit (only sensible for small J)                       #
# --------------------------------------------------------------------------- #

def _all_patterns(J):
    return np.array(list(product([0, 1], repeat=J)), float)

def pattern_prob(patterns, lam, delta):
    lc = _log_comp(patterns, lam, delta); m = lc.max(1, keepdims=True)
    return np.exp(m[:, 0] + np.log(np.exp(lc - m).sum(1)))

def gof(X, lam, delta):
    X = np.asarray(X, int); N, J = X.shape; pats = _all_patterns(J)
    key = lambda M: (M * (2 ** np.arange(J))).sum(1)
    obs = np.bincount(key(X).astype(int), minlength=2 ** J).astype(float)
    exp = N * pattern_prob(pats, lam, delta); exp_k = np.zeros(2 ** J); exp_k[key(pats).astype(int)] = exp
    mask = exp_k > 0; X2 = float(((obs[mask] - exp_k[mask]) ** 2 / exp_k[mask]).sum())
    o = obs[mask]; e = exp_k[mask]; nz = o > 0; G2 = float(2 * (o[nz] * np.log(o[nz] / e[nz])).sum())
    df = 2 ** J - 1 - (len(lam) - 1 + len(lam) * J)
    return dict(patterns=pats, observed=obs, expected=exp_k, X2=X2, G2=G2, df=df,
                sparsity=float(np.mean(exp_k < 1)))