"""Tests for rbsp.query.learning_to_rank — Learning-to-Rank framework."""

from __future__ import annotations

import math

import pytest

from rbsp.query.learning_to_rank import (
    BM25Feature,
    DocLengthFeature,
    Feature,
    FeatureExtractor,
    FeatureVector,
    FreshnessFeature,
    LinearRanker,
    QueryCoverageFeature,
    RankingPipeline,
    TFFeature,
)


# ===================================================================
# Feature / FeatureVector dataclasses
# ===================================================================


class TestFeature:
    def test_creation(self) -> None:
        f = Feature(name="bm25", value=3.14)
        assert f.name == "bm25"
        assert f.value == 3.14

    def test_frozen(self) -> None:
        f = Feature(name="bm25", value=1.0)
        with pytest.raises(AttributeError):
            f.name = "other"  # type: ignore[misc]


class TestFeatureVector:
    def test_creation(self) -> None:
        fv = FeatureVector(
            doc_id="d1",
            features=[Feature("a", 1.0), Feature("b", 2.0)],
            label=1.0,
        )
        assert fv.doc_id == "d1"
        assert len(fv.features) == 2
        assert fv.label == 1.0

    def test_default_label(self) -> None:
        fv = FeatureVector(doc_id="d1", features=[])
        assert fv.label == 0.0


# ===================================================================
# Feature extractors
# ===================================================================


class TestBM25Feature:
    def test_basic_extraction(self) -> None:
        ext = BM25Feature(k1=1.2, b=0.75)
        doc_data = {
            "term_freqs": {"hello": 3},
            "doc_len": 50,
            "avg_doc_len": 100.0,
            "num_docs": 1000,
            "doc_freqs": {"hello": 10},
        }
        score = ext.extract(["hello"], doc_data)
        assert score > 0

    def test_name(self) -> None:
        assert BM25Feature().name == "bm25"

    def test_missing_term_zero_score(self) -> None:
        ext = BM25Feature()
        doc_data = {
            "term_freqs": {},
            "doc_len": 50,
            "avg_doc_len": 100.0,
            "num_docs": 1000,
            "doc_freqs": {},
        }
        assert ext.extract(["hello"], doc_data) == 0.0

    def test_zero_avg_doc_len_handled(self) -> None:
        ext = BM25Feature()
        doc_data = {
            "term_freqs": {"a": 1},
            "doc_len": 10,
            "avg_doc_len": 0.0,
            "num_docs": 10,
            "doc_freqs": {"a": 1},
        }
        # Should not raise
        score = ext.extract(["a"], doc_data)
        assert isinstance(score, float)

    def test_multiple_query_terms(self) -> None:
        ext = BM25Feature()
        doc_data = {
            "term_freqs": {"a": 2, "b": 3},
            "doc_len": 100,
            "avg_doc_len": 100.0,
            "num_docs": 100,
            "doc_freqs": {"a": 5, "b": 10},
        }
        score = ext.extract(["a", "b"], doc_data)
        assert score > ext.extract(["a"], doc_data)


class TestFreshnessFeature:
    def test_name(self) -> None:
        assert FreshnessFeature().name == "freshness"

    def test_current_doc(self) -> None:
        ext = FreshnessFeature(halflife=86400)
        score = ext.extract([], {"timestamp": 1000.0, "now": 1000.0})
        assert abs(score - 1.0) < 1e-10

    def test_old_doc(self) -> None:
        ext = FreshnessFeature(halflife=100)
        score = ext.extract([], {"timestamp": 0.0, "now": 100.0})
        assert abs(score - 0.5) < 0.01

    def test_zero_halflife(self) -> None:
        ext = FreshnessFeature(halflife=0)
        assert ext.extract([], {"timestamp": 1000.0, "now": 1000.0}) == 1.0
        assert ext.extract([], {"timestamp": 999.0, "now": 1000.0}) == 0.0


class TestDocLengthFeature:
    def test_name(self) -> None:
        assert DocLengthFeature().name == "doc_length"

    def test_extraction(self) -> None:
        ext = DocLengthFeature()
        assert ext.extract([], {"doc_len": 42}) == 42.0

    def test_missing_key(self) -> None:
        ext = DocLengthFeature()
        assert ext.extract([], {}) == 0.0


class TestQueryCoverageFeature:
    def test_name(self) -> None:
        assert QueryCoverageFeature().name == "query_coverage"

    def test_full_coverage(self) -> None:
        ext = QueryCoverageFeature()
        doc_data = {"term_freqs": {"a": 1, "b": 1}}
        assert ext.extract(["a", "b"], doc_data) == 1.0

    def test_no_coverage(self) -> None:
        ext = QueryCoverageFeature()
        doc_data = {"term_freqs": {"x": 1}}
        assert ext.extract(["a", "b"], doc_data) == 0.0

    def test_partial_coverage(self) -> None:
        ext = QueryCoverageFeature()
        doc_data = {"term_freqs": {"a": 1}}
        assert ext.extract(["a", "b"], doc_data) == 0.5

    def test_empty_query(self) -> None:
        ext = QueryCoverageFeature()
        assert ext.extract([], {"term_freqs": {"a": 1}}) == 0.0

    def test_dedup_query_terms(self) -> None:
        ext = QueryCoverageFeature()
        doc_data = {"term_freqs": {"a": 1}}
        # "a" appears twice but unique coverage is 1/1 = 1.0
        assert ext.extract(["a", "a"], doc_data) == 1.0


class TestTFFeature:
    def test_name(self) -> None:
        assert TFFeature().name == "tf"

    def test_basic(self) -> None:
        ext = TFFeature()
        doc_data = {"term_freqs": {"a": 3, "b": 5}}
        assert ext.extract(["a", "b"], doc_data) == 8.0

    def test_missing_term(self) -> None:
        ext = TFFeature()
        doc_data = {"term_freqs": {"a": 3}}
        assert ext.extract(["a", "z"], doc_data) == 3.0

    def test_empty_query(self) -> None:
        ext = TFFeature()
        assert ext.extract([], {"term_freqs": {"a": 1}}) == 0.0


class TestFeatureExtractorBase:
    def test_raises_not_implemented(self) -> None:
        ext = FeatureExtractor()
        with pytest.raises(NotImplementedError):
            ext.extract([], {})


# ===================================================================
# LinearRanker
# ===================================================================


class TestLinearRanker:
    def test_score_basic(self) -> None:
        r = LinearRanker(weights={"a": 2.0, "b": 3.0})
        features = [Feature("a", 1.0), Feature("b", 2.0)]
        assert r.score(features) == 8.0  # 2*1 + 3*2

    def test_score_missing_weight(self) -> None:
        r = LinearRanker(weights={"a": 1.0})
        features = [Feature("a", 1.0), Feature("b", 2.0)]
        assert r.score(features) == 1.0  # b has weight 0

    def test_empty_features(self) -> None:
        r = LinearRanker(weights={"a": 1.0})
        assert r.score([]) == 0.0

    def test_no_weights(self) -> None:
        r = LinearRanker()
        assert r.score([Feature("a", 5.0)]) == 0.0

    def test_set_weights(self) -> None:
        r = LinearRanker()
        r.set_weights({"a": 1.0})
        assert r.weights == {"a": 1.0}

    def test_weights_property_returns_copy(self) -> None:
        r = LinearRanker(weights={"a": 1.0})
        w = r.weights
        w["a"] = 999.0
        assert r.weights["a"] == 1.0  # original unchanged


class TestLinearRankerTraining:
    def test_train_basic(self) -> None:
        data = [
            FeatureVector("d1", [Feature("f1", 1.0)], label=2.0),
            FeatureVector("d2", [Feature("f1", 0.0)], label=0.0),
        ]
        r = LinearRanker()
        r.train(data, learning_rate=0.1, epochs=50)
        # After training, f1 should have positive weight
        assert r.weights.get("f1", 0.0) > 0

    def test_train_fewer_than_two_samples(self) -> None:
        r = LinearRanker()
        r.train([FeatureVector("d1", [Feature("f1", 1.0)], label=1.0)])
        assert r.weights == {}

    def test_train_equal_labels_no_change(self) -> None:
        data = [
            FeatureVector("d1", [Feature("f1", 1.0)], label=1.0),
            FeatureVector("d2", [Feature("f1", 2.0)], label=1.0),
        ]
        r = LinearRanker()
        r.train(data, epochs=10)
        # Equal labels → no pairs → weights stay 0
        assert r.weights.get("f1", 0.0) == 0.0

    def test_train_learns_correct_direction(self) -> None:
        data = [
            FeatureVector("good", [Feature("quality", 10.0), Feature("spam", 0.0)], label=2.0),
            FeatureVector("bad", [Feature("quality", 0.0), Feature("spam", 10.0)], label=0.0),
        ]
        r = LinearRanker()
        r.train(data, learning_rate=0.01, epochs=200)
        # quality should get positive weight, spam negative
        assert r.weights["quality"] > r.weights["spam"]


# ===================================================================
# RankingPipeline
# ===================================================================


class TestRankingPipeline:
    def _make_pipeline(self) -> RankingPipeline:
        extractors = [DocLengthFeature(), TFFeature()]
        ranker = LinearRanker(weights={"doc_length": 0.0, "tf": 1.0})
        return RankingPipeline(extractors, ranker)

    def test_extract_features(self) -> None:
        pipe = self._make_pipeline()
        features = pipe.extract_features(
            ["a"], {"term_freqs": {"a": 3}, "doc_len": 50}
        )
        assert len(features) == 2
        names = {f.name for f in features}
        assert "doc_length" in names
        assert "tf" in names

    def test_score(self) -> None:
        pipe = self._make_pipeline()
        s = pipe.score(["a"], {"term_freqs": {"a": 5}, "doc_len": 100})
        # tf weight=1.0 → score = 5.0, doc_length weight=0 → +0
        assert abs(s - 5.0) < 1e-10

    def test_rerank(self) -> None:
        pipe = self._make_pipeline()
        candidates = [
            {"doc_id": "d1", "term_freqs": {"a": 1}, "doc_len": 10},
            {"doc_id": "d2", "term_freqs": {"a": 5}, "doc_len": 20},
            {"doc_id": "d3", "term_freqs": {"a": 3}, "doc_len": 15},
        ]
        ranked = pipe.rerank(["a"], candidates)
        ids = [doc_id for doc_id, _ in ranked]
        assert ids == ["d2", "d3", "d1"]  # sorted by tf descending

    def test_rerank_empty(self) -> None:
        pipe = self._make_pipeline()
        assert pipe.rerank(["a"], []) == []

    def test_rerank_single_candidate(self) -> None:
        pipe = self._make_pipeline()
        candidates = [{"doc_id": "d1", "term_freqs": {"a": 1}, "doc_len": 10}]
        ranked = pipe.rerank(["a"], candidates)
        assert len(ranked) == 1
        assert ranked[0][0] == "d1"
