"""Tests for rbsp.core.bm25 -- BM25Scorer and compute_term_freqs."""

import math
import os
import sys
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
from rbsp.core.bm25 import BM25Scorer, compute_term_freqs


class TestComputeTermFreqs(unittest.TestCase):
    """Tests for the compute_term_freqs helper."""

    def test_compute_term_freqs_basic(self):
        result = compute_term_freqs(["hello", "world", "hello"])
        self.assertEqual(result, {"hello": 2, "world": 1})

    def test_compute_term_freqs_empty(self):
        result = compute_term_freqs([])
        self.assertEqual(result, {})

    def test_compute_term_freqs_single_token(self):
        result = compute_term_freqs(["only"])
        self.assertEqual(result, {"only": 1})

    def test_compute_term_freqs_all_same(self):
        result = compute_term_freqs(["a", "a", "a"])
        self.assertEqual(result, {"a": 3})


class TestBM25IDF(unittest.TestCase):
    """Tests for IDF computation."""

    def _make_scorer(self, num_docs, doc_freqs):
        scorer = BM25Scorer()
        avg_dl = 10.0
        scorer.update_stats(num_docs, avg_dl, doc_freqs)
        return scorer

    def test_idf_rare_term_higher(self):
        # "rare" appears in 1 doc, "common" appears in 9 docs, corpus has 10
        scorer = self._make_scorer(10, {"rare": 1, "common": 9})
        self.assertGreater(scorer.idf("rare"), scorer.idf("common"))

    def test_idf_absent_term(self):
        # Term not in corpus: df=0, IDF = ln((N - 0 + 0.5)/(0 + 0.5) + 1)
        scorer = self._make_scorer(10, {"existing": 5})
        expected = math.log((10 - 0 + 0.5) / (0 + 0.5) + 1.0)
        self.assertAlmostEqual(scorer.idf("absent"), expected)

    def test_idf_non_negative(self):
        # Even a term in every doc should have non-negative IDF with +1 smoothing
        scorer = self._make_scorer(10, {"everywhere": 10})
        self.assertGreaterEqual(scorer.idf("everywhere"), 0.0)

    def test_idf_increases_with_rarity(self):
        scorer = self._make_scorer(100, {"df1": 1, "df10": 10, "df50": 50})
        self.assertGreater(scorer.idf("df1"), scorer.idf("df10"))
        self.assertGreater(scorer.idf("df10"), scorer.idf("df50"))


class TestBM25Score(unittest.TestCase):
    """Tests for score, term_score, and related methods."""

    def setUp(self):
        self.scorer = BM25Scorer(k1=1.2, b=0.75)
        self.scorer.update_stats(
            num_docs=100,
            avg_doc_len=50.0,
            doc_freqs={"python": 10, "java": 20, "code": 80, "rare": 1},
        )

    def test_score_zero_for_no_match(self):
        # Query terms not present in document
        score = self.scorer.score(
            ["python", "java"],
            {"unrelated": 5, "other": 3},
            doc_len=50,
        )
        self.assertEqual(score, 0.0)

    def test_score_positive_for_match(self):
        score = self.scorer.score(
            ["python"],
            {"python": 3, "other": 2},
            doc_len=50,
        )
        self.assertGreater(score, 0.0)

    def test_score_higher_for_more_matches(self):
        # Doc with both query terms should score higher than doc with only one
        score_one = self.scorer.score(
            ["python", "java"],
            {"python": 2},
            doc_len=50,
        )
        score_both = self.scorer.score(
            ["python", "java"],
            {"python": 2, "java": 2},
            doc_len=50,
        )
        self.assertGreater(score_both, score_one)

    def test_score_tf_saturation(self):
        # Increasing TF gives diminishing returns due to k1 saturation.
        # Going from tf=1 to tf=2 should give a bigger boost than tf=10 to tf=11.
        score_tf1 = self.scorer.term_score("python", tf=1, doc_len=50)
        score_tf2 = self.scorer.term_score("python", tf=2, doc_len=50)
        score_tf10 = self.scorer.term_score("python", tf=10, doc_len=50)
        score_tf11 = self.scorer.term_score("python", tf=11, doc_len=50)

        gain_low = score_tf2 - score_tf1
        gain_high = score_tf11 - score_tf10
        self.assertGreater(gain_low, gain_high)

    def test_score_length_normalization(self):
        # Shorter doc scores higher than longer doc for the same TF
        score_short = self.scorer.score(
            ["python"],
            {"python": 3},
            doc_len=20,
        )
        score_long = self.scorer.score(
            ["python"],
            {"python": 3},
            doc_len=200,
        )
        self.assertGreater(score_short, score_long)

    def test_term_score_zero_for_zero_tf(self):
        score = self.scorer.term_score("python", tf=0, doc_len=50)
        self.assertEqual(score, 0.0)

    def test_term_score_negative_tf_returns_zero(self):
        score = self.scorer.term_score("python", tf=-1, doc_len=50)
        self.assertEqual(score, 0.0)


class TestBM25ScoreNormalized(unittest.TestCase):
    """Tests for score_normalized."""

    def setUp(self):
        self.scorer = BM25Scorer(k1=1.2, b=0.75)
        self.scorer.update_stats(
            num_docs=100,
            avg_doc_len=50.0,
            doc_freqs={"python": 10, "java": 20},
        )

    def test_score_normalized_in_range(self):
        norm = self.scorer.score_normalized(
            ["python", "java"],
            {"python": 5, "java": 3},
            doc_len=50,
        )
        self.assertGreaterEqual(norm, 0.0)
        self.assertLessEqual(norm, 1.0)

    def test_score_normalized_zero_when_no_match(self):
        norm = self.scorer.score_normalized(
            ["python"],
            {"other": 5},
            doc_len=50,
        )
        self.assertEqual(norm, 0.0)

    def test_score_normalized_zero_for_empty_query(self):
        norm = self.scorer.score_normalized(
            [],
            {"python": 5},
            doc_len=50,
        )
        self.assertEqual(norm, 0.0)


class TestBM25ScoreUpperBound(unittest.TestCase):
    """Tests for score_upper_bound (WAND pruning)."""

    def setUp(self):
        self.scorer = BM25Scorer(k1=1.2, b=0.75)
        self.scorer.update_stats(
            num_docs=100,
            avg_doc_len=50.0,
            doc_freqs={"python": 10, "java": 20, "code": 80},
        )

    def test_score_upper_bound_geq_score(self):
        query = ["python", "java", "code"]
        doc_tf = {"python": 5, "java": 3, "code": 10}
        doc_len = 50

        actual = self.scorer.score(query, doc_tf, doc_len)
        upper = self.scorer.score_upper_bound(query, doc_len)
        self.assertGreaterEqual(upper, actual)

    def test_upper_bound_geq_score_short_doc(self):
        query = ["python", "java"]
        doc_tf = {"python": 20, "java": 15}
        doc_len = 5

        actual = self.scorer.score(query, doc_tf, doc_len)
        upper = self.scorer.score_upper_bound(query, doc_len)
        self.assertGreaterEqual(upper, actual)

    def test_upper_bound_geq_score_long_doc(self):
        query = ["python"]
        doc_tf = {"python": 1}
        doc_len = 500

        actual = self.scorer.score(query, doc_tf, doc_len)
        upper = self.scorer.score_upper_bound(query, doc_len)
        self.assertGreaterEqual(upper, actual)


class TestBM25UpdateStats(unittest.TestCase):
    """Tests for update_stats and add_document."""

    def test_update_stats(self):
        scorer = BM25Scorer()
        scorer.update_stats(
            num_docs=50,
            avg_doc_len=25.0,
            doc_freqs={"alpha": 10, "beta": 5},
        )
        self.assertEqual(scorer.num_docs, 50)
        self.assertAlmostEqual(scorer.avg_doc_len, 25.0)

    def test_update_stats_zero_avg_doc_len(self):
        # avg_doc_len=0 should be clamped to 1.0 to avoid division by zero
        scorer = BM25Scorer()
        scorer.update_stats(num_docs=0, avg_doc_len=0.0, doc_freqs={})
        self.assertAlmostEqual(scorer.avg_doc_len, 1.0)

    def test_add_document_incremental(self):
        scorer = BM25Scorer()

        scorer.add_document(doc_len=10, term_freqs={"alpha": 3, "beta": 1})
        self.assertEqual(scorer.num_docs, 1)
        self.assertAlmostEqual(scorer.avg_doc_len, 10.0)

        scorer.add_document(doc_len=20, term_freqs={"alpha": 2, "gamma": 5})
        self.assertEqual(scorer.num_docs, 2)
        self.assertAlmostEqual(scorer.avg_doc_len, 15.0)

        # IDF should reflect doc_freqs: alpha in 2 docs, beta in 1, gamma in 1
        self.assertAlmostEqual(
            scorer.idf("beta"), scorer.idf("gamma"),
            msg="beta and gamma should have same IDF (both in 1 doc)",
        )
        self.assertGreater(
            scorer.idf("beta"), scorer.idf("alpha"),
            msg="beta (df=1) should have higher IDF than alpha (df=2)",
        )

    def test_add_document_three_docs(self):
        scorer = BM25Scorer()
        scorer.add_document(doc_len=6, term_freqs={"x": 2})
        scorer.add_document(doc_len=12, term_freqs={"x": 1, "y": 3})
        scorer.add_document(doc_len=9, term_freqs={"y": 4, "z": 1})

        self.assertEqual(scorer.num_docs, 3)
        self.assertAlmostEqual(scorer.avg_doc_len, 9.0)


class TestBM25Parameters(unittest.TestCase):
    """Tests for different k1 and b parameter values."""

    def test_default_parameters(self):
        scorer = BM25Scorer()
        self.assertAlmostEqual(scorer.k1, 1.2)
        self.assertAlmostEqual(scorer.b, 0.75)

    def test_custom_parameters(self):
        scorer = BM25Scorer(k1=2.0, b=0.5)
        self.assertAlmostEqual(scorer.k1, 2.0)
        self.assertAlmostEqual(scorer.b, 0.5)

    def test_b_zero_no_length_normalization(self):
        # With b=0, document length should not affect score
        scorer = BM25Scorer(k1=1.2, b=0.0)
        scorer.update_stats(num_docs=10, avg_doc_len=50.0, doc_freqs={"t": 3})

        score_short = scorer.term_score("t", tf=5, doc_len=10)
        score_long = scorer.term_score("t", tf=5, doc_len=200)
        self.assertAlmostEqual(score_short, score_long)


if __name__ == "__main__":
    unittest.main()
