# llmindex.io — psychometrics tests: parameter recovery on synthetic data
# Author:  Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# License: Proprietary — © Simon-Pierre Boucher, all rights reserved

import numpy as np
import pytest

from llmindex_psycho.bt import fit_bradley_terry
from llmindex_psycho.calibration import brier_score, expected_calibration_error
from llmindex_psycho.irt import fit_2pl


def _synthetic_2pl(M=12, I=80, seed=7):
    rng = np.random.default_rng(seed)
    theta = rng.normal(0, 1, M)
    a = np.exp(rng.normal(0, 0.4, I))
    b = rng.normal(0, 1.2, I)
    P = 1 / (1 + np.exp(-a[None, :] * (theta[:, None] - b[None, :])))
    X = (rng.random((M, I)) < P).astype(float)
    return theta, a, b, X


def test_2pl_recovers_ability_ordering():
    theta_true, _, _, X = _synthetic_2pl()
    result = fit_2pl(X, max_iterations=800)
    corr = np.corrcoef(theta_true, result.theta)[0, 1]
    assert corr > 0.85, f"theta recovery correlation too low: {corr:.3f}"
    assert np.all(result.theta_se > 0)
    assert np.all(result.a > 0)


def test_2pl_handles_missing_cells():
    _, _, _, X = _synthetic_2pl()
    Xm = X.copy()
    rng = np.random.default_rng(1)
    Xm[rng.random(X.shape) < 0.3] = np.nan
    result = fit_2pl(Xm, max_iterations=600)
    assert np.isfinite(result.theta).all()
    assert np.isfinite(result.theta_se).all()


def test_2pl_se_shrinks_with_more_items():
    _, _, _, X = _synthetic_2pl(M=8, I=120, seed=3)
    few = fit_2pl(X[:, :15], max_iterations=600)
    many = fit_2pl(X, max_iterations=600)
    assert many.theta_se.mean() < few.theta_se.mean()


def test_2pl_rejects_empty():
    with pytest.raises(ValueError):
        fit_2pl(np.full((3, 3), np.nan))


def test_bradley_terry_recovers_ordering():
    rng = np.random.default_rng(11)
    strength = np.array([2.0, 1.0, 0.0, -1.0, -2.0])
    M = len(strength)
    wins = np.zeros((M, M))
    for i in range(M):
        for j in range(M):
            if i == j:
                continue
            p = 1 / (1 + np.exp(-(strength[i] - strength[j])))
            wins[i, j] = rng.binomial(40, p)
    result = fit_bradley_terry(wins)
    assert result.converged
    assert list(np.argsort(-result.log_strength)) == [0, 1, 2, 3, 4]


def test_calibration_metrics():
    conf = np.array([0.9, 0.9, 0.1, 0.1])
    correct = np.array([1.0, 1.0, 0.0, 0.0])
    assert brier_score(conf, correct) == pytest.approx(0.01)
    assert expected_calibration_error(conf, correct) == pytest.approx(0.1)
    overconfident = np.full(100, 0.99)
    outcomes = np.zeros(100)
    assert expected_calibration_error(overconfident, outcomes) > 0.9