"""Tests for rbsp.query.ensemble -- FusionMethod, normalize_scores, ScoreFusion."""

import math
import os
import sys
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
from rbsp.query.ensemble import FusionMethod, normalize_scores, ScoreFusion


# ---------------------------------------------------------------------------
# FusionMethod enum
# ---------------------------------------------------------------------------

class TestFusionMethod(unittest.TestCase):
    """Tests for the FusionMethod enum."""

    def test_all_methods_exist(self):
        methods = [
            FusionMethod.LINEAR,
            FusionMethod.RECIPROCAL_RANK,
            FusionMethod.COMBSCORE,
            FusionMethod.COMBMAX,
            FusionMethod.BORDA,
            FusionMethod.MIN,
        ]
        self.assertEqual(len(methods), 6)

    def test_unique_values(self):
        values = [m.value for m in FusionMethod]
        self.assertEqual(len(values), len(set(values)))


# ---------------------------------------------------------------------------
# normalize_scores
# ---------------------------------------------------------------------------

class TestNormalizeScoresMinMax(unittest.TestCase):
    """Tests for normalize_scores with minmax method."""

    def test_basic_minmax(self):
        ranking = [("a", 1.0), ("b", 5.0), ("c", 3.0)]
        normed = normalize_scores(ranking, method="minmax")
        normed_dict = dict(normed)
        self.assertAlmostEqual(normed_dict["a"], 0.0)
        self.assertAlmostEqual(normed_dict["b"], 1.0)
        self.assertAlmostEqual(normed_dict["c"], 0.5)

    def test_empty_ranking(self):
        self.assertEqual(normalize_scores([], method="minmax"), [])

    def test_single_doc(self):
        normed = normalize_scores([("a", 5.0)], method="minmax")
        # span = 0 -> all scores become 0.0
        self.assertAlmostEqual(normed[0][1], 0.0)

    def test_identical_scores(self):
        normed = normalize_scores([("a", 3.0), ("b", 3.0)], method="minmax")
        for _, score in normed:
            self.assertAlmostEqual(score, 0.0)

    def test_preserves_order(self):
        ranking = [("c", 1.0), ("a", 5.0), ("b", 3.0)]
        normed = normalize_scores(ranking, method="minmax")
        doc_ids = [d for d, _ in normed]
        self.assertEqual(doc_ids, ["c", "a", "b"])


class TestNormalizeScoresZScore(unittest.TestCase):
    """Tests for normalize_scores with zscore method."""

    def test_basic_zscore(self):
        ranking = [("a", 1.0), ("b", 5.0), ("c", 3.0)]
        normed = normalize_scores(ranking, method="zscore")
        for _, score in normed:
            self.assertGreaterEqual(score, 0.0)
            self.assertLessEqual(score, 1.0)

    def test_empty_ranking(self):
        self.assertEqual(normalize_scores([], method="zscore"), [])

    def test_identical_scores_zscore(self):
        normed = normalize_scores([("a", 3.0), ("b", 3.0)], method="zscore")
        for _, score in normed:
            self.assertAlmostEqual(score, 0.0)

    def test_unknown_method_raises(self):
        with self.assertRaises(ValueError):
            normalize_scores([("a", 1.0)], method="unknown")


# ---------------------------------------------------------------------------
# ScoreFusion.linear
# ---------------------------------------------------------------------------

class TestScoreFusionLinear(unittest.TestCase):
    """Tests for ScoreFusion.linear."""

    def test_basic_linear(self):
        r1 = [("a", 1.0), ("b", 0.5)]
        r2 = [("a", 0.5), ("b", 1.0)]
        result = ScoreFusion.linear([r1, r2], weights=[1.0, 1.0])
        result_dict = dict(result)
        # Both docs appear in both rankings
        self.assertIn("a", result_dict)
        self.assertIn("b", result_dict)

    def test_weights_affect_ranking(self):
        r1 = [("a", 10.0), ("b", 1.0)]
        r2 = [("a", 1.0), ("b", 10.0)]
        # Heavy weight on r1 -> "a" should rank higher
        result = ScoreFusion.linear([r1, r2], weights=[10.0, 1.0])
        self.assertEqual(result[0][0], "a")

    def test_mismatched_lengths_raises(self):
        with self.assertRaises(ValueError):
            ScoreFusion.linear([[], []], weights=[1.0])

    def test_zero_weights_raises(self):
        with self.assertRaises(ValueError):
            ScoreFusion.linear([[("a", 1.0)]], weights=[0.0])

    def test_empty_rankings(self):
        result = ScoreFusion.linear([], weights=[])
        self.assertEqual(result, [])

    def test_sorted_descending(self):
        r1 = [("a", 1.0), ("b", 3.0), ("c", 2.0)]
        result = ScoreFusion.linear([r1], weights=[1.0])
        scores = [s for _, s in result]
        self.assertEqual(scores, sorted(scores, reverse=True))


# ---------------------------------------------------------------------------
# ScoreFusion.reciprocal_rank_fusion
# ---------------------------------------------------------------------------

class TestScoreFusionRRF(unittest.TestCase):
    """Tests for ScoreFusion.reciprocal_rank_fusion."""

    def test_basic_rrf(self):
        r1 = [("a", 3.0), ("b", 2.0), ("c", 1.0)]
        r2 = [("b", 3.0), ("a", 2.0), ("c", 1.0)]
        result = ScoreFusion.reciprocal_rank_fusion([r1, r2])
        result_dict = dict(result)
        # "a" is rank 1 in r1, rank 2 in r2 -> 1/(60+1) + 1/(60+2)
        # "b" is rank 2 in r1, rank 1 in r2 -> same total
        self.assertIn("a", result_dict)
        self.assertIn("b", result_dict)
        # a and b should have the same RRF score (symmetric ranks)
        self.assertAlmostEqual(result_dict["a"], result_dict["b"])

    def test_empty_rankings(self):
        self.assertEqual(ScoreFusion.reciprocal_rank_fusion([]), [])

    def test_single_ranking(self):
        r1 = [("a", 5.0), ("b", 3.0)]
        result = ScoreFusion.reciprocal_rank_fusion([r1], k=60)
        # "a" rank 1 -> 1/61, "b" rank 2 -> 1/62
        result_dict = dict(result)
        self.assertGreater(result_dict["a"], result_dict["b"])

    def test_custom_k(self):
        r1 = [("a", 5.0)]
        result_k1 = ScoreFusion.reciprocal_rank_fusion([r1], k=1)
        result_k60 = ScoreFusion.reciprocal_rank_fusion([r1], k=60)
        # With k=1: score = 1/(1+1) = 0.5
        # With k=60: score = 1/(60+1) ~ 0.0164
        self.assertGreater(dict(result_k1)["a"], dict(result_k60)["a"])

    def test_sorted_descending(self):
        r1 = [("a", 1.0), ("b", 3.0), ("c", 2.0)]
        result = ScoreFusion.reciprocal_rank_fusion([r1])
        scores = [s for _, s in result]
        self.assertEqual(scores, sorted(scores, reverse=True))


# ---------------------------------------------------------------------------
# ScoreFusion.combscore
# ---------------------------------------------------------------------------

class TestScoreFusionCombScore(unittest.TestCase):
    """Tests for ScoreFusion.combscore."""

    def test_basic_combscore(self):
        r1 = [("a", 2.0), ("b", 1.0)]
        r2 = [("a", 3.0), ("c", 4.0)]
        result = ScoreFusion.combscore([r1, r2])
        result_dict = dict(result)
        self.assertAlmostEqual(result_dict["a"], 5.0)
        self.assertAlmostEqual(result_dict["b"], 1.0)
        self.assertAlmostEqual(result_dict["c"], 4.0)

    def test_empty_rankings(self):
        self.assertEqual(ScoreFusion.combscore([]), [])

    def test_single_doc_single_ranking(self):
        result = ScoreFusion.combscore([[("a", 7.0)]])
        self.assertEqual(len(result), 1)
        self.assertAlmostEqual(result[0][1], 7.0)

    def test_sorted_descending(self):
        r1 = [("a", 1.0), ("b", 3.0), ("c", 2.0)]
        result = ScoreFusion.combscore([r1])
        scores = [s for _, s in result]
        self.assertEqual(scores, sorted(scores, reverse=True))


# ---------------------------------------------------------------------------
# ScoreFusion.combmax
# ---------------------------------------------------------------------------

class TestScoreFusionCombMax(unittest.TestCase):
    """Tests for ScoreFusion.combmax."""

    def test_basic_combmax(self):
        r1 = [("a", 2.0), ("b", 5.0)]
        r2 = [("a", 8.0), ("b", 1.0)]
        result = ScoreFusion.combmax([r1, r2])
        result_dict = dict(result)
        self.assertAlmostEqual(result_dict["a"], 8.0)
        self.assertAlmostEqual(result_dict["b"], 5.0)

    def test_empty_rankings(self):
        self.assertEqual(ScoreFusion.combmax([]), [])

    def test_single_ranking(self):
        result = ScoreFusion.combmax([[("a", 3.0), ("b", 7.0)]])
        result_dict = dict(result)
        self.assertAlmostEqual(result_dict["a"], 3.0)
        self.assertAlmostEqual(result_dict["b"], 7.0)

    def test_sorted_descending(self):
        r1 = [("a", 1.0), ("b", 3.0), ("c", 2.0)]
        result = ScoreFusion.combmax([r1])
        scores = [s for _, s in result]
        self.assertEqual(scores, sorted(scores, reverse=True))


# ---------------------------------------------------------------------------
# ScoreFusion.borda_count
# ---------------------------------------------------------------------------

class TestScoreFusionBordaCount(unittest.TestCase):
    """Tests for ScoreFusion.borda_count."""

    def test_basic_borda(self):
        r1 = [("a", 3.0), ("b", 2.0), ("c", 1.0)]
        result = ScoreFusion.borda_count([r1])
        result_dict = dict(result)
        # n=3: rank 0 -> 2 pts, rank 1 -> 1 pt, rank 2 -> 0 pts
        self.assertAlmostEqual(result_dict["a"], 2.0)
        self.assertAlmostEqual(result_dict["b"], 1.0)
        self.assertAlmostEqual(result_dict["c"], 0.0)

    def test_two_rankings(self):
        r1 = [("a", 3.0), ("b", 2.0), ("c", 1.0)]
        r2 = [("c", 3.0), ("b", 2.0), ("a", 1.0)]
        result = ScoreFusion.borda_count([r1, r2])
        result_dict = dict(result)
        # a: 2 + 0 = 2, b: 1 + 1 = 2, c: 0 + 2 = 2 (all tied)
        self.assertAlmostEqual(result_dict["a"], result_dict["b"])
        self.assertAlmostEqual(result_dict["b"], result_dict["c"])

    def test_empty_rankings(self):
        self.assertEqual(ScoreFusion.borda_count([]), [])

    def test_sorted_descending(self):
        r1 = [("a", 5.0), ("b", 3.0), ("c", 1.0)]
        result = ScoreFusion.borda_count([r1])
        scores = [s for _, s in result]
        self.assertEqual(scores, sorted(scores, reverse=True))


# ---------------------------------------------------------------------------
# ScoreFusion.combmin
# ---------------------------------------------------------------------------

class TestScoreFusionCombMin(unittest.TestCase):
    """Tests for ScoreFusion.combmin."""

    def test_basic_combmin(self):
        r1 = [("a", 2.0), ("b", 5.0)]
        r2 = [("a", 8.0), ("b", 1.0)]
        result = ScoreFusion.combmin([r1, r2])
        result_dict = dict(result)
        self.assertAlmostEqual(result_dict["a"], 2.0)
        self.assertAlmostEqual(result_dict["b"], 1.0)

    def test_empty_rankings(self):
        self.assertEqual(ScoreFusion.combmin([]), [])

    def test_single_ranking(self):
        result = ScoreFusion.combmin([[("a", 3.0)]])
        self.assertAlmostEqual(result[0][1], 3.0)

    def test_doc_in_single_ranking_uses_that_score(self):
        r1 = [("a", 5.0)]
        r2 = [("b", 3.0)]
        result = ScoreFusion.combmin([r1, r2])
        result_dict = dict(result)
        self.assertAlmostEqual(result_dict["a"], 5.0)
        self.assertAlmostEqual(result_dict["b"], 3.0)

    def test_sorted_descending(self):
        r1 = [("a", 1.0), ("b", 3.0), ("c", 2.0)]
        result = ScoreFusion.combmin([r1])
        scores = [s for _, s in result]
        self.assertEqual(scores, sorted(scores, reverse=True))


# ---------------------------------------------------------------------------
# ScoreFusion.fuse dispatch
# ---------------------------------------------------------------------------

class TestScoreFusionFuse(unittest.TestCase):
    """Tests for ScoreFusion.fuse dispatch method."""

    def test_fuse_linear(self):
        r1 = [("a", 1.0), ("b", 2.0)]
        result = ScoreFusion.fuse(
            FusionMethod.LINEAR, [r1], weights=[1.0]
        )
        self.assertIsInstance(result, list)

    def test_fuse_rrf(self):
        r1 = [("a", 1.0)]
        result = ScoreFusion.fuse(FusionMethod.RECIPROCAL_RANK, [r1])
        self.assertIsInstance(result, list)

    def test_fuse_combscore(self):
        r1 = [("a", 1.0)]
        result = ScoreFusion.fuse(FusionMethod.COMBSCORE, [r1])
        self.assertIsInstance(result, list)

    def test_fuse_combmax(self):
        r1 = [("a", 1.0)]
        result = ScoreFusion.fuse(FusionMethod.COMBMAX, [r1])
        self.assertIsInstance(result, list)

    def test_fuse_borda(self):
        r1 = [("a", 1.0)]
        result = ScoreFusion.fuse(FusionMethod.BORDA, [r1])
        self.assertIsInstance(result, list)

    def test_fuse_min(self):
        r1 = [("a", 1.0)]
        result = ScoreFusion.fuse(FusionMethod.MIN, [r1])
        self.assertIsInstance(result, list)

    def test_fuse_linear_requires_weights(self):
        with self.assertRaises(ValueError):
            ScoreFusion.fuse(FusionMethod.LINEAR, [[("a", 1.0)]])

    def test_fuse_rrf_custom_k(self):
        r1 = [("a", 5.0)]
        result = ScoreFusion.fuse(FusionMethod.RECIPROCAL_RANK, [r1], k=10)
        score = result[0][1]
        # k=10, rank=1 -> 1/(10+1)
        self.assertAlmostEqual(score, 1.0 / 11.0)


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------

class TestScoreFusionEdgeCases(unittest.TestCase):
    """Tests for edge cases in ScoreFusion."""

    def test_empty_rankings_all_methods(self):
        for method in FusionMethod:
            if method is FusionMethod.LINEAR:
                result = ScoreFusion.fuse(method, [], weights=[])
            else:
                result = ScoreFusion.fuse(method, [])
            self.assertEqual(result, [], f"Failed for {method}")

    def test_single_doc_all_methods(self):
        ranking = [[("only_doc", 5.0)]]
        for method in FusionMethod:
            if method is FusionMethod.LINEAR:
                result = ScoreFusion.fuse(method, ranking, weights=[1.0])
            else:
                result = ScoreFusion.fuse(method, ranking)
            self.assertEqual(len(result), 1, f"Failed for {method}")
            self.assertEqual(result[0][0], "only_doc")

    def test_disjoint_rankings(self):
        r1 = [("a", 5.0)]
        r2 = [("b", 3.0)]
        result = ScoreFusion.combscore([r1, r2])
        result_dict = dict(result)
        self.assertAlmostEqual(result_dict["a"], 5.0)
        self.assertAlmostEqual(result_dict["b"], 3.0)

    def test_many_rankings(self):
        rankings = [[("a", float(i))] for i in range(10)]
        result = ScoreFusion.combscore(rankings)
        result_dict = dict(result)
        # sum of 0..9 = 45
        self.assertAlmostEqual(result_dict["a"], 45.0)


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