# llmindex.io — custom 2PL IRT fit (numpy MAP via Adam, missing-aware)
# Author:  Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# License: Proprietary — © Simon-Pierre Boucher, all rights reserved
#
# 2PL: P(correct) = sigmoid(a_i * (theta_m - b_i))
# MAP estimation with priors:
#   theta ~ N(0, 1)      b ~ N(0, 1.5)      log a ~ N(0, 0.5)
# Missing responses are masked. Theta SE comes from the Fisher information
# (plus prior precision). Numpy-based (no torch) so the fit runs anywhere.

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np


def _sigmoid(x: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-np.clip(x, -30, 30)))


@dataclass
class Fit2PLResult:
    theta: np.ndarray  # (M,)
    theta_se: np.ndarray  # (M,)
    a: np.ndarray  # (I,) discrimination
    b: np.ndarray  # (I,) difficulty (logits)
    iterations: int
    converged: bool
    final_loglik: float
    diagnostics: dict = field(default_factory=dict)


def fit_2pl(
    responses: np.ndarray,
    max_iterations: int = 500,
    tolerance: float = 1e-6,
    lr: float = 0.05,
    prior_theta_sd: float = 1.0,
    prior_b_sd: float = 1.5,
    prior_log_a_sd: float = 0.5,
    seed: int = 0,
) -> Fit2PLResult:
    """Fit a 2PL model on an (M models × I items) matrix of {0,1,nan}."""
    X = np.asarray(responses, dtype=float)
    if X.ndim != 2:
        raise ValueError("responses must be a 2D matrix (models × items)")
    M, I = X.shape
    mask = ~np.isnan(X)
    if mask.sum() == 0:
        raise ValueError("no observed responses")
    Xf = np.nan_to_num(X, nan=0.0)

    rng = np.random.default_rng(seed)
    # Warm start: theta from row accuracy, b from item difficulty (logit of failure rate).
    row_acc = np.where(mask.sum(1) > 0, Xf.sum(1) / np.maximum(mask.sum(1), 1), 0.5)
    col_acc = np.where(mask.sum(0) > 0, Xf.sum(0) / np.maximum(mask.sum(0), 1), 0.5)
    theta = np.clip(np.log(row_acc + 1e-3) - np.log(1 - row_acc + 1e-3), -2, 2)
    theta = theta - theta.mean()
    b = np.clip(-(np.log(col_acc + 1e-3) - np.log(1 - col_acc + 1e-3)), -2.5, 2.5)
    log_a = rng.normal(0.0, 0.01, size=I)

    # Adam state
    params = [theta, b, log_a]
    m_state = [np.zeros_like(p) for p in params]
    v_state = [np.zeros_like(p) for p in params]
    beta1, beta2, eps = 0.9, 0.999, 1e-8

    prev_obj = -np.inf
    converged = False
    it = 0
    for it in range(1, max_iterations + 1):
        a = np.exp(log_a)
        Z = a[None, :] * (theta[:, None] - b[None, :])
        P = _sigmoid(Z)
        R = np.where(mask, Xf - P, 0.0)  # residuals on observed cells

        g_theta = (R * a[None, :]).sum(1) - theta / prior_theta_sd**2
        g_b = (-(R * a[None, :])).sum(0) - b / prior_b_sd**2
        g_a = (R * (theta[:, None] - b[None, :])).sum(0)
        g_log_a = g_a * a - log_a / prior_log_a_sd**2

        grads = [g_theta, g_b, g_log_a]
        for j, (p, g) in enumerate(zip(params, grads)):
            m_state[j] = beta1 * m_state[j] + (1 - beta1) * g
            v_state[j] = beta2 * v_state[j] + (1 - beta2) * g * g
            mhat = m_state[j] / (1 - beta1**it)
            vhat = v_state[j] / (1 - beta2**it)
            p += lr * mhat / (np.sqrt(vhat) + eps)

        # Identification: center abilities each step (scale is pinned by priors).
        shift = theta.mean()
        theta -= shift
        b -= shift

        with np.errstate(divide="ignore", invalid="ignore"):
            ll = np.where(mask, Xf * np.log(P + 1e-12) + (1 - Xf) * np.log(1 - P + 1e-12), 0.0).sum()
        obj = (
            ll
            - 0.5 * (theta**2).sum() / prior_theta_sd**2
            - 0.5 * (b**2).sum() / prior_b_sd**2
            - 0.5 * (log_a**2).sum() / prior_log_a_sd**2
        )
        if abs(obj - prev_obj) < tolerance * (1 + abs(prev_obj)):
            converged = True
            break
        prev_obj = obj

    a = np.exp(log_a)
    Z = a[None, :] * (theta[:, None] - b[None, :])
    P = _sigmoid(Z)
    info = (mask * (a[None, :] ** 2) * P * (1 - P)).sum(1) + 1.0 / prior_theta_sd**2
    theta_se = 1.0 / np.sqrt(info)
    final_ll = float(
        np.where(mask, Xf * np.log(P + 1e-12) + (1 - Xf) * np.log(1 - P + 1e-12), 0.0).sum()
    )

    return Fit2PLResult(
        theta=theta,
        theta_se=theta_se,
        a=a,
        b=b,
        iterations=it,
        converged=converged,
        final_loglik=final_ll,
        diagnostics={
            "observed_cells": int(mask.sum()),
            "models": int(M),
            "items": int(I),
        },
    )