"""
huiwalter.py -- From-scratch HUI-WALTER latent class model for DIAGNOSTIC TESTING
without a gold standard.

Backs the notebooks in  "Diagnostic Testing Without a Gold Standard"  (the diagnostic-testing
project of the Latent Class Analysis arc).

When several imperfect binary tests are applied but the true disease status is never
observed, the Hui-Walter model treats that status as a LATENT CLASS and estimates,
jointly, the disease PREVALENCE and every test's SENSITIVITY and SPECIFICITY. It is a
two-class LCA with a clinical reading:

    T_i = 1 (diseased) with probability prevalence_g  (g = subject i's population)
    x_ij | T_i = 1 ~ Bernoulli(Se_j)          Se_j = sensitivity  = P(test j + | diseased)
    x_ij | T_i = 0 ~ Bernoulli(1 - Sp_j)       Sp_j = specificity  = P(test j - | healthy)

under LOCAL (conditional) INDEPENDENCE of the tests given true status.

Identifiability is the crux: with 2 tests in 1 population there are only 3 degrees of
freedom for 5 parameters (under-identified). Hui & Walter (1980) solve this with TWO
populations of different prevalence sharing the same test characteristics (6 df, 6
parameters); alternatively 3+ tests in one population is identified. Informative Beta
priors on Se/Sp can also resolve weak identifiability (Joseph, Gyorkos & Coupal 1995).

We fit by a conjugate data-augmentation Gibbs sampler (impute T; prevalence | T ~ Beta
per population; Se, Sp | T ~ Beta). Labels are anchored by requiring the tests to be
informative on average (Se + Sp > 1), which fixes the diseased/healthy labelling.
"""

import numpy as np


def hw_gibbs(X, pop, rng, se_ab=(1.0, 1.0), sp_ab=(1.0, 1.0), prev_ab=(1.0, 1.0),
             draws=5000, burn=1500, anchor=True):
    """Data-augmentation Gibbs for the Hui-Walter model.
    X: (N,J) binary test results ; pop: (N,) population index in 0..G-1.
    Priors: Se_j~Beta(se_ab), Sp_j~Beta(sp_ab), prevalence_g~Beta(prev_ab).
    Returns posterior draws of prevalence (draws,G), Se (draws,J), Sp (draws,J)."""
    X = np.asarray(X, float); N, J = X.shape
    pop = np.asarray(pop, int); G = int(pop.max()) + 1
    Se = rng.uniform(0.5, 0.9, J); Sp = rng.uniform(0.5, 0.9, J); prev = rng.uniform(0.2, 0.6, G)
    PREV = np.empty((draws, G)); SE = np.empty((draws, J)); SP = np.empty((draws, J))
    for t in range(draws + burn):
        # augment the latent disease status T_i
        l1 = np.log(prev[pop]) + X @ np.log(Se) + (1 - X) @ np.log1p(-Se)
        l0 = np.log1p(-prev[pop]) + 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
        # prevalence per population
        for g in range(G):
            m = pop == g; nd = int(T[m].sum())
            prev[g] = rng.beta(prev_ab[0] + nd, prev_ab[1] + int(m.sum()) - nd)
        # sensitivity: successes = positives among the diseased
        sd = X[dis].sum(0); nd = int(dis.sum())
        Se = rng.beta(se_ab[0] + sd, se_ab[1] + nd - sd)
        # specificity: successes = negatives among the healthy
        sh = X[hea].sum(0); nh = int(hea.sum())
        Sp = rng.beta(sp_ab[0] + (nh - sh), sp_ab[1] + sh)
        if anchor and (Se.mean() + Sp.mean() < 1.0):           # enforce 'tests informative on average'
            Se, Sp, prev = 1 - Sp, 1 - Se, 1 - prev
        if t >= burn:
            i = t - burn; PREV[i] = prev; SE[i] = Se; SP[i] = Sp
    return dict(prev=PREV, Se=SE, Sp=SP)


def simulate_hw(n_per_pop, prev, Se, Sp, rng):
    """Simulate Hui-Walter data. prev: (G,) prevalence per population; Se,Sp: (J,).
    Returns X (N,J), pop (N,), and the true disease status T (N,)."""
    prev = np.atleast_1d(np.asarray(prev, float)); Se = np.asarray(Se, float); Sp = np.asarray(Sp, float)
    G = len(prev); J = len(Se); n_per_pop = np.atleast_1d(n_per_pop)
    if len(n_per_pop) == 1: n_per_pop = np.repeat(n_per_pop, G)
    Xs, pops, Ts = [], [], []
    for g in range(G):
        n = int(n_per_pop[g]); T = (rng.random(n) < prev[g]).astype(int)
        p_pos = np.where(T[:, None] == 1, Se[None, :], 1 - Sp[None, :])
        Xs.append((rng.random((n, J)) < p_pos).astype(int)); pops.append(np.full(n, g)); Ts.append(T)
    return np.vstack(Xs), np.concatenate(pops), np.concatenate(Ts)


def summary(post, names=None):
    """Posterior mean and 95% credible interval for prevalence, Se and Sp."""
    import pandas as pd
    rows = []
    for g in range(post["prev"].shape[1]):
        v = post["prev"][:, g]; rows.append((f"prevalence[pop{g+1}]", v.mean(), *np.percentile(v, [2.5, 97.5])))
    J = post["Se"].shape[1]; names = names or [f"test{j+1}" for j in range(J)]
    for j in range(J):
        v = post["Se"][:, j]; rows.append((f"Se[{names[j]}]", v.mean(), *np.percentile(v, [2.5, 97.5])))
    for j in range(J):
        v = post["Sp"][:, j]; rows.append((f"Sp[{names[j]}]", v.mean(), *np.percentile(v, [2.5, 97.5])))
    return pd.DataFrame(rows, columns=["parameter", "mean", "lo95", "hi95"]).set_index("parameter")


def youden(post):
    """Youden's J = Se + Sp - 1 for each test (posterior draws), a single accuracy summary."""
    return post["Se"] + post["Sp"] - 1