"""
lca.py -- From-scratch LATENT CLASS ANALYSIS for binary items.

Backs the notebooks in  "Latent Class Analysis -- Foundations"  (Project 1 of the Latent Class
Analysis arc).

The latent class model for N subjects, J binary items, C classes:

    T_i  ~ Categorical(lambda)                         (latent class of subject i)
    x_ij | T_i=c ~ Bernoulli(delta[c,j])               (LOCAL INDEPENDENCE)

so the marginal likelihood of a response pattern is a finite mixture of
product-Bernoulli components:

    P(x_i) = sum_c  lambda_c  prod_j  delta[c,j]^{x_ij} (1-delta[c,j])^{1-x_ij}.

Two from-scratch estimators are provided:
    em_lca     -- maximum likelihood by Expectation-Maximization (with random restarts)
    gibbs_lca  -- Bayesian posterior by DATA-AUGMENTATION Gibbs (the latent T_i is the
                  augmentation variable; lambda|T ~ Dirichlet, delta|T ~ Beta are conjugate)

plus helpers: log-likelihood, posterior class membership (responsibilities), the
entropy R^2 classification-quality index, expected pattern frequencies and the
G^2 / X^2 goodness-of-fit statistics, and a simulator for validation.

Label switching (the mixture-model identifiability issue) is handled by RELABELLING:
`order_by_prevalence` sorts classes so lambda is decreasing, giving a canonical labelling.
"""

import numpy as np
from itertools import product


# --------------------------------------------------------------------------- #
#  Core likelihood pieces                                                       #
# --------------------------------------------------------------------------- #

def _log_comp(X, lam, delta):
    """Per-subject, per-class log-joint  log lambda_c + sum_j log Bernoulli(x_ij; delta[c,j]).
    X: (N,J) 0/1 ; lam: (C,) ; delta: (C,J).  Returns (N,C)."""
    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))
    # (N,1,J) against (1,C,J), summed over the J items -> (N,C)
    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):
    """Posterior class-membership probabilities r[i,c] = P(T_i=c | x_i, lambda, 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):
    """Observed-data log-likelihood sum_i log P(x_i)."""
    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())


# --------------------------------------------------------------------------- #
#  Estimator 1 -- Expectation-Maximization (maximum likelihood)                 #
# --------------------------------------------------------------------------- #

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 it in range(max_iter):
        r = responsibilities(X, lam, delta)                      # E-step
        Nc = r.sum(0)                                            # M-step
        lam = Nc / N
        delta = (r.T @ X) / Nc[:, None]
        delta = np.clip(delta, 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):
    """MLE by EM with `n_starts` random restarts (the LCA likelihood is multimodal);
    returns the best (lambda, delta, loglik, responsibilities), relabelled by prevalence."""
    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])


# --------------------------------------------------------------------------- #
#  Estimator 2 -- Data-augmentation Gibbs sampler (Bayesian)                    #
# --------------------------------------------------------------------------- #

def gibbs_lca(X, C, rng, draws=4000, burn=1000, a_lam=1.0, a_delta=1.0, b_delta=1.0,
              relabel=True):
    """Bayesian LCA by data augmentation. Each sweep:
        T_i    ~ Categorical( responsibilities )                (impute the latent class)
        lambda ~ Dirichlet( a_lam + class counts )              (conjugate)
        delta[c,j] ~ Beta( a_delta + successes, b_delta + failures )   (conjugate)
    Returns posterior draws of lambda (draws,C) and delta (draws,C,J).
    If relabel, each draw is reordered so lambda is decreasing (kills label switching)."""
    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)); Traw = None
    total = draws + burn
    for t in range(total):
        r = responsibilities(X, lam, delta)                     # augment T
        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)                     # lambda | T
        for c in range(C):                                      # delta | T, x
            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)


# --------------------------------------------------------------------------- #
#  Label switching, fit, simulation                                             #
# --------------------------------------------------------------------------- #

def order_by_prevalence(lam, delta):
    """Canonical labelling: sort classes so lambda is decreasing."""
    o = np.argsort(lam)[::-1]
    return np.asarray(lam)[o], np.asarray(delta)[o]

def entropy_R2(resp):
    """Entropy-based classification quality in [0,1]; 1 = perfectly separated classes.
    Undefined for a single class (returns NaN)."""
    r = np.clip(resp, 1e-12, 1); N, C = r.shape
    if C < 2:
        return np.nan
    E = -(r * np.log(r)).sum()
    return 1 - E / (N * np.log(C))

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

def pattern_prob(patterns, lam, delta):
    """Model probability of each response pattern (rows of `patterns`)."""
    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):
    """Goodness of fit: observed vs expected frequencies over all 2^J patterns,
    with Pearson X^2 and likelihood-ratio G^2 (only sensible for small J)."""
    X = np.asarray(X, int); N, J = X.shape
    pats = _all_patterns(J)
    # observed counts per pattern
    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)

def simulate_lca(N, lam, delta, rng):
    """Draw an N x J binary data set from the LCA with the given lambda, delta."""
    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