"""Two-parameter logistic IRT (marginal maximum likelihood, EM with Gauss–Hermite quadrature).

Items are binary (0/1); missing entries contribute nothing to the likelihood (missing-at-random by design, no
imputation). Returns item discriminations a_j, difficulties b_j, EAP scores and posterior SDs for each unit.
Deterministic (no randomness). Written for transparency rather than speed (N ≈ 600 × J ≈ 30 is instant).
"""
from __future__ import annotations
import numpy as np
from scipy.special import expit


def fit_2pl(X: np.ndarray, n_quad: int = 41, max_iter: int = 1000, tol: float = 1e-6, min_items: int = 3,
            prior_sd_log_a: float = 0.5, prior_sd_b: float = 2.0):
    """MAP variant: log-normal prior on a (mean 0, sd prior_sd_log_a) and normal prior on b (sd prior_sd_b) to keep
    items that are nearly constant in the sample from running to the boundary (Heywood-type cases)."""
    X = np.asarray(X, dtype=float)
    N, J = X.shape
    obs = ~np.isnan(X)
    Xz = np.where(obs, X, 0.0)
    nodes, weights = np.polynomial.hermite_e.hermegauss(n_quad)  # standard normal prior
    weights = weights / weights.sum()
    a = np.ones(J); b = np.zeros(J)
    ll_old = -np.inf
    for it in range(max_iter):
        # E-step: posterior over nodes for each unit
        P = expit(a[None, :] * (nodes[:, None] - b[None, :]))          # Q x J
        logP, log1P = np.log(P + 1e-300), np.log(1 - P + 1e-300)
        # log-lik of unit i at node q = sum_j obs_ij [x_ij logP_qj + (1-x_ij) log(1-P_qj)]
        L = Xz @ logP.T + ((1 - Xz) * obs) @ log1P.T                   # N x Q
        L += np.log(weights)[None, :]
        m = L.max(axis=1, keepdims=True)
        post = np.exp(L - m); norm = post.sum(axis=1, keepdims=True); post /= norm
        ll = float((m.squeeze() + np.log(norm.squeeze())).sum())
        # M-step: Newton updates per item on expected counts
        r = post.T @ Xz                                                 # Q x J expected correct
        n = post.T @ obs.astype(float)                                  # Q x J expected tries
        for j in range(J):
            for _ in range(10):
                p = expit(a[j] * (nodes - b[j]))
                g_a = np.sum((r[:, j] - n[:, j] * p) * (nodes - b[j])) - (np.log(a[j]) / prior_sd_log_a ** 2) / a[j]
                g_b = np.sum((r[:, j] - n[:, j] * p) * (-a[j])) - b[j] / prior_sd_b ** 2
                w = n[:, j] * p * (1 - p)
                H_aa = -np.sum(w * (nodes - b[j]) ** 2) - (1 - np.log(a[j])) / (prior_sd_log_a ** 2 * a[j] ** 2)
                H_bb = -np.sum(w * a[j] ** 2) - 1 / prior_sd_b ** 2
                H_ab = np.sum(w * a[j] * (nodes - b[j]))
                H = np.array([[H_aa, H_ab], [H_ab, H_bb]]) - 1e-6 * np.eye(2)
                step = np.linalg.solve(H, np.array([g_a, g_b]))
                a[j] -= step[0]; b[j] -= step[1]
                a[j] = float(np.clip(a[j], 0.05, 8.0)); b[j] = float(np.clip(b[j], -6, 6))
                if np.abs(step).max() < 1e-6:
                    break
        if abs(ll - ll_old) < tol:
            break
        ll_old = ll
    # final EAP
    P = expit(a[None, :] * (nodes[:, None] - b[None, :]))
    L = Xz @ np.log(P + 1e-300).T + ((1 - Xz) * obs) @ np.log(1 - P + 1e-300).T + np.log(weights)[None, :]
    m = L.max(axis=1, keepdims=True); post = np.exp(L - m); post /= post.sum(axis=1, keepdims=True)
    eap = post @ nodes
    psd = np.sqrt(np.maximum(post @ nodes ** 2 - eap ** 2, 0))
    n_items = obs.sum(axis=1)
    eap = np.where(n_items >= min_items, eap, np.nan); psd = np.where(n_items >= min_items, psd, np.nan)
    return {"a": a, "b": b, "eap": eap, "psd": psd, "loglik": ll, "iterations": it + 1, "n_items": n_items}