"""Tests for rbsp.core.blockwand -- Block-Max WAND algorithm."""

import os
import sys
import unittest

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
from rbsp.core.blockwand import (
    BLOCK_SIZE,
    Block,
    BlockMaxWAND,
    BlockPostingList,
    build_block_posting_list,
)


class TestBlock(unittest.TestCase):
    """Tests for the Block dataclass."""

    def test_block_creation(self):
        b = Block(doc_ids=[1, 2, 3], scores=[1.0, 2.0, 0.5], max_score=2.0)
        self.assertEqual(b.doc_ids, [1, 2, 3])
        self.assertEqual(b.scores, [1.0, 2.0, 0.5])
        self.assertEqual(b.max_score, 2.0)

    def test_block_single_entry(self):
        b = Block(doc_ids=[42], scores=[3.14], max_score=3.14)
        self.assertEqual(len(b.doc_ids), 1)
        self.assertEqual(b.max_score, 3.14)


class TestBlockPostingList(unittest.TestCase):
    """Tests for BlockPostingList."""

    def _make_pl(self, blocks):
        return BlockPostingList(blocks)

    def test_max_score_multiple_blocks(self):
        blocks = [
            Block([1, 2], [1.0, 0.5], 1.0),
            Block([5, 6], [3.0, 2.0], 3.0),
            Block([10, 11], [0.1, 0.2], 0.2),
        ]
        pl = self._make_pl(blocks)
        self.assertEqual(pl.max_score(), 3.0)

    def test_max_score_empty(self):
        pl = self._make_pl([])
        self.assertEqual(pl.max_score(), 0.0)

    def test_advance_to_block_first(self):
        blocks = [
            Block([1, 2, 3], [1.0, 1.0, 1.0], 1.0),
            Block([10, 11, 12], [2.0, 2.0, 2.0], 2.0),
        ]
        pl = self._make_pl(blocks)
        block = pl.advance_to_block(0)
        self.assertIsNotNone(block)
        self.assertEqual(block.doc_ids, [1, 2, 3])

    def test_advance_to_block_skips(self):
        blocks = [
            Block([1, 2, 3], [1.0, 1.0, 1.0], 1.0),
            Block([10, 11, 12], [2.0, 2.0, 2.0], 2.0),
            Block([20, 21, 22], [3.0, 3.0, 3.0], 3.0),
        ]
        pl = self._make_pl(blocks)
        block = pl.advance_to_block(10)
        self.assertIsNotNone(block)
        self.assertEqual(block.doc_ids, [10, 11, 12])

    def test_advance_to_block_within_block(self):
        blocks = [
            Block([1, 2, 3], [1.0, 1.0, 1.0], 1.0),
            Block([10, 11, 12], [2.0, 2.0, 2.0], 2.0),
        ]
        pl = self._make_pl(blocks)
        # doc_id=2 is within the first block (last doc = 3 >= 2)
        block = pl.advance_to_block(2)
        self.assertIsNotNone(block)
        self.assertEqual(block.doc_ids[0], 1)

    def test_advance_to_block_past_all(self):
        blocks = [
            Block([1, 2], [1.0, 1.0], 1.0),
            Block([5, 6], [2.0, 2.0], 2.0),
        ]
        pl = self._make_pl(blocks)
        block = pl.advance_to_block(100)
        self.assertIsNone(block)

    def test_current_block_max(self):
        blocks = [
            Block([1, 2], [1.0, 0.5], 1.0),
            Block([5, 6], [3.0, 2.0], 3.0),
        ]
        pl = self._make_pl(blocks)
        self.assertEqual(pl.current_block_max(), 1.0)
        pl.advance_to_block(5)
        self.assertEqual(pl.current_block_max(), 3.0)

    def test_current_block_max_exhausted(self):
        blocks = [Block([1], [1.0], 1.0)]
        pl = self._make_pl(blocks)
        pl.advance_to_block(100)
        self.assertEqual(pl.current_block_max(), 0.0)

    def test_exhausted_property(self):
        blocks = [Block([1], [1.0], 1.0)]
        pl = self._make_pl(blocks)
        self.assertFalse(pl.exhausted)
        pl.advance_to_block(100)
        self.assertTrue(pl.exhausted)

    def test_reset(self):
        blocks = [
            Block([1, 2], [1.0, 0.5], 1.0),
            Block([5, 6], [3.0, 2.0], 3.0),
        ]
        pl = self._make_pl(blocks)
        pl.advance_to_block(5)
        self.assertEqual(pl.current_block_max(), 3.0)
        pl.reset()
        self.assertEqual(pl.current_block_max(), 1.0)

    def test_current_block_property(self):
        blocks = [Block([1, 2], [1.0, 0.5], 1.0)]
        pl = self._make_pl(blocks)
        self.assertIsNotNone(pl.current_block)
        self.assertEqual(pl.current_block.doc_ids, [1, 2])

    def test_current_block_none_when_empty(self):
        pl = self._make_pl([])
        self.assertIsNone(pl.current_block)


class TestBuildBlockPostingList(unittest.TestCase):
    """Tests for the build_block_posting_list helper."""

    def test_basic_build(self):
        doc_ids = [1, 2, 3, 4, 5]
        scores = [1.0, 2.0, 3.0, 4.0, 5.0]
        pl = build_block_posting_list(doc_ids, scores, block_size=3)
        self.assertEqual(pl.max_score(), 5.0)

    def test_block_size_splits(self):
        doc_ids = list(range(10))
        scores = [1.0] * 10
        pl = build_block_posting_list(doc_ids, scores, block_size=3)
        # Should create ceil(10/3) = 4 blocks
        self.assertEqual(len(pl._blocks), 4)

    def test_exact_block_boundary(self):
        doc_ids = list(range(6))
        scores = [1.0] * 6
        pl = build_block_posting_list(doc_ids, scores, block_size=3)
        self.assertEqual(len(pl._blocks), 2)

    def test_single_element(self):
        pl = build_block_posting_list([42], [3.14], block_size=128)
        self.assertEqual(pl.max_score(), 3.14)
        self.assertEqual(len(pl._blocks), 1)

    def test_empty_input(self):
        pl = build_block_posting_list([], [], block_size=128)
        self.assertEqual(pl.max_score(), 0.0)
        self.assertEqual(len(pl._blocks), 0)

    def test_mismatched_lengths_raises(self):
        with self.assertRaises(ValueError):
            build_block_posting_list([1, 2], [1.0], block_size=128)

    def test_default_block_size(self):
        self.assertEqual(BLOCK_SIZE, 128)
        # Verify default is used
        doc_ids = list(range(300))
        scores = [1.0] * 300
        pl = build_block_posting_list(doc_ids, scores)
        # ceil(300 / 128) = 3
        self.assertEqual(len(pl._blocks), 3)

    def test_block_max_scores_correct(self):
        doc_ids = [0, 1, 2, 3, 4, 5]
        scores = [1.0, 5.0, 2.0, 3.0, 4.0, 6.0]
        pl = build_block_posting_list(doc_ids, scores, block_size=3)
        self.assertEqual(pl._blocks[0].max_score, 5.0)
        self.assertEqual(pl._blocks[1].max_score, 6.0)


class TestBlockMaxWAND(unittest.TestCase):
    """Tests for the BlockMaxWAND algorithm."""

    def test_empty_posting_lists(self):
        wand = BlockMaxWAND({}, top_k=5)
        self.assertEqual(wand.evaluate(), [])

    def test_single_term(self):
        pl = build_block_posting_list(
            [1, 3, 5, 7, 9],
            [1.0, 3.0, 2.0, 5.0, 4.0],
            block_size=3,
        )
        wand = BlockMaxWAND({"term": pl}, top_k=3)
        results = wand.evaluate()
        self.assertEqual(len(results), 3)
        # Top-3: doc 7 (5.0), doc 9 (4.0), doc 3 (3.0)
        doc_ids = [doc_id for doc_id, _score in results]
        self.assertEqual(doc_ids[0], 7)
        self.assertEqual(doc_ids[1], 9)
        self.assertEqual(doc_ids[2], 3)

    def test_two_terms_with_overlap(self):
        pl_a = build_block_posting_list(
            [1, 3, 5, 7, 9],
            [1.0, 2.0, 0.5, 3.0, 1.5],
            block_size=3,
        )
        pl_b = build_block_posting_list(
            [2, 3, 5, 8, 9],
            [0.8, 1.5, 2.0, 0.3, 1.0],
            block_size=3,
        )
        wand = BlockMaxWAND({"a": pl_a, "b": pl_b}, top_k=3)
        results = wand.evaluate()
        self.assertEqual(len(results), 3)
        # doc 3: 2.0 + 1.5 = 3.5
        # doc 7: 3.0
        # doc 5: 0.5 + 2.0 = 2.5 or doc 9: 1.5 + 1.0 = 2.5
        scores = {doc_id: score for doc_id, score in results}
        self.assertAlmostEqual(scores[3], 3.5)
        self.assertAlmostEqual(scores[7], 3.0)

    def test_top_k_limits_results(self):
        pl = build_block_posting_list(
            list(range(100)),
            [float(i) for i in range(100)],
            block_size=10,
        )
        wand = BlockMaxWAND({"term": pl}, top_k=5)
        results = wand.evaluate()
        self.assertEqual(len(results), 5)
        # Highest doc_ids should have highest scores
        doc_ids = [doc_id for doc_id, _score in results]
        self.assertIn(99, doc_ids)
        self.assertIn(98, doc_ids)

    def test_results_sorted_by_score_desc(self):
        pl = build_block_posting_list(
            [1, 2, 3, 4, 5],
            [5.0, 1.0, 3.0, 2.0, 4.0],
            block_size=3,
        )
        wand = BlockMaxWAND({"term": pl}, top_k=5)
        results = wand.evaluate()
        scores = [s for _, s in results]
        self.assertEqual(scores, sorted(scores, reverse=True))

    def test_disjoint_posting_lists(self):
        pl_a = build_block_posting_list([1, 2, 3], [1.0, 2.0, 3.0], block_size=3)
        pl_b = build_block_posting_list([4, 5, 6], [4.0, 5.0, 6.0], block_size=3)
        wand = BlockMaxWAND({"a": pl_a, "b": pl_b}, top_k=6)
        results = wand.evaluate()
        self.assertEqual(len(results), 6)
        scores = {doc_id: score for doc_id, score in results}
        self.assertAlmostEqual(scores[6], 6.0)
        self.assertAlmostEqual(scores[1], 1.0)

    def test_all_same_scores(self):
        pl = build_block_posting_list(
            [1, 2, 3, 4, 5],
            [1.0, 1.0, 1.0, 1.0, 1.0],
            block_size=2,
        )
        wand = BlockMaxWAND({"term": pl}, top_k=3)
        results = wand.evaluate()
        self.assertEqual(len(results), 3)
        for _, score in results:
            self.assertAlmostEqual(score, 1.0)

    def test_top_k_larger_than_corpus(self):
        pl = build_block_posting_list([1, 2, 3], [1.0, 2.0, 3.0], block_size=2)
        wand = BlockMaxWAND({"term": pl}, top_k=10)
        results = wand.evaluate()
        self.assertEqual(len(results), 3)

    def test_three_terms(self):
        pl_a = build_block_posting_list([1, 3, 5], [1.0, 1.0, 1.0], block_size=3)
        pl_b = build_block_posting_list([1, 2, 5], [2.0, 2.0, 2.0], block_size=3)
        pl_c = build_block_posting_list([1, 4, 5], [3.0, 3.0, 3.0], block_size=3)
        wand = BlockMaxWAND({"a": pl_a, "b": pl_b, "c": pl_c}, top_k=2)
        results = wand.evaluate()
        scores = {doc_id: score for doc_id, score in results}
        # doc 1: 1+2+3 = 6.0, doc 5: 1+2+3 = 6.0
        self.assertIn(1, scores)
        self.assertIn(5, scores)
        self.assertAlmostEqual(scores[1], 6.0)
        self.assertAlmostEqual(scores[5], 6.0)

    def test_large_block_size(self):
        # All postings in a single block
        doc_ids = list(range(50))
        scores = [float(i) for i in range(50)]
        pl = build_block_posting_list(doc_ids, scores, block_size=1000)
        self.assertEqual(len(pl._blocks), 1)
        wand = BlockMaxWAND({"term": pl}, top_k=3)
        results = wand.evaluate()
        top_ids = [doc_id for doc_id, _ in results]
        self.assertEqual(top_ids[0], 49)

    def test_single_doc_per_block(self):
        doc_ids = [10, 20, 30]
        scores = [3.0, 1.0, 2.0]
        pl = build_block_posting_list(doc_ids, scores, block_size=1)
        self.assertEqual(len(pl._blocks), 3)
        wand = BlockMaxWAND({"term": pl}, top_k=2)
        results = wand.evaluate()
        top_ids = [doc_id for doc_id, _ in results]
        self.assertEqual(top_ids[0], 10)
        self.assertEqual(top_ids[1], 30)

    def test_score_accumulation_correctness(self):
        # Verify exact score accumulation for documents in multiple lists
        pl_a = build_block_posting_list([100], [1.5], block_size=128)
        pl_b = build_block_posting_list([100], [2.5], block_size=128)
        pl_c = build_block_posting_list([100], [0.5], block_size=128)
        wand = BlockMaxWAND({"a": pl_a, "b": pl_b, "c": pl_c}, top_k=1)
        results = wand.evaluate()
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0][0], 100)
        self.assertAlmostEqual(results[0][1], 4.5)


class TestBlockMaxWANDPruning(unittest.TestCase):
    """Tests that verify pruning behavior doesn't affect correctness."""

    def test_pruning_does_not_miss_top_results(self):
        # Create a scenario where pruning is likely: many low-score docs
        # and a few high-score docs scattered across blocks
        n = 500
        doc_ids = list(range(n))
        scores = [0.01] * n
        # Plant high scores at specific positions
        scores[42] = 10.0
        scores[250] = 9.0
        scores[499] = 8.0

        pl = build_block_posting_list(doc_ids, scores, block_size=64)
        wand = BlockMaxWAND({"term": pl}, top_k=3)
        results = wand.evaluate()
        result_ids = [doc_id for doc_id, _ in results]
        self.assertIn(42, result_ids)
        self.assertIn(250, result_ids)
        self.assertIn(499, result_ids)

    def test_correctness_with_multi_term_pruning(self):
        # Two terms: term_a has high scores in early docs, term_b in late docs
        n = 200
        ids_a = list(range(n))
        scores_a = [float(n - i) for i in range(n)]  # decreasing
        ids_b = list(range(n))
        scores_b = [float(i) for i in range(n)]  # increasing

        pl_a = build_block_posting_list(ids_a, scores_a, block_size=32)
        pl_b = build_block_posting_list(ids_b, scores_b, block_size=32)

        wand = BlockMaxWAND({"a": pl_a, "b": pl_b}, top_k=5)
        results = wand.evaluate()

        # Brute force: compute all scores
        brute = {}
        for i in range(n):
            brute[i] = scores_a[i] + scores_b[i]
        top5_brute = sorted(brute.items(), key=lambda x: x[1], reverse=True)[:5]
        top5_brute_ids = {doc_id for doc_id, _ in top5_brute}

        result_ids = {doc_id for doc_id, _ in results}
        self.assertEqual(result_ids, top5_brute_ids)


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