"""Learning-to-Rank (LTR) framework for RBSP.

Combines multiple scoring signals (BM25, freshness, proximity, etc.) into
a single ranking model.  Modern search engines use LTR to achieve their
best results by learning optimal feature weights from relevance data.

Components:
    - Feature / FeatureVector:   Structured feature representation.
    - FeatureExtractor:          Base class for extracting a feature value.
    - BM25Feature:               BM25 score as a ranking feature.
    - FreshnessFeature:          Recency decay as a ranking feature.
    - DocLengthFeature:          Document length as a ranking feature.
    - QueryCoverageFeature:      Fraction of query terms present in doc.
    - TFFeature:                 Sum of raw term frequencies.
    - LinearRanker:              Weighted linear combination with pairwise
                                 gradient descent training.
    - RankingPipeline:           End-to-end extract -> score -> rerank.
"""

from __future__ import annotations

import math
import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------

@dataclass(frozen=True, slots=True)
class Feature:
    """A single named feature value.

    Attributes:
        name:  Identifier for this feature (e.g. "bm25", "freshness").
        value: Numeric feature value.
    """
    name: str
    value: float


@dataclass(slots=True)
class FeatureVector:
    """Feature vector for a single document.

    Attributes:
        doc_id:   Unique document identifier.
        features: Extracted feature values for this document.
        label:    Relevance label for training (higher = more relevant).
    """
    doc_id: str
    features: list[Feature]
    label: float = 0.0


# ---------------------------------------------------------------------------
# Feature extractors
# ---------------------------------------------------------------------------

class FeatureExtractor:
    """Base class for feature extraction.

    Subclasses implement ``extract`` to compute a single numeric feature
    from query terms and document data.
    """

    name: str = "base"

    def extract(self, query_terms: list[str], doc_data: dict) -> float:
        """Extract a feature value.

        Args:
            query_terms: Tokenized query terms.
            doc_data:    Dictionary with document fields.  Expected keys
                         depend on the concrete extractor.

        Returns:
            Numeric feature value.
        """
        raise NotImplementedError


class BM25Feature(FeatureExtractor):
    """BM25 relevance score as a feature.

    Expected ``doc_data`` keys:
        - ``term_freqs``: ``Dict[str, int]`` — term -> count in doc.
        - ``doc_len``:    ``int`` — number of tokens in the document.
        - ``avg_doc_len``: ``float`` — average document length in corpus.
        - ``num_docs``:   ``int`` — total documents in corpus.
        - ``doc_freqs``:  ``Dict[str, int]`` — term -> number of docs containing term.
    """

    name: str = "bm25"

    def __init__(self, k1: float = 1.2, b: float = 0.75) -> None:
        self.k1 = k1
        self.b = b

    def extract(self, query_terms: list[str], doc_data: dict) -> float:
        term_freqs: Dict[str, int] = doc_data.get("term_freqs", {})
        doc_len: int = doc_data.get("doc_len", 0)
        avg_doc_len: float = doc_data.get("avg_doc_len", 1.0)
        num_docs: int = doc_data.get("num_docs", 1)
        corpus_df: Dict[str, int] = doc_data.get("doc_freqs", {})

        if avg_doc_len <= 0:
            avg_doc_len = 1.0

        total = 0.0
        for term in query_terms:
            tf = term_freqs.get(term, 0)
            if tf <= 0:
                continue
            df = corpus_df.get(term, 0)
            idf = math.log((num_docs - df + 0.5) / (df + 0.5) + 1.0)
            numerator = tf * (self.k1 + 1.0)
            denominator = tf + self.k1 * (
                1.0 - self.b + self.b * doc_len / avg_doc_len
            )
            total += idf * (numerator / denominator)
        return total


class FreshnessFeature(FeatureExtractor):
    """Recency as a feature via exponential decay.

    Expected ``doc_data`` keys:
        - ``timestamp``: ``float`` — document creation/update epoch.

    Optionally supply ``now`` in ``doc_data`` to override the current time.
    """

    name: str = "freshness"

    def __init__(self, halflife: float = 86400.0) -> None:
        self.halflife = halflife

    def extract(self, query_terms: list[str], doc_data: dict) -> float:
        timestamp: float = doc_data.get("timestamp", 0.0)
        now: float = doc_data.get("now", time.time())
        age = max(now - timestamp, 0.0)
        if self.halflife <= 0:
            return 1.0 if age == 0 else 0.0
        return 2.0 ** (-age / self.halflife)


class DocLengthFeature(FeatureExtractor):
    """Document length as a feature.

    Returns ``doc_len`` directly — useful when combined with other features
    so the ranker can learn whether longer or shorter docs are preferred.

    Expected ``doc_data`` keys:
        - ``doc_len``: ``int``
    """

    name: str = "doc_length"

    def extract(self, query_terms: list[str], doc_data: dict) -> float:
        return float(doc_data.get("doc_len", 0))


class QueryCoverageFeature(FeatureExtractor):
    """Fraction of unique query terms that appear in the document.

    Expected ``doc_data`` keys:
        - ``term_freqs``: ``Dict[str, int]``
    """

    name: str = "query_coverage"

    def extract(self, query_terms: list[str], doc_data: dict) -> float:
        if not query_terms:
            return 0.0
        term_freqs: Dict[str, int] = doc_data.get("term_freqs", {})
        unique_terms = set(query_terms)
        matched = sum(1 for t in unique_terms if term_freqs.get(t, 0) > 0)
        return matched / len(unique_terms)


class TFFeature(FeatureExtractor):
    """Sum of raw term frequencies for query terms.

    Expected ``doc_data`` keys:
        - ``term_freqs``: ``Dict[str, int]``
    """

    name: str = "tf"

    def extract(self, query_terms: list[str], doc_data: dict) -> float:
        term_freqs: Dict[str, int] = doc_data.get("term_freqs", {})
        total = 0
        for term in query_terms:
            total += term_freqs.get(term, 0)
        return float(total)


# ---------------------------------------------------------------------------
# Linear ranker with pairwise training
# ---------------------------------------------------------------------------

class LinearRanker:
    """Weighted linear combination of features.

    Supports manual weight setting and simple pairwise gradient descent
    training to learn optimal weights from labeled data.

    Args:
        weights: Optional mapping of feature name to weight.  Features
                 not present in the map receive a default weight of 0.0.
    """

    def __init__(self, weights: Optional[dict[str, float]] = None) -> None:
        self._weights: dict[str, float] = dict(weights) if weights else {}

    @property
    def weights(self) -> dict[str, float]:
        return dict(self._weights)

    def set_weights(self, weights: dict[str, float]) -> None:
        """Replace weights entirely."""
        self._weights = dict(weights)

    def score(self, features: list[Feature]) -> float:
        """Compute weighted linear combination of features.

        Args:
            features: List of Feature values.

        Returns:
            Weighted sum.  Features missing from the weight map contribute 0.
        """
        total = 0.0
        for f in features:
            total += self._weights.get(f.name, 0.0) * f.value
        return total

    def train(
        self,
        training_data: list[FeatureVector],
        learning_rate: float = 0.01,
        epochs: int = 100,
    ) -> None:
        """Train weights via pairwise gradient descent.

        Minimizes pairwise hinge loss: for every pair (i, j) where
        label_i > label_j, the model should score i higher than j.

        Loss for a violating pair: max(0, 1 - (score_i - score_j))
        Gradient: pushes weights toward (features_i - features_j).

        Args:
            training_data:  List of FeatureVector with relevance labels.
            learning_rate:  Step size for gradient updates.
            epochs:         Number of full passes over all pairs.
        """
        if len(training_data) < 2:
            return

        # Collect all feature names and initialize missing weights
        all_names: set[str] = set()
        for fv in training_data:
            for f in fv.features:
                all_names.add(f.name)
        for name in all_names:
            if name not in self._weights:
                self._weights[name] = 0.0

        # Build pairs: (higher_label_vec, lower_label_vec)
        pairs: list[tuple[FeatureVector, FeatureVector]] = []
        for i in range(len(training_data)):
            for j in range(i + 1, len(training_data)):
                a, b = training_data[i], training_data[j]
                if a.label > b.label:
                    pairs.append((a, b))
                elif b.label > a.label:
                    pairs.append((b, a))
                # Equal labels — no preference, skip

        if not pairs:
            return

        for _epoch in range(epochs):
            for higher, lower in pairs:
                score_h = self.score(higher.features)
                score_l = self.score(lower.features)
                margin = score_h - score_l

                if margin < 1.0:
                    # Build feature diffs
                    h_map: dict[str, float] = {
                        f.name: f.value for f in higher.features
                    }
                    l_map: dict[str, float] = {
                        f.name: f.value for f in lower.features
                    }
                    for name in all_names:
                        diff = h_map.get(name, 0.0) - l_map.get(name, 0.0)
                        self._weights[name] += learning_rate * diff


# ---------------------------------------------------------------------------
# Ranking pipeline
# ---------------------------------------------------------------------------

class RankingPipeline:
    """End-to-end feature extraction, scoring, and reranking.

    Args:
        extractors: Ordered list of feature extractors.
        ranker:     A LinearRanker instance for scoring.
    """

    def __init__(
        self,
        extractors: list[FeatureExtractor],
        ranker: LinearRanker,
    ) -> None:
        self.extractors = extractors
        self.ranker = ranker

    def extract_features(
        self,
        query_terms: list[str],
        doc_data: dict,
    ) -> list[Feature]:
        """Run all extractors on a single document.

        Args:
            query_terms: Tokenized query.
            doc_data:    Document data dictionary.

        Returns:
            List of Feature values, one per extractor.
        """
        features: list[Feature] = []
        for ext in self.extractors:
            value = ext.extract(query_terms, doc_data)
            features.append(Feature(name=ext.name, value=value))
        return features

    def score(
        self,
        query_terms: list[str],
        doc_data: dict,
    ) -> float:
        """Extract features and score a single document.

        Args:
            query_terms: Tokenized query.
            doc_data:    Document data dictionary.

        Returns:
            Linear ranker score.
        """
        features = self.extract_features(query_terms, doc_data)
        return self.ranker.score(features)

    def rerank(
        self,
        query_terms: list[str],
        candidates: list[dict],
    ) -> list[tuple[str, float]]:
        """Score and rerank a list of candidate documents.

        Each candidate dict must include a ``"doc_id"`` key.  All other
        keys are passed through as ``doc_data`` to the extractors.

        Args:
            query_terms: Tokenized query.
            candidates:  List of candidate document dicts.

        Returns:
            List of (doc_id, score) tuples sorted by descending score.
        """
        scored: list[tuple[float, int, str]] = []
        for idx, candidate in enumerate(candidates):
            doc_id = candidate.get("doc_id", "")
            s = self.score(query_terms, candidate)
            scored.append((s, idx, doc_id))

        # Sort descending by score, stable on insertion order
        scored.sort(key=lambda x: (-x[0], x[1]))

        return [(doc_id, s) for s, _idx, doc_id in scored]
