"""Tests for rbsp.query.language_model — Language model scorers."""

from __future__ import annotations

import math

import pytest

from rbsp.query.language_model import (
    AbsoluteDiscountLMScorer,
    DirichletLMScorer,
    JelinekMercerLMScorer,
)


# ===================================================================
# DirichletLMScorer
# ===================================================================


class TestDirichletInit:
    def test_default_mu(self) -> None:
        d = DirichletLMScorer()
        assert d.mu == 2000.0

    def test_custom_mu(self) -> None:
        d = DirichletLMScorer(mu=500.0)
        assert d.mu == 500.0

    def test_total_corpus_tokens_initially_zero(self) -> None:
        d = DirichletLMScorer()
        assert d.total_corpus_tokens == 0


class TestDirichletScore:
    def _scorer(self) -> DirichletLMScorer:
        s = DirichletLMScorer(mu=100.0)
        s.update_stats(1000, {"hello": 50, "world": 30, "test": 20})
        return s

    def test_zero_corpus_returns_zero(self) -> None:
        s = DirichletLMScorer()
        assert s.score(["hello"], {"hello": 1}, 10) == 0.0

    def test_basic_score(self) -> None:
        s = self._scorer()
        score = s.score(["hello"], {"hello": 5}, 50)
        assert isinstance(score, float)

    def test_missing_term_in_corpus_skipped(self) -> None:
        s = self._scorer()
        score = s.score(["nonexistent"], {"nonexistent": 1}, 50)
        # Only the |q| * log(mu/(|d|+mu)) part remains
        expected = 1 * math.log(100.0 / (50 + 100.0))
        assert abs(score - expected) < 1e-10

    def test_higher_tf_gives_higher_score(self) -> None:
        s = self._scorer()
        low = s.score(["hello"], {"hello": 1}, 100)
        high = s.score(["hello"], {"hello": 10}, 100)
        assert high > low

    def test_update_stats(self) -> None:
        s = DirichletLMScorer()
        s.update_stats(5000, {"a": 100})
        assert s.total_corpus_tokens == 5000

    def test_multi_term_query(self) -> None:
        s = self._scorer()
        score = s.score(["hello", "world"], {"hello": 3, "world": 2}, 50)
        assert isinstance(score, float)


class TestDirichletNormalized:
    def test_returns_between_zero_and_one(self) -> None:
        s = DirichletLMScorer(mu=100.0)
        s.update_stats(1000, {"a": 50})
        n = s.score_normalized(["a"], {"a": 5}, 50)
        assert 0.0 <= n <= 1.0

    def test_zero_corpus_returns_half(self) -> None:
        s = DirichletLMScorer()
        n = s.score_normalized(["a"], {"a": 1}, 10)
        # raw=0 -> sigmoid(0) = 0.5
        assert abs(n - 0.5) < 1e-10

    def test_high_score_clamps(self) -> None:
        s = DirichletLMScorer(mu=10.0)
        s.update_stats(100, {"a": 10})
        # Very high tf relative to corpus should give high normalized
        n = s.score_normalized(["a"], {"a": 50}, 50)
        assert n >= 0.5


# ===================================================================
# JelinekMercerLMScorer
# ===================================================================


class TestJelinekMercerInit:
    def test_default_lambda(self) -> None:
        j = JelinekMercerLMScorer()
        assert j.lambda_param == 0.1

    def test_custom_lambda(self) -> None:
        j = JelinekMercerLMScorer(lambda_param=0.5)
        assert j.lambda_param == 0.5

    def test_invalid_lambda_zero(self) -> None:
        with pytest.raises(ValueError):
            JelinekMercerLMScorer(lambda_param=0.0)

    def test_invalid_lambda_one(self) -> None:
        with pytest.raises(ValueError):
            JelinekMercerLMScorer(lambda_param=1.0)

    def test_invalid_lambda_negative(self) -> None:
        with pytest.raises(ValueError):
            JelinekMercerLMScorer(lambda_param=-0.5)


class TestJelinekMercerScore:
    def _scorer(self) -> JelinekMercerLMScorer:
        s = JelinekMercerLMScorer(lambda_param=0.3)
        s.update_stats(1000, {"hello": 50, "world": 30})
        return s

    def test_zero_corpus_returns_zero(self) -> None:
        s = JelinekMercerLMScorer()
        assert s.score(["hello"], {"hello": 1}, 10) == 0.0

    def test_zero_doc_len_returns_zero(self) -> None:
        s = self._scorer()
        assert s.score(["hello"], {"hello": 1}, 0) == 0.0

    def test_basic_score(self) -> None:
        s = self._scorer()
        score = s.score(["hello"], {"hello": 5}, 50)
        assert isinstance(score, float)
        assert score < 0  # log-probabilities are negative

    def test_higher_tf_gives_higher_score(self) -> None:
        s = self._scorer()
        low = s.score(["hello"], {"hello": 1}, 100)
        high = s.score(["hello"], {"hello": 10}, 100)
        assert high > low

    def test_term_not_in_doc_nor_corpus(self) -> None:
        s = self._scorer()
        score = s.score(["xyz"], {}, 50)
        # prob = lambda*0 + (1-lambda)*0 = 0 → skipped → 0.0
        assert score == 0.0


class TestJelinekMercerNormalized:
    def test_returns_between_zero_and_one(self) -> None:
        s = JelinekMercerLMScorer(lambda_param=0.3)
        s.update_stats(1000, {"a": 50})
        n = s.score_normalized(["a"], {"a": 5}, 50)
        assert 0.0 <= n <= 1.0


# ===================================================================
# AbsoluteDiscountLMScorer
# ===================================================================


class TestAbsoluteDiscountInit:
    def test_default_delta(self) -> None:
        a = AbsoluteDiscountLMScorer()
        assert a.delta == 0.7

    def test_custom_delta(self) -> None:
        a = AbsoluteDiscountLMScorer(delta=0.3)
        assert a.delta == 0.3

    def test_invalid_delta_zero(self) -> None:
        with pytest.raises(ValueError):
            AbsoluteDiscountLMScorer(delta=0.0)

    def test_invalid_delta_one(self) -> None:
        with pytest.raises(ValueError):
            AbsoluteDiscountLMScorer(delta=1.0)

    def test_invalid_delta_negative(self) -> None:
        with pytest.raises(ValueError):
            AbsoluteDiscountLMScorer(delta=-0.1)


class TestAbsoluteDiscountScore:
    def _scorer(self) -> AbsoluteDiscountLMScorer:
        s = AbsoluteDiscountLMScorer(delta=0.5)
        s.update_stats(1000, {"hello": 50, "world": 30})
        return s

    def test_zero_corpus_returns_zero(self) -> None:
        s = AbsoluteDiscountLMScorer()
        assert s.score(["hello"], {"hello": 1}, 10) == 0.0

    def test_zero_doc_len_returns_zero(self) -> None:
        s = self._scorer()
        assert s.score(["hello"], {"hello": 1}, 0) == 0.0

    def test_basic_score(self) -> None:
        s = self._scorer()
        score = s.score(["hello"], {"hello": 5}, 50)
        assert isinstance(score, float)

    def test_higher_tf_gives_higher_score(self) -> None:
        s = self._scorer()
        low = s.score(["hello"], {"hello": 1}, 100)
        high = s.score(["hello"], {"hello": 10}, 100)
        assert high > low

    def test_term_not_in_doc_or_corpus(self) -> None:
        s = self._scorer()
        score = s.score(["xyz"], {}, 50)
        assert score == 0.0


class TestAbsoluteDiscountNormalized:
    def test_returns_between_zero_and_one(self) -> None:
        s = AbsoluteDiscountLMScorer(delta=0.5)
        s.update_stats(1000, {"a": 50})
        n = s.score_normalized(["a"], {"a": 5}, 50)
        assert 0.0 <= n <= 1.0


# ===================================================================
# Cross-scorer comparisons
# ===================================================================


class TestCrossScorer:
    def test_all_scorers_produce_different_scores(self) -> None:
        stats = (1000, {"hello": 50, "world": 30})
        q = ["hello"]
        tf = {"hello": 5}
        dl = 50

        d = DirichletLMScorer(mu=100.0)
        d.update_stats(*stats)

        j = JelinekMercerLMScorer(lambda_param=0.3)
        j.update_stats(*stats)

        a = AbsoluteDiscountLMScorer(delta=0.5)
        a.update_stats(*stats)

        scores = {d.score(q, tf, dl), j.score(q, tf, dl), a.score(q, tf, dl)}
        # All three methods should produce different scores
        assert len(scores) == 3

    def test_all_normalized_in_range(self) -> None:
        stats = (1000, {"hello": 50})
        q = ["hello"]
        tf = {"hello": 5}
        dl = 50

        for Cls in [DirichletLMScorer, JelinekMercerLMScorer, AbsoluteDiscountLMScorer]:
            if Cls is DirichletLMScorer:
                s = Cls(mu=100.0)
            elif Cls is JelinekMercerLMScorer:
                s = Cls(lambda_param=0.3)
            else:
                s = Cls(delta=0.5)
            s.update_stats(*stats)
            n = s.score_normalized(q, tf, dl)
            assert 0.0 <= n <= 1.0, f"{Cls.__name__} normalized out of range: {n}"

    def test_repeated_calls_consistent(self) -> None:
        s = DirichletLMScorer(mu=100.0)
        s.update_stats(1000, {"a": 50})
        s1 = s.score(["a"], {"a": 5}, 50)
        s2 = s.score(["a"], {"a": 5}, 50)
        assert s1 == s2
