"""Metric tests against hand-computed values.

Eval code that is silently wrong is worse than no eval, so every metric is
checked against a case whose answer can be worked out on paper.
"""

from __future__ import annotations

import numpy as np
import pytest

from mss.evaluate import (
    average_precision,
    evaluate_retrieval,
    markdown_table,
    ndcg_at_k,
    random_baseline,
    rank_matrix,
    recall_at_k,
    reciprocal_rank,
)


def test_recall_counts_relevant_items_found():
    ranking = np.array([3, 1, 4, 0, 2])
    assert recall_at_k(ranking, {3}, 1) == 1.0
    assert recall_at_k(ranking, {4}, 1) == 0.0
    assert recall_at_k(ranking, {4}, 3) == 1.0
    # two relevant, one inside the cut -> half of them recalled
    assert recall_at_k(ranking, {3, 2}, 3) == 0.5


def test_reciprocal_rank_uses_the_first_hit():
    ranking = np.array([3, 1, 4, 0, 2])
    assert reciprocal_rank(ranking, {3}) == 1.0
    assert reciprocal_rank(ranking, {1}) == pytest.approx(1 / 2)
    assert reciprocal_rank(ranking, {2}) == pytest.approx(1 / 5)
    # first of several relevant items is what counts
    assert reciprocal_rank(ranking, {4, 2}) == pytest.approx(1 / 3)
    assert reciprocal_rank(ranking, set()) == 0.0


def test_average_precision_hand_computed():
    ranking = np.array([0, 1, 2, 3, 4])
    # relevant at positions 1 and 3 (1-based): precision 1/1 and 2/3
    assert average_precision(ranking, {0, 2}) == pytest.approx((1.0 + 2 / 3) / 2)
    # a single relevant item makes AP equal to the reciprocal rank
    assert average_precision(ranking, {2}) == pytest.approx(1 / 3)


def test_average_precision_penalises_unretrieved_relevants():
    # only the top 3 can ever be hit here; the 4th relevant item is missing
    ranking = np.array([0, 1, 2])
    ap = average_precision(ranking, {0, 1, 2, 99})
    assert ap == pytest.approx((1.0 + 1.0 + 1.0) / 4)


def test_ndcg_is_one_for_a_perfect_ranking():
    ranking = np.array([0, 1, 2, 3])
    assert ndcg_at_k(ranking, {0, 1}, 4) == pytest.approx(1.0)


def test_ndcg_hand_computed_for_an_imperfect_ranking():
    ranking = np.array([0, 1, 2, 3])
    # one relevant item at position 2 -> DCG = 1/log2(3), IDCG = 1/log2(2) = 1
    assert ndcg_at_k(ranking, {1}, 4) == pytest.approx(1 / np.log2(3))


def test_ndcg_ignores_relevants_beyond_the_cut():
    ranking = np.array([9, 8, 7, 0])
    assert ndcg_at_k(ranking, {0}, 3) == 0.0


def test_rank_matrix_orders_by_cosine():
    items = np.eye(4, dtype=np.float32)
    queries = np.array([[0, 0, 1, 0]], dtype=np.float32)
    assert rank_matrix(queries, items)[0][0] == 2


def test_rank_matrix_is_deterministic_under_ties():
    items = np.ones((5, 3), dtype=np.float32)
    queries = np.ones((1, 3), dtype=np.float32)
    a = rank_matrix(queries, items)
    b = rank_matrix(queries, items)
    np.testing.assert_array_equal(a, b)
    np.testing.assert_array_equal(a[0], np.arange(5))


def test_perfect_retrieval_scores_one():
    vecs = np.eye(6, dtype=np.float32)
    res = evaluate_retrieval(vecs, vecs, [{i} for i in range(6)])
    assert res.metrics["R@1"] == pytest.approx(1.0)
    assert res.metrics["MRR"] == pytest.approx(1.0)
    assert res.metrics["mAP"] == pytest.approx(1.0)
    assert res.metrics["median_rank"] == 1.0


def test_metrics_are_bounded_and_ordered():
    rng = np.random.default_rng(0)
    q = rng.standard_normal((30, 32)).astype(np.float32)
    it = rng.standard_normal((80, 32)).astype(np.float32)
    q /= np.linalg.norm(q, axis=1, keepdims=True)
    it /= np.linalg.norm(it, axis=1, keepdims=True)
    res = evaluate_retrieval(q, it, [{i} for i in range(30)])
    for key, val in res.metrics.items():
        if key != "median_rank":
            assert 0.0 <= val <= 1.0, key
    # recall cannot shrink as the cut widens
    assert res.metrics["R@1"] <= res.metrics["R@5"] <= res.metrics["R@10"]


def test_random_baseline_lands_near_theory():
    n_items = 1000
    res = random_baseline(n_items, [{i} for i in range(400)], seed=1)
    # with one relevant item among n, expected R@k is k/n
    assert res.metrics["R@10"] == pytest.approx(10 / n_items, abs=0.01)
    assert res.metrics["median_rank"] == n_items / 2


def test_markdown_table_has_a_row_per_result():
    rng = np.random.default_rng(2)
    v = rng.standard_normal((5, 8)).astype(np.float32)
    v /= np.linalg.norm(v, axis=1, keepdims=True)
    r1 = evaluate_retrieval(v, v, [{i} for i in range(5)], name="a")
    r2 = random_baseline(5, [{i} for i in range(5)])
    table = markdown_table([r1, r2])
    assert table.count("\n") == 3  # header + separator + 2 rows
    assert "| a |" in table