"""
grm.py -- POLYTOMOUS IRT: the Graded Response Model (from scratch).

Backs the notebooks in  "Polytomous IRT -- the Graded Response Model".

When test items are ORDERED categories (Likert: strongly disagree ... strongly agree), the
2PL is not enough -- we need a curve for each category. Samejima's GRADED RESPONSE MODEL (1969)
is the ordered-probit item response model: a person with ability theta_i produces a latent
propensity z_ij = a_j theta_i + noise, and the observed category is the interval that z falls in,

    x_ij = k   iff   gamma_{j,k-1} < z_ij <= gamma_{j,k},     z_ij ~ N(a_j theta_i, 1),

with item DISCRIMINATION a_j and ordered category thresholds gamma_{j,1} < ... < gamma_{j,K-1}
(and gamma_{j,0} = -inf, gamma_{j,K} = +inf). Equivalently the cumulative curve is
P(x_ij >= k) = Phi(a_j theta_i - gamma_{j,k-1}); differencing adjacent cumulatives gives the
CATEGORY RESPONSE CURVES. This is the polytomous generalisation of the 2PL (K=2 recovers it).

The from-scratch sampler is the ORDERED-PROBIT data augmentation (Albert & Chib 1993) -- the same
truncated-normal trick as the dichotomous IRT and the sequential/ordinal-probit models: draw the
latent z_ij truncated to the observed category's interval, after which theta and a are conjugate
Gaussian and the thresholds gamma have simple order-constrained conditionals. Scale and location
are fixed by theta ~ N(0,1), the discrimination sign by keeping sum_j a_j > 0. The Rasch-family
alternative (Partial Credit / Generalised Partial Credit) uses adjacent-category rather than
cumulative logits -- fitted with the packages in the R notebook.
"""

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


def _rtruncnorm(mean, lo, hi, rng):
    """draw N(mean, 1) truncated to (lo, hi], elementwise."""
    a = Phi(lo - mean); b = Phi(hi - mean)
    u = a + rng.random(mean.shape) * (b - a)
    return mean + Phinv(np.clip(u, 1e-12, 1 - 1e-12))


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

def simulate_grm(n, a, gamma, rng, theta=None):
    """simulate ordered responses (categories 1..K). a:(J,), gamma:(J,K-1) interior thresholds."""
    a = np.asarray(a, float); gamma = np.asarray(gamma, float); J = len(a)
    if theta is None:
        theta = rng.standard_normal(n)
    z = np.outer(theta, a) + rng.standard_normal((n, J))
    X = np.ones((n, J), int)
    for j in range(J):
        for c in gamma[j]:
            X[:, j] += (z[:, j] > c).astype(int)
    return X, theta


# --------------------------------------------------------------------------- #
#  Gibbs sampler (ordered-probit augmentation)                                  #
# --------------------------------------------------------------------------- #

def grm_gibbs(X, rng, draws=3000, burn=1500, a_sd=3.0):
    """graded response model via ordered-probit augmentation. X:(N,J) integer categories 1..K.
    Returns posterior draws of discriminations a:(draws,J), interior thresholds
    gamma:(draws,J,K-1) and abilities theta:(draws,N)."""
    X = np.asarray(X, int); N, J = X.shape; K = int(X.max())
    theta = rng.standard_normal(N); a = np.ones(J)
    gam = np.tile(np.linspace(-1.0, 1.0, K - 1), (J, 1)).astype(float)   # (J, K-1) interior
    A = np.empty((draws, J)); G = np.empty((draws, J, K - 1)); TH = np.empty((draws, N))
    for it in range(draws + burn):
        # full cutpoint arrays with +-inf sentinels: (J, K+1)
        full = np.concatenate([np.full((J, 1), -1e6), gam, np.full((J, 1), 1e6)], axis=1)
        lo = full[np.arange(J)[None, :], X - 1]
        hi = full[np.arange(J)[None, :], X]
        mu = np.outer(theta, a)
        z = _rtruncnorm(mu, lo, hi, rng)
        # theta | z, a  (N(0,1) prior)
        prec = 1.0 + np.sum(a ** 2)
        theta = (z @ a) / prec + rng.standard_normal(N) / np.sqrt(prec)
        # a_j | z, theta  (regression through the origin, positive)
        sa = 1.0 / a_sd ** 2
        prec_a = sa + np.sum(theta ** 2)
        a = (theta @ z) / prec_a + rng.standard_normal(J) / np.sqrt(prec_a)
        if a.sum() < 0:
            a = -a; theta = -theta
        # thresholds gamma_{j,c} | z, X  (order-constrained uniform)
        for j in range(J):
            zj = z[:, j]; xj = X[:, j]
            for c in range(1, K):
                L = gam[j, c - 2] if c >= 2 else -1e6
                U = gam[j, c] if c <= K - 2 else 1e6
                below = zj[xj == c]; above = zj[xj == c + 1]
                if below.size: L = max(L, below.max())
                if above.size: U = min(U, above.min())
                if U > L:
                    gam[j, c - 1] = rng.uniform(L, U)
        if it >= burn:
            A[it - burn] = a; G[it - burn] = gam; TH[it - burn] = theta
    return dict(a=A, gamma=G, theta=TH, K=K)


# --------------------------------------------------------------------------- #
#  category response curves + information                                       #
# --------------------------------------------------------------------------- #

def category_probs(theta, a, gamma):
    """P(x = k | theta) for k = 1..K.  gamma: (K-1,) interior thresholds. Returns (len(theta), K)."""
    gamma = np.asarray(gamma, float); K = len(gamma) + 1
    cuts = np.concatenate([[-1e6], gamma, [1e6]])
    cum = Phi(a * theta[:, None] - cuts[None, :])        # P(z > cut) = P(x >= k+1)
    return cum[:, :-1] - cum[:, 1:]                       # differences give P(x=k)

def item_information(theta, a, gamma):
    """graded-model item information I_j(theta) = sum_k (dP_k/dtheta)^2 / P_k."""
    gamma = np.asarray(gamma, float); cuts = np.concatenate([[-1e6], gamma, [1e6]])
    z = a * theta[:, None] - cuts[None, :]
    phi = np.exp(-0.5 * z ** 2) / np.sqrt(2 * np.pi)
    dcum = a * phi                                        # d/dtheta P(z>cut)
    P = category_probs(theta, a, gamma)
    dP = dcum[:, :-1] - dcum[:, 1:]
    return np.sum(dP ** 2 / np.clip(P, 1e-9, None), axis=1)

def test_information(theta, a, gamma):
    a = np.asarray(a); TIF = np.zeros_like(theta, float)
    for j in range(len(a)):
        TIF += item_information(theta, a[j], gamma[j])
    return TIF