"""
tests/test_locomo_metrics.py — Unit tests for LoCoMo benchmark metrics.

Covers:
  - stemmed_f1_score
  - multi_hop_f1
  - adversarial_accuracy
  - retrieval_recall
  - score_by_category / CATEGORY_METRICS dispatcher
"""
import pytest

from cortex.eval.metrics import (
    adversarial_accuracy,
    multi_hop_f1,
    retrieval_recall,
    score_by_category,
    stemmed_f1_score,
    CATEGORY_METRICS,
)


# ---------------------------------------------------------------------------
# stemmed_f1_score
# ---------------------------------------------------------------------------

class TestStemmedF1Score:
    def test_case_insensitive(self):
        """'Amsterdam' vs 'amsterdam' should be 1.0."""
        assert stemmed_f1_score("Amsterdam", "amsterdam") == pytest.approx(1.0)

    def test_stemming_fox(self):
        """'the quick brown fox' vs 'quick brown foxes' — stemming should produce high score."""
        score = stemmed_f1_score("the quick brown fox", "quick brown foxes")
        assert score >= 0.7, f"Expected high score with stemming, got {score}"

    def test_wrong_answer_low_score(self):
        """'completely wrong answer' vs 'right answer' should give low score."""
        score = stemmed_f1_score("completely wrong answer", "right answer")
        assert score < 0.5, f"Expected low score for wrong answer, got {score}"

    def test_empty_prediction(self):
        """Empty prediction → 0.0."""
        assert stemmed_f1_score("", "some reference") == 0.0

    def test_empty_reference(self):
        """Empty reference → 0.0."""
        assert stemmed_f1_score("some prediction", "") == 0.0

    def test_both_empty(self):
        """Both empty → 0.0."""
        assert stemmed_f1_score("", "") == 0.0

    def test_exact_match(self):
        """Exact same string → 1.0."""
        assert stemmed_f1_score("the cat sat", "the cat sat") == pytest.approx(1.0)

    def test_articles_removed(self):
        """Articles 'a', 'an', 'the' should be stripped before scoring."""
        assert stemmed_f1_score("a quick fox", "the quick fox") == pytest.approx(1.0)

    def test_partial_overlap(self):
        """Partial word overlap returns score between 0 and 1."""
        score = stemmed_f1_score("cat sat on mat", "cat sat")
        assert 0.0 < score < 1.0

    def test_no_overlap(self):
        """Zero word overlap returns 0.0."""
        assert stemmed_f1_score("apple orange", "car truck") == 0.0


# ---------------------------------------------------------------------------
# multi_hop_f1
# ---------------------------------------------------------------------------

class TestMultiHopF1:
    def test_two_sub_answers_match(self):
        """'Alice and Bob' vs 'Alice, Bob' → high score."""
        score = multi_hop_f1("Alice and Bob", "Alice, Bob")
        assert score >= 0.5, f"Expected high score, got {score}"

    def test_one_of_three_sub_answers(self):
        """'Alice' vs 'Alice, Bob, Carol' → ~0.33 (only 1/3 sub-answers match)."""
        score = multi_hop_f1("Alice", "Alice, Bob, Carol")
        assert score == pytest.approx(1.0 / 3.0, abs=0.05), f"Expected ~0.33, got {score}"

    def test_all_three_match(self):
        """'Alice, Bob, Carol' vs 'Alice, Bob, Carol'.

        Each single-token sub-answer scored against 3-token prediction:
          precision=1/3, recall=1.0 → F1=0.5 per sub-answer → mean=0.5
        """
        score = multi_hop_f1("Alice, Bob, Carol", "Alice, Bob, Carol")
        assert score == pytest.approx(0.5, abs=0.05)

    def test_empty_reference(self):
        """Empty reference → 0.0."""
        assert multi_hop_f1("Alice", "") == 0.0

    def test_empty_prediction(self):
        """Empty prediction → 0.0."""
        score = multi_hop_f1("", "Alice, Bob")
        assert score == 0.0

    def test_single_sub_answer(self):
        """Single sub-answer — equivalent to stemmed_f1_score."""
        score_multi = multi_hop_f1("Paris", "Paris")
        score_direct = stemmed_f1_score("Paris", "Paris")
        assert score_multi == pytest.approx(score_direct)


# ---------------------------------------------------------------------------
# adversarial_accuracy
# ---------------------------------------------------------------------------

class TestAdversarialAccuracy:
    def test_no_information_available(self):
        assert adversarial_accuracy("There is no information available about this") == 1.0

    def test_not_mentioned(self):
        assert adversarial_accuracy("The answer is not mentioned in the conversation") == 1.0

    def test_not_available(self):
        assert adversarial_accuracy("This data is not available") == 1.0

    def test_cannot_be_determined(self):
        assert adversarial_accuracy("That cannot be determined from context") == 1.0

    def test_no_evidence(self):
        assert adversarial_accuracy("There is no evidence for this claim") == 1.0

    def test_positive_answer(self):
        assert adversarial_accuracy("Paris is the capital of France") == 0.0

    def test_empty_string(self):
        assert adversarial_accuracy("") == 0.0

    def test_case_insensitive(self):
        assert adversarial_accuracy("NO INFORMATION found") == 1.0

    def test_unrelated_negative_sentence(self):
        assert adversarial_accuracy("I don't know the exact date") == 0.0


# ---------------------------------------------------------------------------
# retrieval_recall
# ---------------------------------------------------------------------------

class TestRetrievalRecall:
    def _items(self, *ids: str):
        return [{"dia_id": d} for d in ids]

    def test_perfect_recall(self):
        items = self._items("D1:3", "D5:12")
        assert retrieval_recall(items, ["D1:3", "D5:12"]) == pytest.approx(1.0)

    def test_partial_recall(self):
        items = self._items("D1:3")
        assert retrieval_recall(items, ["D1:3", "D5:12"]) == pytest.approx(0.5)

    def test_zero_recall(self):
        items = self._items("D9:99")
        assert retrieval_recall(items, ["D1:3"]) == pytest.approx(0.0)

    def test_empty_evidence(self):
        """Empty evidence → 1.0 (vacuously true)."""
        items = self._items("D1:3")
        assert retrieval_recall(items, []) == pytest.approx(1.0)

    def test_empty_retrieved(self):
        assert retrieval_recall([], ["D1:3", "D5:12"]) == pytest.approx(0.0)

    def test_custom_id_field(self):
        items = [{"msg_id": "M1"}, {"msg_id": "M2"}]
        assert retrieval_recall(items, ["M1", "M2"], id_field="msg_id") == pytest.approx(1.0)

    def test_items_without_id_field(self):
        items = [{"other_field": "D1:3"}]
        assert retrieval_recall(items, ["D1:3"]) == pytest.approx(0.0)

    def test_superset_retrieved(self):
        """Retrieved superset → recall still 1.0."""
        items = self._items("D1:3", "D5:12", "D9:99", "D3:7")
        assert retrieval_recall(items, ["D1:3", "D5:12"]) == pytest.approx(1.0)


# ---------------------------------------------------------------------------
# score_by_category / CATEGORY_METRICS dispatcher
# ---------------------------------------------------------------------------

class TestScoreByCategory:
    def test_category_1_routes_to_multi_hop_f1(self):
        assert CATEGORY_METRICS[1] is multi_hop_f1
        score_direct = multi_hop_f1("Alice and Bob", "Alice, Bob")
        score_dispatch = score_by_category("Alice and Bob", "Alice, Bob", 1)
        assert score_dispatch == pytest.approx(score_direct)

    def test_category_2_routes_to_stemmed_f1(self):
        assert CATEGORY_METRICS[2] is stemmed_f1_score
        score_direct = stemmed_f1_score("Amsterdam", "amsterdam")
        score_dispatch = score_by_category("Amsterdam", "amsterdam", 2)
        assert score_dispatch == pytest.approx(score_direct)

    def test_category_3_routes_to_stemmed_f1(self):
        assert CATEGORY_METRICS[3] is stemmed_f1_score
        score_direct = stemmed_f1_score("January 2020", "january 2020")
        score_dispatch = score_by_category("January 2020", "january 2020", 3)
        assert score_dispatch == pytest.approx(score_direct)

    def test_category_4_routes_to_stemmed_f1(self):
        assert CATEGORY_METRICS[4] is stemmed_f1_score
        score_direct = stemmed_f1_score("the quick fox", "quick fox")
        score_dispatch = score_by_category("the quick fox", "quick fox", 4)
        assert score_dispatch == pytest.approx(score_direct)

    def test_category_5_routes_to_adversarial(self):
        assert CATEGORY_METRICS[5] is adversarial_accuracy
        score = score_by_category("no information available", "some_reference", 5)
        assert score == pytest.approx(1.0)

    def test_category_5_ignores_reference(self):
        score_a = score_by_category("no information", "ref_a", 5)
        score_b = score_by_category("no information", "totally_different_ref", 5)
        assert score_a == score_b

    def test_category_5_wrong_prediction(self):
        score = score_by_category("Paris is the capital", "doesn't matter", 5)
        assert score == 0.0

    def test_invalid_category_raises(self):
        with pytest.raises(KeyError):
            score_by_category("pred", "ref", 99)

    def test_all_categories_covered(self):
        assert set(CATEGORY_METRICS.keys()) == {1, 2, 3, 4, 5}
