"""
condlca.py -- LATENT CLASS ANALYSIS with CONDITIONAL DEPENDENCE (from scratch).

Backs the notebooks in  "Relaxing Local Independence".

Basic LCA / Hui-Walter assumes LOCAL INDEPENDENCE -- the tests (items) are independent
given the true class. When that fails (correlated tests), the naive model biases the
sensitivities/specificities and can invent spurious extra classes. We relax the
assumption with a RANDOM-EFFECTS (probit factor) model (Qu, Tan & Kutner 1996): each
subject carries a continuous latent 'severity' b_i that shifts all their test
probabilities together, inducing within-class correlation:

    T_i ~ Categorical(pi)                         (latent class / disease status)
    b_i ~ Normal(0, 1)                            (subject random effect = a latent trait)
    x_ij | T_i=c, b_i ~ Bernoulli( Phi( a_{cj} + beta_j b_i ) )     (probit)

beta_j is test j's LOADING on the shared effect: beta_j = 0 for all j recovers ordinary
LCA (local independence). The marginal test accuracies integrate the random effect out
in closed form (the probit-normal integral):

    Se_j = Phi( a_{1j} / sqrt(1 + beta_j^2) ),   1 - Sp_j = Phi( a_{0j} / sqrt(1 + beta_j^2) ).

The from-scratch sampler is an ALBERT-CHIB data-augmentation Gibbs (the same truncated-
normal augmentation used for probit / Tobit): introduce z_ij ~ N(a_{T_i,j}+beta_j b_i, 1)
truncated by the sign of x_ij, after which a, beta, b and T all have conjugate Gaussian /
categorical full conditionals.
"""

import numpy as np
from scipy.special import ndtr as Phi          # standard normal CDF
from scipy.special import ndtri as Phinv        # inverse


# --------------------------------------------------------------------------- #
#  Truncated-normal helper (Albert-Chib latent z)                               #
# --------------------------------------------------------------------------- #

def _rtruncnorm_sign(mean, positive, rng):
    """Draw N(mean,1) truncated to (0, inf) if positive else (-inf, 0), by inverse-CDF."""
    lo = np.where(positive, Phi(-mean), 0.0)
    hi = np.where(positive, 1.0, Phi(-mean))
    u = lo + rng.random(mean.shape) * (hi - lo)
    return mean + Phinv(np.clip(u, 1e-12, 1 - 1e-12))


# --------------------------------------------------------------------------- #
#  Random-effects (probit factor) LCA -- the conditional-dependence model       #
# --------------------------------------------------------------------------- #

def re_lca_gibbs(X, rng, draws=4000, burn=1500, a_sd=3.0, beta_sd=0.4, anchor=True):
    """Albert-Chib Gibbs for the 2-class probit random-effects LCA.
    Returns posterior draws of prevalence, marginal Se, Sp, and the loadings beta.
    NOTE: beta_sd is a half-normal-scale REGULARISER on the loadings. A diffuse prior
    (beta_sd>=1) lets the sampler collapse into a degenerate 'one big random effect'
    mode (a factor-mixture identifiability trap), inflating every loading and deflating
    Se; beta_sd=0.4 anchors it and recovers the truth while still detecting real
    dependence (large true loadings overcome the mild shrinkage)."""
    X = np.asarray(X, float); N, J = X.shape
    # init
    T = (X.mean(1) > X.mean()).astype(int)
    alpha = np.zeros((2, J)); beta = rng.uniform(0.2, 0.6, J); b = rng.standard_normal(N)
    pi = 0.5
    A_prec = 1.0 / a_sd ** 2; B_prec = 1.0 / beta_sd ** 2
    PREV = np.empty(draws); SE = np.empty((draws, J)); SP = np.empty((draws, J)); BETA = np.empty((draws, J))
    for it in range(draws + burn):
        mu = alpha[T] + beta[None, :] * b[:, None]              # (N,J)
        z = _rtruncnorm_sign(mu, X > 0.5, rng)                  # augment latent z
        # b_i | z, alpha, beta, T
        r = z - alpha[T]                                        # (N,J)
        prec_b = 1.0 + np.sum(beta ** 2)
        mean_b = (r @ beta) / prec_b
        b = mean_b + rng.standard_normal(N) / np.sqrt(prec_b)
        # alpha[c,j] | z, b, T, beta
        for c in (0, 1):
            m = T == c; nc = int(m.sum())
            resid = z[m] - beta[None, :] * b[m, None]           # (nc,J)
            prec = A_prec + nc
            mean = resid.sum(0) / prec
            alpha[c] = mean + rng.standard_normal(J) / np.sqrt(prec)
        # beta_j | z, alpha, b, T
        rj = z - alpha[T]                                       # (N,J)
        prec = B_prec + np.sum(b ** 2)
        mean = (b[:, None] * rj).sum(0) / prec
        beta = mean + rng.standard_normal(J) / np.sqrt(prec)
        # T_i | z, alpha, beta, b, pi   (Gaussian likelihood of z given class)
        d1 = z - (alpha[1] + beta[None, :] * b[:, None]); d0 = z - (alpha[0] + beta[None, :] * b[:, None])
        ll1 = np.log(pi) - 0.5 * (d1 ** 2).sum(1); ll0 = np.log1p(-pi) - 0.5 * (d0 ** 2).sum(1)
        p1 = 1.0 / (1.0 + np.exp(np.clip(ll0 - ll1, -700, 700)))
        T = (rng.random(N) < p1).astype(int)
        pi = rng.beta(1 + int(T.sum()), 1 + N - int(T.sum()))
        # identify sign of the shared effect: keep sum(beta) >= 0
        if beta.sum() < 0:
            beta = -beta; b = -b
        # marginal accuracies (integrate out b)
        s = np.sqrt(1 + beta ** 2)
        Se = Phi(alpha[1] / s); Sp = 1 - Phi(alpha[0] / s)
        if anchor and (Se.mean() + Sp.mean() < 1):              # anchor disease/healthy labelling
            alpha = alpha[::-1].copy(); T = 1 - T; pi = 1 - pi
            Se = Phi(alpha[1] / s); Sp = 1 - Phi(alpha[0] / s)
        if it >= burn:
            i = it - burn; PREV[i] = pi; SE[i] = Se; SP[i] = Sp; BETA[i] = beta
    return dict(prev=PREV, Se=SE, Sp=SP, beta=BETA)


# --------------------------------------------------------------------------- #
#  Marginal likelihood + multi-restart (the factor-mixture posterior is multimodal)#
# --------------------------------------------------------------------------- #

def marginal_loglik(X, prev, Se, Sp, beta, n_quad=20):
    """Observed-data log-likelihood of the random-effects model, integrating the random
    effect out by Gauss-Hermite quadrature. Used to compare restarts."""
    X = np.asarray(X, float); gx, gw = np.polynomial.hermite_e.hermegauss(n_quad)
    gw = gw / np.sqrt(2 * np.pi)
    s = np.sqrt(1 + np.asarray(beta) ** 2)
    a1 = Phinv(np.clip(Se, 1e-6, 1 - 1e-6)) * s; a0 = Phinv(np.clip(1 - np.asarray(Sp), 1e-6, 1 - 1e-6)) * s
    p1 = Phi(a1[None, :] + beta[None, :] * gx[:, None]); p0 = Phi(a0[None, :] + beta[None, :] * gx[:, None])
    Li = np.zeros(len(X))
    for q in range(len(gx)):
        c1 = np.prod(np.where(X == 1, p1[q], 1 - p1[q]), 1); c0 = np.prod(np.where(X == 1, p0[q], 1 - p0[q]), 1)
        Li += gw[q] * (prev * c1 + (1 - prev) * c0)
    return float(np.log(Li + 1e-300).sum())

def re_lca_gibbs_ms(X, rng, restarts=5, draws=2000, burn=1200, beta_sd=0.4):
    """Random-effects LCA with MULTIPLE RESTARTS, returning the run with the best marginal
    fit. The factor-mixture posterior is multimodal -- a diffuse or unlucky run collapses
    toward a single dominant factor -- so restarts with marginal-likelihood selection are
    needed for reliable recovery. Returns the best run's full posterior draws."""
    best = None
    for _ in range(restarts):
        res = re_lca_gibbs(X, rng, draws=draws, burn=burn, beta_sd=beta_sd)
        ll = marginal_loglik(X, res["prev"].mean(), res["Se"].mean(0), res["Sp"].mean(0), res["beta"].mean(0))
        if best is None or ll > best[0]:
            best = (ll, res)
    out = best[1]; out["marginal_loglik"] = best[0]
    return out


# --------------------------------------------------------------------------- #
#  Naive LOCAL-INDEPENDENCE 2-class model (Hui-Walter), for comparison          #
# --------------------------------------------------------------------------- #

def naive_lca_gibbs(X, rng, draws=4000, burn=1500, ab=(1.0, 1.0), anchor=True):
    """Ordinary 2-class LCA / Hui-Walter (conditional independence assumed)."""
    X = np.asarray(X, float); N, J = X.shape
    Se = rng.uniform(0.5, 0.9, J); Sp = rng.uniform(0.5, 0.9, J); pi = 0.4
    PREV = np.empty(draws); SE = np.empty((draws, J)); SP = np.empty((draws, J))
    for it in range(draws + burn):
        l1 = np.log(pi) + X @ np.log(Se) + (1 - X) @ np.log1p(-Se)
        l0 = np.log1p(-pi) + X @ np.log1p(-Sp) + (1 - X) @ np.log(Sp)
        p1 = 1.0 / (1.0 + np.exp(np.clip(l0 - l1, -700, 700)))
        T = (rng.random(N) < p1).astype(int); dis = T == 1; hea = ~dis
        pi = rng.beta(1 + int(dis.sum()), 1 + int(hea.sum()))
        sd = X[dis].sum(0); nd = int(dis.sum()); Se = rng.beta(ab[0] + sd, ab[1] + nd - sd)
        sh = X[hea].sum(0); nh = int(hea.sum()); Sp = rng.beta(ab[0] + (nh - sh), ab[1] + sh)
        if anchor and (Se.mean() + Sp.mean() < 1): Se, Sp, pi = 1 - Sp, 1 - Se, 1 - pi
        if it >= burn:
            i = it - burn; PREV[i] = pi; SE[i] = Se; SP[i] = Sp
    return dict(prev=PREV, Se=SE, Sp=SP)


# --------------------------------------------------------------------------- #
#  Simulation with known conditional dependence                                 #
# --------------------------------------------------------------------------- #

def simulate_dep(N, prev, Se, Sp, beta, rng):
    """Simulate 2-class data with a shared random effect of loadings `beta`
    (beta=0 -> conditionally independent). Target marginal Se, Sp are hit exactly."""
    Se = np.asarray(Se, float); Sp = np.asarray(Sp, float); beta = np.asarray(beta, float); J = len(Se)
    s = np.sqrt(1 + beta ** 2)
    a1 = Phinv(Se) * s; a0 = Phinv(1 - Sp) * s                  # back out probit intercepts
    T = (rng.random(N) < prev).astype(int); b = rng.standard_normal(N)
    A = np.where(T[:, None] == 1, a1[None, :], a0[None, :])
    p = Phi(A + beta[None, :] * b[:, None])
    X = (rng.random((N, J)) < p).astype(int)
    return X, T, b


def conditional_corr(X, T):
    """Empirical within-class correlation between tests (averaged over the two classes):
    the diagnostic of conditional dependence that local independence assumes is zero."""
    X = np.asarray(X, float); J = X.shape[1]
    cc = np.zeros((J, J))
    for c in (0, 1):
        Xc = X[T == c]
        if len(Xc) > 2: cc += np.corrcoef(Xc.T)
    return cc / 2