import math

import pytest

from leanbench.evalstats import auc, cohen_kappa, confusion, wilson


def test_wilson_matches_known_values():
    lo, hi = wilson(5, 10)
    assert math.isclose(lo, 0.2366, abs_tol=1e-3) and math.isclose(hi, 0.7634, abs_tol=1e-3)
    assert wilson(0, 0)[0] == 0.0 and wilson(0, 0)[1] == 1.0
    assert wilson(10, 10)[1] == 1.0


def test_confusion_counts_and_rates():
    c = confusion([True, True, False, False, True], [True, False, True, False, True])
    assert (c.tp, c.fp, c.fn, c.tn) == (2, 1, 1, 1)
    assert math.isclose(c.recall, 2 / 3) and math.isclose(c.precision, 2 / 3)
    assert math.isclose(c.false_flag_rate, 0.5)
    assert c.as_dict()["n"] == 5


def test_confusion_length_mismatch():
    with pytest.raises(ValueError):
        confusion([True], [True, False])


def test_empty_classes_give_nan_not_a_crash():
    c = confusion([False, False], [False, False])
    assert math.isnan(c.recall) and math.isnan(c.precision)


def test_auc():
    assert auc([0.9, 0.8, 0.2, 0.1], [True, True, False, False]) == 1.0
    assert auc([0.1, 0.2, 0.8, 0.9], [True, True, False, False]) == 0.0
    assert auc([0.5, 0.5], [True, False]) == 0.5
    assert math.isnan(auc([0.5], [True]))


def test_kappa():
    assert cohen_kappa(["a", "b", "a"], ["a", "b", "a"]) == 1.0
    assert math.isclose(cohen_kappa(["a", "a", "b", "b"], ["a", "b", "a", "b"]), 0.0)
    with pytest.raises(ValueError):
        cohen_kappa([], [])