"""
dif.py -- DIFFERENTIAL ITEM FUNCTIONING and explanatory IRT (from scratch).

Backs the notebooks in  "Differential Item Functioning & Explanatory IRT".

Two questions about groups in a test:

  * IMPACT (latent regression): do the groups differ in the TRAIT itself? (Are men more
    verbally aggressive than women?) This is a legitimate group difference, modelled by letting
    the ability distribution depend on a covariate:  theta_i ~ N(beta * g_i, 1).

  * DIFFERENTIAL ITEM FUNCTIONING (DIF): does an ITEM behave differently across groups AT THE
    SAME trait level? An item with DIF is biased -- two equally-aggressive people of different
    groups endorse it at different rates. Uniform DIF is a group-specific shift in item difficulty:

        P(x_ij = 1 | theta_i) = Phi( a_j theta_i + d_j - delta_j * g_i ),

    where delta_j is item j's DIF (how much harder it is for group g=1 at the same theta).

Separating impact (beta) from bias (delta) is the whole point: a group can be genuinely more
aggressive AND some items can still be biased. From scratch it is the ALBERT-CHIB augmentation
again: draw z_ij truncated by the sign of x_ij, after which theta, the item parameters (a_j,d_j),
the DIF shifts delta_j and the regression coefficient beta all have conjugate Gaussian conditionals.
Ability scale/location are fixed by theta ~ N(beta g, 1); the DIF parameters are identified by
centring them to sum to zero (DIF is relative to the average item). Items whose delta_j credible
interval excludes zero are flagged as functioning differentially -- the psychometric fairness check.
Explanatory IRT (the LLTM idea) then regresses the item difficulties on ITEM FEATURES to explain
what makes items hard.
"""

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


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))


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

def simulate_dif(N, a, d, delta, beta, rng, g=None):
    """a,d,delta:(J,); beta scalar (group trait difference); g:(N,) group indicator 0/1."""
    a = np.asarray(a, float); d = np.asarray(d, float); delta = np.asarray(delta, float); J = len(a)
    if g is None:
        g = (rng.random(N) < 0.5).astype(float)
    theta = beta * g + rng.standard_normal(N)
    eta = np.outer(theta, a) + d[None, :] - delta[None, :] * g[:, None]
    X = (Phi(eta) > rng.random((N, J))).astype(float)
    return X, theta, g


# --------------------------------------------------------------------------- #
#  Gibbs sampler                                                                #
# --------------------------------------------------------------------------- #

def dif_gibbs(X, g, rng, draws=4000, burn=2000, a_sd=3.0, dif_sd=1.0, beta_sd=3.0):
    """2PL with uniform DIF (per-item shift by group) and a latent-regression group effect.
    g:(N,) group indicator 0/1 (the focal group = 1). Returns posterior draws of discriminations
    a, easinesses d, DIF shifts delta:(draws,J), the group trait difference beta, and abilities."""
    X = np.asarray(X, float); g = np.asarray(g, float); N, J = X.shape
    n1 = g.sum()
    theta = rng.standard_normal(N); a = np.ones(J); d = np.zeros(J); delta = np.zeros(J); beta = 0.0
    A = np.empty((draws, J)); D = np.empty((draws, J)); DE = np.empty((draws, J)); B = np.empty(draws); TH = np.empty((draws, N))
    for it in range(draws + burn):
        mu = np.outer(theta, a) + d[None, :] - delta[None, :] * g[:, None]
        z = _rtrunc_sign(mu, X > 0.5, rng)
        # theta_i | z, a, d, delta, beta   (prior N(beta g_i, 1))
        prec = 1.0 + np.sum(a ** 2)
        adjr = (z - d[None, :] + delta[None, :] * g[:, None]) @ a       # sum_j a_j (z - d + delta g)
        theta = (adjr + beta * g) / prec + rng.standard_normal(N) / np.sqrt(prec)
        # (a_j, d_j) | z, theta   (regress z + delta g on [theta, 1])
        zc = z + delta[None, :] * g[:, None]
        Xd = np.column_stack([theta, np.ones(N)]); Pp = np.diag([1.0 / a_sd ** 2, 1e-6])
        V = np.linalg.inv(Xd.T @ Xd + Pp); M = V @ (Xd.T @ zc); 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; beta = -beta
        # delta_j | z, theta, a, d   (only focal-group residuals inform it)
        r = z - np.outer(theta, a) - d[None, :]                        # residual = -delta_j g_i + eps
        r1sum = (r * g[:, None]).sum(0)                                # sum over g=1
        precd = n1 + 1.0 / dif_sd ** 2
        delta = -r1sum / precd + rng.standard_normal(J) / np.sqrt(precd)
        delta -= delta.mean()                                          # identify: DIF relative to average item
        # beta | theta, g   (regress theta on g)
        precb = g @ g + 1.0 / beta_sd ** 2
        beta = (theta @ g) / precb + rng.standard_normal() / np.sqrt(precb)
        if it >= burn:
            i = it - burn; A[i] = a; D[i] = d; DE[i] = delta; B[i] = beta; TH[i] = theta
    return dict(a=A, d=D, delta=DE, beta=B, theta=TH)