"""Okapi BM25 Scorer for RBSP.

State-of-the-art probabilistic ranking function for information retrieval.
Accounts for term frequency saturation, inverse document frequency, and
document length normalization.

Formula:
    score(q, d) = Σ IDF(t) * (tf(t,d) * (k1 + 1)) / (tf(t,d) + k1 * (1 - b + b * |d| / avgdl))

where:
    IDF(t) = ln((N - df(t) + 0.5) / (df(t) + 0.5) + 1)
    tf(t,d) = frequency of term t in document d
    |d| = length of document d (in tokens)
    avgdl = average document length across corpus
    k1 = term frequency saturation (default 1.2)
    b = length normalization (default 0.75)
"""

from __future__ import annotations

import math
from typing import Dict, Iterable, List


def compute_term_freqs(tokens: Iterable[str]) -> Dict[str, int]:
    """Count term frequencies from a token iterable.

    Args:
        tokens: Iterable of string tokens.

    Returns:
        Dict mapping each unique token to its count.
    """
    freqs: Dict[str, int] = {}
    for t in tokens:
        freqs[t] = freqs.get(t, 0) + 1
    return freqs


class BM25Scorer:
    """Okapi BM25 relevance scorer.

    Computes BM25 scores given pre-computed corpus statistics.  Designed
    to be updated incrementally as documents are added to the index.

    Args:
        k1: Term frequency saturation parameter.  Higher values give
            more weight to repeated terms. Default 1.2 (standard).
        b:  Document length normalization.  0 = no normalization,
            1 = full normalization.  Default 0.75 (standard).
    """

    def __init__(self, k1: float = 1.2, b: float = 0.75) -> None:
        self.k1 = k1
        self.b = b
        self._num_docs: int = 0
        self._avg_doc_len: float = 0.0
        self._total_doc_len: int = 0
        self._doc_freqs: Dict[str, int] = {}

    @property
    def num_docs(self) -> int:
        return self._num_docs

    @property
    def avg_doc_len(self) -> float:
        return self._avg_doc_len

    def update_stats(
        self,
        num_docs: int,
        avg_doc_len: float,
        doc_freqs: Dict[str, int],
    ) -> None:
        """Bulk-update corpus statistics.

        Called after indexing to set IDF values.

        Args:
            num_docs:    Total documents in corpus.
            avg_doc_len: Average document length (tokens).
            doc_freqs:   {term: number_of_documents_containing_term}.
        """
        self._num_docs = num_docs
        self._avg_doc_len = avg_doc_len if avg_doc_len > 0 else 1.0
        self._doc_freqs = doc_freqs

    def add_document(self, doc_len: int, term_freqs: Dict[str, int]) -> None:
        """Incrementally update stats when a single document is added.

        Args:
            doc_len:    Number of tokens in the document.
            term_freqs: {term: count} for the document.
        """
        self._total_doc_len += doc_len
        self._num_docs += 1
        self._avg_doc_len = self._total_doc_len / self._num_docs
        for term in term_freqs:
            self._doc_freqs[term] = self._doc_freqs.get(term, 0) + 1

    def idf(self, term: str) -> float:
        """Compute Inverse Document Frequency for a term.

        Uses the Robertson-Sparck Jones IDF variant with +1 smoothing
        to avoid negative values:

            IDF(t) = ln((N - df + 0.5) / (df + 0.5) + 1)

        Args:
            term: The query term.

        Returns:
            Non-negative IDF value.  Rare terms score higher.
        """
        n = self._num_docs
        df = self._doc_freqs.get(term, 0)
        return math.log((n - df + 0.5) / (df + 0.5) + 1.0)

    def term_score(
        self,
        term: str,
        tf: int,
        doc_len: int,
    ) -> float:
        """Score a single term's contribution to a document.

        Args:
            term:    The query term (for IDF lookup).
            tf:      Term frequency in the document.
            doc_len: Total tokens in the document.

        Returns:
            BM25 contribution of this term.
        """
        if tf <= 0:
            return 0.0
        idf_val = self.idf(term)
        numerator = tf * (self.k1 + 1.0)
        denominator = tf + self.k1 * (
            1.0 - self.b + self.b * doc_len / self._avg_doc_len
        )
        return idf_val * (numerator / denominator)

    def score(
        self,
        query_terms: List[str],
        doc_term_freqs: Dict[str, int],
        doc_len: int,
    ) -> float:
        """Score a document against a query using full BM25.

        Args:
            query_terms:    List of query terms (duplicates boost weight).
            doc_term_freqs: {term: count} for the candidate document.
            doc_len:        Total tokens in the document.

        Returns:
            BM25 relevance score (non-negative, unbounded).
        """
        total = 0.0
        for term in query_terms:
            tf = doc_term_freqs.get(term, 0)
            if tf > 0:
                total += self.term_score(term, tf, doc_len)
        return total

    def score_upper_bound(
        self,
        query_terms: List[str],
        doc_len: int,
    ) -> float:
        """Compute an upper bound on the BM25 score for WAND pruning.

        Assumes maximum possible TF for each query term.  Used to skip
        documents that cannot possibly reach the current top-K threshold.

        Args:
            query_terms: List of query terms.
            doc_len:     Total tokens in the document.

        Returns:
            Upper bound BM25 score.
        """
        total = 0.0
        for term in query_terms:
            idf_val = self.idf(term)
            # Maximum TF contribution approaches idf * (k1 + 1) as tf -> inf
            max_tf_contribution = idf_val * (self.k1 + 1.0)
            total += max_tf_contribution
        return total

    def score_normalized(
        self,
        query_terms: List[str],
        doc_term_freqs: Dict[str, int],
        doc_len: int,
    ) -> float:
        """BM25 score normalized to approximately [0, 1].

        Divides by the theoretical maximum (all terms present with
        saturated TF).  Useful for combining with other signals.

        Args:
            query_terms:    List of query terms.
            doc_term_freqs: {term: count} for the document.
            doc_len:        Total tokens in the document.

        Returns:
            Normalized score in approximately [0, 1].
        """
        raw = self.score(query_terms, doc_term_freqs, doc_len)
        upper = self.score_upper_bound(query_terms, doc_len)
        if upper <= 0:
            return 0.0
        return min(raw / upper, 1.0)
