"""
irtcompare.py -- IRT MODEL COMPARISON: 1PL vs 2PL vs 3PL (from scratch).

Backs the notebooks in  "IRT Model Comparison -- 1PL vs 2PL vs 3PL".

The three dichotomous IRT models are NESTED:

    1PL:  P(x=1) = Phi(a*theta_i + d_j)             one common discrimination a
    2PL:  P(x=1) = Phi(a_j*theta_i + d_j)           per-item discrimination
    3PL:  P(x=1) = c_j + (1-c_j)*Phi(a_j*theta_i + d_j)   plus a lower asymptote (guessing)

Each adds parameters, so each fits the training data at least as well -- the in-sample log-
likelihood can only rise. The question model comparison answers is whether the extra parameters
earn their keep OUT of sample. WAIC and PSIS-LOO estimate the expected log pointwise predictive
density (elpd) with a complexity penalty, so a richer model is preferred only when it predicts
new responses better. The classic lesson: on the near-Rasch LSAT items the 1PL/2PL suffice and
3PL guessing is wasted, whereas on multiple-choice SAT12 the 3PL guessing parameter pays off.

All three samplers use the ALBERT-CHIB augmentation (the shared probit backbone): draw a latent
z_ij truncated by the response, after which theta and the item parameters are conjugate Gaussian.
The 3PL adds a latent "knows-it" indicator K_ij (Beguin-Glas): a correct response is either genuine
knowledge (K=1, prob Phi(eta)) or a lucky guess (K=0, prob c_j); the augmentation z is then driven
by K, and c_j is Beta among the non-knowers. Every sampler also stores the POINTWISE log-likelihood
log p(x_ij | draw) at each retained draw -- the raw material WAIC and LOO need.
"""

import numpy as np
from scipy.special import ndtr as Phi, ndtri as Phinv, logsumexp


def _rtrunc_sign(mean, pos, rng):
    lo = np.where(pos, Phi(-mean), 0.0); hi = np.where(pos, 1.0, Phi(-mean))
    u = lo + rng.random(mean.shape) * (hi - lo)
    return mean + Phinv(np.clip(u, 1e-12, 1 - 1e-12))


def _ll_probit(eta, X):
    """pointwise log-lik for P = Phi(eta)."""
    p = np.clip(Phi(eta), 1e-9, 1 - 1e-9)
    return np.where(X > 0.5, np.log(p), np.log1p(-p))


# --------------------------------------------------------------------------- #
#  1PL -- one common discrimination                                             #
# --------------------------------------------------------------------------- #

def irt1pl_gibbs(X, rng, draws=1500, burn=800, a_sd=3.0, d_sd=3.0):
    """P = Phi(a*theta + d_j), single slope a, theta ~ N(0,1)."""
    X = np.asarray(X, float); N, J = X.shape
    theta = rng.standard_normal(N); a = 1.0; d = np.zeros(J)
    A = np.empty(draws); D = np.empty((draws, J)); TH = np.empty((draws, N)); LL = np.empty((draws, N, J), np.float32)
    for it in range(draws + burn):
        eta = a * theta[:, None] + d[None, :]
        z = _rtrunc_sign(eta, X > 0.5, rng)
        prec = 1.0 + J * a * a                                        # theta_i | .   prior N(0,1)
        theta = (a * (z - d[None, :]).sum(1)) / prec + rng.standard_normal(N) / np.sqrt(prec)
        precd = N + 1.0 / d_sd ** 2                                   # d_j | .
        d = (z - a * theta[:, None]).sum(0) / precd + rng.standard_normal(J) / np.sqrt(precd)
        precA = J * (theta @ theta) + 1.0 / a_sd ** 2                 # a | .   (pooled regression of z-d on theta)
        a = (theta[:, None] * (z - d[None, :])).sum() / precA + rng.standard_normal() / np.sqrt(precA)
        if a < 0:
            a = -a; theta = -theta
        if it >= burn:
            i = it - burn; A[i] = a; D[i] = d; TH[i] = theta
            LL[i] = _ll_probit(a * theta[:, None] + d[None, :], X).astype(np.float32)
    return dict(a=A, d=D, theta=TH, loglik=LL)


# --------------------------------------------------------------------------- #
#  2PL -- per-item discrimination                                               #
# --------------------------------------------------------------------------- #

def irt2pl_gibbs(X, rng, draws=1500, burn=800, a_sd=3.0, d_sd=3.0):
    """P = Phi(a_j*theta + d_j), theta ~ N(0,1)."""
    X = np.asarray(X, float); N, J = X.shape
    theta = rng.standard_normal(N); a = np.ones(J); d = np.zeros(J)
    A = np.empty((draws, J)); D = np.empty((draws, J)); TH = np.empty((draws, N)); LL = np.empty((draws, N, J), np.float32)
    Pp = np.diag([1.0 / a_sd ** 2, 1.0 / d_sd ** 2])
    for it in range(draws + burn):
        eta = theta[:, None] * a[None, :] + d[None, :]
        z = _rtrunc_sign(eta, X > 0.5, rng)
        prec = 1.0 + a @ a                                            # theta_i | .   prior N(0,1)
        theta = ((z - d[None, :]) @ a) / prec + rng.standard_normal(N) / np.sqrt(prec)
        Xd = np.column_stack([theta, np.ones(N)])                    # (a_j,d_j) | .  regress z on [theta,1]
        V = np.linalg.inv(Xd.T @ Xd + Pp); M = V @ (Xd.T @ z); L = np.linalg.cholesky(V)
        dr = M + L @ rng.standard_normal((2, J)); a, d = dr[0], dr[1]
        if a.sum() < 0:
            a = -a; theta = -theta
        if it >= burn:
            i = it - burn; A[i] = a; D[i] = d; TH[i] = theta
            LL[i] = _ll_probit(theta[:, None] * a[None, :] + d[None, :], X).astype(np.float32)
    return dict(a=A, d=D, theta=TH, loglik=LL)


# --------------------------------------------------------------------------- #
#  3PL -- add a guessing lower asymptote                                        #
# --------------------------------------------------------------------------- #

def irt3pl_gibbs(X, rng, draws=1500, burn=800, a_sd=3.0, d_sd=3.0, c_a=1.0, c_b=4.0):
    """P = c_j + (1-c_j)*Phi(a_j*theta + d_j).  Beta(c_a,c_b) prior on guessing (mean ~ 0.2)."""
    X = np.asarray(X, float); N, J = X.shape
    theta = rng.standard_normal(N); a = np.ones(J); d = np.zeros(J); c = np.full(J, 0.2)
    A = np.empty((draws, J)); D = np.empty((draws, J)); C = np.empty((draws, J)); TH = np.empty((draws, N)); LL = np.empty((draws, N, J), np.float32)
    Pp = np.diag([1.0 / a_sd ** 2, 1.0 / d_sd ** 2])
    for it in range(draws + burn):
        eta = theta[:, None] * a[None, :] + d[None, :]; ph = Phi(eta)
        # latent "knows-it" indicator K: correct answers split into knowledge vs lucky guess
        pk = ph / (ph + c[None, :] * (1 - ph))                        # P(K=1 | x=1)
        K = np.where(X > 0.5, (rng.random((N, J)) < pk).astype(float), 0.0)
        z = _rtrunc_sign(eta, K > 0.5, rng)                          # augmentation driven by K
        prec = 1.0 + a @ a
        theta = ((z - d[None, :]) @ a) / prec + rng.standard_normal(N) / np.sqrt(prec)
        Xd = np.column_stack([theta, np.ones(N)])
        V = np.linalg.inv(Xd.T @ Xd + Pp); M = V @ (Xd.T @ z); L = np.linalg.cholesky(V)
        dr = M + L @ rng.standard_normal((2, J)); a, d = dr[0], dr[1]
        if a.sum() < 0:
            a = -a; theta = -theta
        # c_j | .  Beta among NON-knowers (K=0): guessed-right vs guessed-wrong
        nonk = (K < 0.5)
        g1 = (nonk & (X > 0.5)).sum(0); g0 = (nonk & (X < 0.5)).sum(0)
        c = rng.beta(c_a + g1, c_b + g0)
        if it >= burn:
            i = it - burn; A[i] = a; D[i] = d; C[i] = c; TH[i] = theta
            P = np.clip(c[None, :] + (1 - c[None, :]) * Phi(theta[:, None] * a[None, :] + d[None, :]), 1e-9, 1 - 1e-9)
            LL[i] = np.where(X > 0.5, np.log(P), np.log1p(-P)).astype(np.float32)
    return dict(a=A, d=D, c=C, theta=TH, loglik=LL)


# --------------------------------------------------------------------------- #
#  marginal (integrated) likelihood -- the RIGHT unit for IRT model comparison  #
# --------------------------------------------------------------------------- #

def marginal_loglik(X, res, model, n_gh=25):
    """Per-PERSON marginal log-likelihood at each draw, integrating theta out by Gauss-Hermite
    quadrature:  L_i = int prod_j p(x_ij | theta) N(theta;0,1) dtheta.  Returns (draws, N).

    Comparing IRT models by the CONDITIONAL likelihood (treating each response as the unit and
    counting the ability parameters) is badly behaved -- p_eff is dominated by the N abilities and
    the ranking is noisy. Integrating theta out makes the PERSON the exchangeable unit, so WAIC/LOO
    penalise only the item-parameter structure. Same Gauss-Hermite marginalisation used to score the
    latent-class models (see Latent Class Analysis -- Choosing the Number of Classes)."""
    nodes, wts = np.polynomial.hermite_e.hermegauss(n_gh)             # weight exp(-x^2/2)
    logw = np.log(wts / np.sqrt(2 * np.pi))                           # normalise to E under N(0,1)
    X = np.asarray(X, float); N, J = X.shape; S = len(res["theta"]); LLm = np.empty((S, N))
    for s in range(S):
        if model == "1pl":
            P = Phi(res["a"][s] * nodes[:, None] + res["d"][s][None, :])
        elif model == "2pl":
            P = Phi(nodes[:, None] * res["a"][s][None, :] + res["d"][s][None, :])
        else:
            c = res["c"][s]; P = c[None, :] + (1 - c[None, :]) * Phi(nodes[:, None] * res["a"][s][None, :] + res["d"][s][None, :])
        P = np.clip(P, 1e-9, 1 - 1e-9)                               # (Q, J)
        ll_iq = X @ np.log(P).T + (1 - X) @ np.log1p(-P).T           # (N, Q) log-lik at each node
        LLm[s] = logsumexp(ll_iq + logw[None, :], axis=1)
    return LLm


# --------------------------------------------------------------------------- #
#  information criteria                                                          #
# --------------------------------------------------------------------------- #

def waic(LL):
    """WAIC from a (draws, N, J) pointwise log-likelihood array (Watanabe / Gelman-Hwang-Vehtari).
    elpd_waic = sum_obs [ log mean_s p_s  -  var_s log p_s ]."""
    S = LL.shape[0]; ll = LL.reshape(S, -1).astype(float)
    lppd = logsumexp(ll, axis=0) - np.log(S)                          # log posterior predictive density per obs
    p_w = ll.var(0)                                                   # effective number of parameters per obs
    elpd = lppd - p_w                                                 # pointwise elpd
    n = elpd.size
    return dict(elpd=elpd.sum(), p_eff=p_w.sum(), waic=-2 * elpd.sum(),
                se=np.sqrt(n) * elpd.std(), pointwise=elpd)


def compare_elpd(a, b):
    """difference in pointwise elpd (model a - model b) with paired SE (McElreath / arviz style)."""
    d = a["pointwise"] - b["pointwise"]
    return dict(delta=d.sum(), se=np.sqrt(d.size) * d.std())


# --------------------------------------------------------------------------- #
#  posterior-predictive item fit                                                #
# --------------------------------------------------------------------------- #

def prob2(a, d, theta):
    return Phi(np.outer(theta, a) + d[None, :])


def prob3(a, d, c, theta):
    return c[None, :] + (1 - c[None, :]) * Phi(np.outer(theta, a) + d[None, :])


def item_fit_curve(theta, P_model, X, item, nbin=6):
    """observed vs model-expected proportion-correct for one item, binned by ability.
    Returns bin ability, observed proportion, model proportion -- the empirical ICC check."""
    order = np.argsort(theta); edges = np.array_split(order, nbin)
    xb = np.array([theta[b].mean() for b in edges])
    ob = np.array([X[b, item].mean() for b in edges])
    mb = np.array([P_model[b, item].mean() for b in edges])
    return xb, ob, mb