"""
dplca.py -- From-scratch NONPARAMETRIC (Dirichlet-process) LATENT CLASS ANALYSIS.

Backs the notebooks in  "Nonparametric Latent Class Analysis"  (the nonparametric project of
the Latent Class Analysis arc).

A DP-LCA replaces the finite Dirichlet prior on C class weights with a DIRICHLET
PROCESS, giving a countably infinite number of latent classes of which only finitely
many are occupied by any finite sample. The number of occupied classes K is then a
posterior quantity -- inferred, not chosen. Model:

    (partition of subjects)   ~ CRP(alpha)                         [Chinese Restaurant Process]
    delta_{k,j}               ~ Beta(a, b)          i.i.d.         [base measure G0]
    x_ij | class k            ~ Bernoulli(delta_{k,j})            [local independence]

We sample it with the COLLAPSED CRP GIBBS sampler (Neal 2000, Algorithm 3): the
class-specific delta are integrated out, so each subject is reassigned using the
Beta-Bernoulli posterior predictive,

    P(join class k) proportional to  n_k  * prod_j  Pred(x_ij ; a+s_kj, b+n_k-s_kj)
    P(open a new class) proportional to alpha * prod_j Pred(x_ij ; a, b),

with Pred(x=1 ; A, B) = A/(A+B). The concentration alpha can be given a Gamma prior
and updated by the Escobar-West (1995) auxiliary-variable step.

Because the class labels are arbitrary, results are summarised by LABEL-INVARIANT
quantities: the posterior of K, and the co-clustering (posterior similarity) matrix.
"""

import numpy as np


# --------------------------------------------------------------------------- #
#  Collapsed CRP Gibbs sampler for the DP mixture of product-Bernoulli          #
# --------------------------------------------------------------------------- #

def dp_lca_gibbs(X, rng, alpha=1.0, a=1.0, b=1.0, draws=3000, burn=1000,
                 sample_alpha=False, alpha_prior=(2.0, 4.0), init_k=1):
    """Collapsed Gibbs (Neal Alg. 3) for DP-LCA. Returns the trace of the number of
    occupied classes K, the concentration alpha, and the (canonically relabelled)
    subject assignments at each retained sweep."""
    X = np.asarray(X, float); N, J = X.shape
    # initialise: init_k clusters at random
    z = rng.integers(0, init_k, N)
    clusters = {}
    for k in range(init_k):
        m = z == k
        if m.any():
            clusters[k] = dict(n=int(m.sum()), s=X[m].sum(0))
    # relabel to compact ids
    z, clusters, next_id = _compact(z, clusters)
    K_tr, al_tr, Z_tr = [], [], []
    for it in range(draws + burn):
        for i in range(N):
            xi = X[i]; ci = z[i]
            clusters[ci]["n"] -= 1; clusters[ci]["s"] = clusters[ci]["s"] - xi
            if clusters[ci]["n"] == 0:
                del clusters[ci]
            keys = list(clusters.keys())
            logp = np.empty(len(keys) + 1)
            for idx, k in enumerate(keys):
                nk = clusters[k]["n"]; sk = clusters[k]["s"]
                p1 = (a + sk) / (a + b + nk)
                logp[idx] = np.log(nk) + np.sum(xi * np.log(p1) + (1 - xi) * np.log1p(-p1))
            p1n = a / (a + b)
            logp[-1] = np.log(alpha) + np.sum(xi * np.log(p1n) + (1 - xi) * np.log(1 - p1n))
            logp -= logp.max(); pr = np.exp(logp); pr /= pr.sum()
            ch = rng.choice(len(keys) + 1, p=pr)
            if ch == len(keys):
                clusters[next_id] = dict(n=1, s=xi.copy()); z[i] = next_id; next_id += 1
            else:
                k = keys[ch]; clusters[k]["n"] += 1; clusters[k]["s"] = clusters[k]["s"] + xi; z[i] = k
        K = len(clusters)
        if sample_alpha:
            alpha = _update_alpha(alpha, K, N, rng, *alpha_prior)
        if it >= burn:
            K_tr.append(K); al_tr.append(alpha); Z_tr.append(_canonical(z))
    return dict(K=np.array(K_tr), alpha=np.array(al_tr), Z=np.array(Z_tr),
                a=a, b=b, X=np.asarray(X, float))


def _compact(z, clusters):
    old = sorted(clusters.keys()); remap = {o: i for i, o in enumerate(old)}
    z = np.array([remap[v] for v in z]); cl = {remap[o]: clusters[o] for o in old}
    return z, cl, len(old)

def _canonical(z):
    """Relabel a partition by order of first appearance (a label-invariant canonical form)."""
    seen = {}; out = np.empty(len(z), int); nxt = 0
    for i, v in enumerate(z):
        if v not in seen: seen[v] = nxt; nxt += 1
        out[i] = seen[v]
    return out


def _update_alpha(alpha, K, N, rng, a0, b0):
    """Escobar-West (1995) update of the DP concentration under a Gamma(a0, b0) prior."""
    eta = rng.beta(alpha + 1, N)
    c = a0 + K - 1; d = N * (b0 - np.log(eta))
    pi = c / (c + d)
    shape = a0 + K if rng.random() < pi else a0 + K - 1
    return rng.gamma(shape, 1.0 / (b0 - np.log(eta)))


# --------------------------------------------------------------------------- #
#  Label-invariant summaries                                                    #
# --------------------------------------------------------------------------- #

def coclustering(Z):
    """Posterior co-clustering (similarity) matrix: P(subjects i,j share a class)."""
    Z = np.asarray(Z); S, N = Z.shape
    P = np.zeros((N, N))
    for s in range(S):
        same = (Z[s][:, None] == Z[s][None, :])
        P += same
    return P / S

def representative_partition(Z, target_K=None):
    """A representative clustering: the retained draw whose partition is closest (in
    co-clustering Hamming distance) to the posterior mean, optionally restricted to
    draws with K = target_K (e.g. the modal K)."""
    Z = np.asarray(Z); P = coclustering(Z); S, N = Z.shape
    Ks = np.array([len(np.unique(z)) for z in Z])
    idx = np.arange(S) if target_K is None else np.where(Ks == target_K)[0]
    best, bd = None, np.inf
    for s in idx:
        same = (Z[s][:, None] == Z[s][None, :]).astype(float)
        d = np.abs(same - P).sum()
        if d < bd: bd, best = d, s
    return _canonical(Z[best])

def cluster_profiles(X, z, a=1.0, b=1.0):
    """Posterior-mean item-response probabilities delta_k for each cluster in partition z,
    ordered by cluster size (descending)."""
    X = np.asarray(X, float); z = np.asarray(z)
    ks, counts = np.unique(z, return_counts=True)
    order = ks[np.argsort(counts)[::-1]]
    lam = np.array([np.mean(z == k) for k in order])
    delta = np.array([(a + X[z == k].sum(0)) / (a + b + (z == k).sum()) for k in order])
    return lam, delta


def adjusted_rand(labels_true, labels_pred):
    """Adjusted Rand Index between two clusterings (label-invariant agreement; 1 = identical)."""
    from math import comb
    a = np.asarray(labels_true); b = np.asarray(labels_pred)
    ca = {v: i for i, v in enumerate(np.unique(a))}; cb = {v: i for i, v in enumerate(np.unique(b))}
    A = np.array([ca[v] for v in a]); B = np.array([cb[v] for v in b])
    cont = np.zeros((len(ca), len(cb)), int)
    for x, y in zip(A, B): cont[x, y] += 1
    sum_c = sum(comb(v, 2) for v in cont.sum(1))
    sum_k = sum(comb(v, 2) for v in cont.sum(0))
    sum_ij = sum(comb(v, 2) for v in cont.flat)
    n2 = comb(len(a), 2); exp = sum_c * sum_k / n2; mx = 0.5 * (sum_c + sum_k)
    return (sum_ij - exp) / (mx - exp) if mx != exp else 1.0


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