"""
irt.py -- ITEM RESPONSE THEORY: 2PL and 3PL models (from scratch).

Backs the notebooks in  "Item Response Theory -- 2PL & 3PL".

Item response theory models a test: each person i has a continuous latent ABILITY theta_i, and
each binary item j has an item characteristic curve giving the probability of a correct answer as
a smooth function of ability. In the normal-ogive parameterisation

    2PL:  P(x_ij = 1 | theta_i) = Phi( a_j theta_i + d_j )
    3PL:  P(x_ij = 1 | theta_i) = c_j + (1 - c_j) Phi( a_j theta_i + d_j )

a_j is the item's DISCRIMINATION (how sharply it separates ability levels), the difficulty is
b_j = -d_j / a_j (the ability at which P = (1+c)/2), and c_j is the GUESSING floor (3PL only).
Setting all a_j equal is the RASCH / 1PL model -- a random-effects logistic regression, which is
why this connects to the panel-logit notebooks; here we read it psychometrically. (Normal-ogive
a's map to the logistic metric by a x 1.7.)

The from-scratch sampler is exactly ALBERT & CHIB's data augmentation (Albert 1992 -- the paper
that INTRODUCED this trick, for IRT, before it spread to probit/Tobit/multivariate-probit). Introduce
z_ij ~ N(a_j theta_i + d_j, 1) truncated by the sign of x_ij; then theta and (a_j, d_j) have
conjugate Gaussian full conditionals (a factor-analysis Gibbs). Scale/location are fixed by the
theta ~ N(0,1) prior, sign by keeping sum_j a_j > 0. The 3PL adds a latent "knows-it" indicator
K_ij ~ Bernoulli(Phi(a theta + d)) (Beguin & Glas 2001): correct answers by non-knowers are guesses,
so c_j gets a Beta posterior and the augmentation z is built on K rather than x.
"""

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


def _rtrunc_sign(mean, positive, rng):
    """draw N(mean,1) truncated to (0,inf) where 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))


# --------------------------------------------------------------------------- #
#  simulation                                                                   #
# --------------------------------------------------------------------------- #

def simulate_irt(n, a, d, rng, c=None, theta=None):
    """simulate a response matrix from a 2PL (c=None) or 3PL item bank."""
    a = np.asarray(a, float); d = np.asarray(d, float); J = len(a)
    if theta is None:
        theta = rng.standard_normal(n)
    P = Phi(np.outer(theta, a) + d[None, :])
    if c is not None:
        c = np.asarray(c, float); P = c[None, :] + (1 - c[None, :]) * P
    X = (rng.random((n, J)) < P).astype(float)
    return X, theta


# --------------------------------------------------------------------------- #
#  2PL Gibbs (Albert-Chib augmentation)                                         #
# --------------------------------------------------------------------------- #

def irt2pl_gibbs(X, rng, draws=3000, burn=1500, a_sd=3.0):
    """normal-ogive 2PL via Albert-Chib augmentation. Returns posterior draws of discriminations
    a:(draws,J), intercepts d:(draws,J) and abilities theta:(draws,N). Difficulty = -d/a."""
    X = np.asarray(X, float); N, J = X.shape
    theta = rng.standard_normal(N); a = np.ones(J); d = np.zeros(J)
    Pp = np.diag([1.0 / a_sd ** 2, 1e-6])              # prior precision for (a_j, d_j)
    A = np.empty((draws, J)); D = np.empty((draws, J)); TH = np.empty((draws, N))
    for it in range(draws + burn):
        mu = np.outer(theta, a) + d[None, :]
        z = _rtrunc_sign(mu, X > 0.5, rng)
        # theta_i | z, a, d  (N(0,1) prior fixes the scale)
        prec = 1.0 + np.sum(a ** 2)
        theta = ((z - d[None, :]) @ a) / prec + rng.standard_normal(N) / np.sqrt(prec)
        # (a_j, d_j) | z, theta  -- Bayesian regression of z[:,j] on [theta, 1]
        Xd = np.column_stack([theta, np.ones(N)])
        V = np.linalg.inv(Xd.T @ Xd + Pp); L = np.linalg.cholesky(V)
        M = V @ (Xd.T @ z)                             # (2, J)
        draw = M + L @ rng.standard_normal((2, J))
        a, d = draw[0], draw[1]
        if a.sum() < 0:                                # identify sign of the trait
            a = -a; theta = -theta
        if it >= burn:
            A[it - burn] = a; D[it - burn] = d; TH[it - burn] = theta
    return dict(a=A, d=D, theta=TH)


# --------------------------------------------------------------------------- #
#  3PL Gibbs (add a latent "knows-it" indicator for guessing)                   #
# --------------------------------------------------------------------------- #

def irt3pl_gibbs(X, rng, draws=4000, burn=2000, a_sd=3.0, cg=(1.0, 4.0)):
    """normal-ogive 3PL. cg = Beta prior on guessing c_j (default Beta(1,4), mean 0.2).
    Returns a, d, theta and guessing c:(draws,J)."""
    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)
    Pp = np.diag([1.0 / a_sd ** 2, 1e-6])
    A = np.empty((draws, J)); D = np.empty((draws, J)); TH = np.empty((draws, N)); C = np.empty((draws, J))
    for it in range(draws + burn):
        Pknow = Phi(np.outer(theta, a) + d[None, :])   # prob of KNOWING each item
        # latent knows-it indicator K: x=0 -> K=0 (would have answered right); x=1 -> Bernoulli
        pK1 = Pknow / (Pknow + (1 - Pknow) * c[None, :] + 1e-12)
        K = np.where(X > 0.5, (rng.random((N, J)) < pK1).astype(float), 0.0)
        # augment z on K (the probit "knows" process)
        mu = np.outer(theta, a) + d[None, :]
        z = _rtrunc_sign(mu, K > 0.5, rng)
        prec = 1.0 + np.sum(a ** 2)
        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); L = np.linalg.cholesky(V)
        M = V @ (Xd.T @ z); draw = M + L @ rng.standard_normal((2, J)); a, d = draw[0], draw[1]
        if a.sum() < 0: a = -a; theta = -theta
        # guessing c_j | data: among non-knowers (K=0), fraction answering correctly
        non = K < 0.5; corr = ((X > 0.5) & non).sum(0); tot = non.sum(0)
        c = rng.beta(cg[0] + corr, cg[1] + (tot - corr))
        if it >= burn:
            A[it - burn] = a; D[it - burn] = d; TH[it - burn] = theta; C[it - burn] = c
    return dict(a=A, d=D, theta=TH, c=C)


# --------------------------------------------------------------------------- #
#  item characteristic curve + information functions                            #
# --------------------------------------------------------------------------- #

def icc(theta, a, d, c=0.0):
    """item characteristic curve P(correct | theta)."""
    P = Phi(a * theta + d)
    return c + (1 - c) * P

def item_information(theta, a, d, c=0.0):
    """Fisher item information I_j(theta) = (P')^2 / (P(1-P)); normal-ogive metric."""
    phi = np.exp(-0.5 * (a * theta + d) ** 2) / np.sqrt(2 * np.pi)
    P = icc(theta, a, d, c); Pp = (1 - c) * a * phi
    return Pp ** 2 / np.clip(P * (1 - P), 1e-9, None)

def test_information(theta, a, d, c=None):
    """total test information (sum over items) and the standard error 1/sqrt(TIF)."""
    a = np.asarray(a); d = np.asarray(d); c = np.zeros_like(a) if c is None else np.asarray(c)
    TIF = np.zeros_like(theta, float)
    for j in range(len(a)):
        TIF += item_information(theta, a[j], d[j], c[j])
    return TIF, 1.0 / np.sqrt(np.clip(TIF, 1e-9, None))