"""Tests for rbsp.query.expand — Query expansion via pseudo-relevance feedback."""

from __future__ import annotations

import math

import pytest

from rbsp.query.expand import CooccurrenceExpander, RocchioExpander, SynonymExpander


# ===================================================================
# RocchioExpander
# ===================================================================


class TestRocchioExpanderBasic:
    def test_empty_query(self) -> None:
        r = RocchioExpander()
        assert r.expand([], [["a", "b"]]) == []

    def test_no_feedback_docs(self) -> None:
        r = RocchioExpander()
        result = r.expand(["hello"], [])
        assert len(result) == 1
        assert result[0][0] == "hello"
        assert result[0][1] == r.alpha

    def test_single_positive_doc(self) -> None:
        r = RocchioExpander()
        result = r.expand(["search"], [["search", "engine", "index"]])
        # Should produce expansion terms from the feedback doc
        terms = [t for t, w in result]
        assert len(terms) >= 1

    def test_expansion_adds_new_terms(self) -> None:
        r = RocchioExpander(num_expand_terms=5)
        result = r.expand(
            ["search"],
            [["search", "engine", "index", "query", "rank"]],
        )
        terms = {t for t, w in result}
        # Should have expansion terms from feedback
        assert len(terms) >= 1

    def test_num_expand_terms_limit(self) -> None:
        r = RocchioExpander(num_expand_terms=2)
        result = r.expand(
            ["a"],
            [["a", "b", "c", "d", "e", "f", "g"]],
        )
        non_original = [t for t, w in result if t != "a"]
        assert len(non_original) <= 2

    def test_num_feedback_docs_limit(self) -> None:
        r = RocchioExpander(num_feedback_docs=1)
        docs = [["a", "b"], ["a", "c"], ["a", "d"]]
        result = r.expand(["a"], docs)
        # Only first doc used, so "c" and "d" shouldn't get high weight
        terms = {t for t, w in result}
        assert "b" in terms

    def test_weights_positive(self) -> None:
        r = RocchioExpander()
        result = r.expand(["search"], [["search", "engine"]])
        for _, w in result:
            assert w > 0


class TestRocchioExpanderNegative:
    def test_negative_feedback_reduces_terms(self) -> None:
        r = RocchioExpander(gamma=1.0, beta=0.0, num_expand_terms=20)
        result = r.expand(
            ["search"],
            [["search", "engine"]],
            negative_doc_tokens=[["search", "engine", "spam"]],
        )
        terms = {t for t, w in result}
        # "spam" should be suppressed or have low weight
        if "spam" in terms:
            spam_w = next(w for t, w in result if t == "spam")
            assert spam_w <= 0.1  # very low weight

    def test_negative_docs_limit(self) -> None:
        r = RocchioExpander(num_feedback_docs=1)
        result = r.expand(
            ["a"],
            [["a", "b"]],
            negative_doc_tokens=[["a", "c"], ["a", "d"]],
        )
        # Should still produce results without error
        assert len(result) >= 1


class TestRocchioExpanderParameters:
    def test_alpha_zero(self) -> None:
        r = RocchioExpander(alpha=0.0, beta=1.0)
        result = r.expand(["q"], [["q", "new"]])
        # Original query weight zeroed out from alpha, but may appear via beta
        assert isinstance(result, list)

    def test_custom_params(self) -> None:
        r = RocchioExpander(alpha=2.0, beta=0.5, gamma=0.1)
        assert r.alpha == 2.0
        assert r.beta == 0.5
        assert r.gamma == 0.1


# ===================================================================
# SynonymExpander
# ===================================================================


class TestSynonymExpanderBasic:
    def test_empty_expand(self) -> None:
        se = SynonymExpander()
        assert se.expand("nothing") == set()

    def test_add_synonym(self) -> None:
        se = SynonymExpander()
        se.add_synonym("big", "large")
        assert "large" in se.expand("big")
        # Directional: reverse should not exist
        assert "big" not in se.expand("large")

    def test_add_synonym_group(self) -> None:
        se = SynonymExpander()
        se.add_synonym_group(["big", "large", "huge"])
        assert "large" in se.expand("big")
        assert "big" in se.expand("large")
        assert "huge" in se.expand("big")

    def test_expand_does_not_include_self(self) -> None:
        se = SynonymExpander()
        se.add_synonym("big", "large")
        assert "big" not in se.expand("big")


class TestSynonymExpanderQuery:
    def test_expand_query_basic(self) -> None:
        se = SynonymExpander()
        se.add_synonym("big", "large")
        result = se.expand_query(["big", "dog"])
        assert result[0] == "big"
        assert "large" in result
        assert "dog" in result

    def test_expand_query_empty(self) -> None:
        se = SynonymExpander()
        assert se.expand_query([]) == []

    def test_expand_query_deduplication(self) -> None:
        se = SynonymExpander()
        se.add_synonym("big", "large")
        se.add_synonym("huge", "large")
        result = se.expand_query(["big", "huge"])
        assert result.count("large") == 1

    def test_expand_query_preserves_order(self) -> None:
        se = SynonymExpander()
        result = se.expand_query(["a", "b", "c"])
        assert result == ["a", "b", "c"]

    def test_expand_query_sorted_synonyms(self) -> None:
        se = SynonymExpander()
        se.add_synonym("a", "z")
        se.add_synonym("a", "m")
        result = se.expand_query(["a"])
        # Synonyms should be sorted: m before z
        idx_m = result.index("m")
        idx_z = result.index("z")
        assert idx_m < idx_z


class TestSynonymExpanderBulkLoad:
    def test_load_from_dict(self) -> None:
        se = SynonymExpander()
        se.load_from_dict({"big": ["large", "huge"], "fast": ["quick"]})
        assert "large" in se.expand("big")
        assert "huge" in se.expand("big")
        assert "quick" in se.expand("fast")

    def test_load_from_dict_empty(self) -> None:
        se = SynonymExpander()
        se.load_from_dict({})
        assert se.expand("anything") == set()

    def test_load_from_dict_merges(self) -> None:
        se = SynonymExpander()
        se.add_synonym("big", "large")
        se.load_from_dict({"big": ["huge"]})
        syns = se.expand("big")
        assert "large" in syns
        assert "huge" in syns


# ===================================================================
# CooccurrenceExpander
# ===================================================================


class TestCooccurrenceExpanderBasic:
    def test_empty_expander(self) -> None:
        ce = CooccurrenceExpander()
        assert ce.expand("hello") == []

    def test_single_document(self) -> None:
        # With only 1 doc, P(joint) = P(x)*P(y), so PMI = 0 → no positive PMI
        # Need multiple docs for meaningful PMI
        ce = CooccurrenceExpander()
        ce.add_document(["search", "engine"])
        ce.add_document(["search", "engine"])
        ce.add_document(["other", "stuff"])
        results = ce.expand("search")
        terms = {t for t, pmi in results}
        assert "engine" in terms

    def test_pmi_positive_only(self) -> None:
        ce = CooccurrenceExpander()
        ce.add_document(["a", "b"])
        ce.add_document(["a", "b"])
        results = ce.expand("a")
        for _, pmi in results:
            assert pmi > 0

    def test_top_k_limit(self) -> None:
        ce = CooccurrenceExpander()
        ce.add_document(["a", "b", "c", "d", "e", "f"])
        results = ce.expand("a", top_k=2)
        assert len(results) <= 2

    def test_nonexistent_term(self) -> None:
        ce = CooccurrenceExpander()
        ce.add_document(["a", "b"])
        assert ce.expand("xyz") == []


class TestCooccurrenceExpanderMultiDoc:
    def test_multiple_documents(self) -> None:
        ce = CooccurrenceExpander()
        ce.add_document(["search", "engine"])
        ce.add_document(["search", "query"])
        ce.add_document(["engine", "index"])
        results = ce.expand("search")
        terms = {t for t, _ in results}
        assert "engine" in terms or "query" in terms

    def test_independent_terms_low_pmi(self) -> None:
        ce = CooccurrenceExpander()
        # "a" and "b" never co-occur
        ce.add_document(["a", "c"])
        ce.add_document(["b", "d"])
        results = ce.expand("a")
        terms = {t for t, _ in results}
        assert "b" not in terms

    def test_pmi_sorted_descending(self) -> None:
        ce = CooccurrenceExpander()
        ce.add_document(["a", "b", "c"])
        ce.add_document(["a", "b"])
        ce.add_document(["a", "c", "d"])
        results = ce.expand("a", top_k=10)
        pmis = [pmi for _, pmi in results]
        assert pmis == sorted(pmis, reverse=True)

    def test_empty_document_ignored(self) -> None:
        ce = CooccurrenceExpander()
        ce.add_document([])
        ce.add_document(["a", "b"])
        results = ce.expand("a")
        assert len(results) >= 1

    def test_duplicate_tokens_in_doc(self) -> None:
        ce = CooccurrenceExpander()
        # Need multiple docs for PMI to be positive
        ce.add_document(["a", "a", "b", "b"])
        ce.add_document(["a", "b"])
        ce.add_document(["c", "d"])
        results = ce.expand("a")
        terms = {t for t, _ in results}
        assert "b" in terms
