"""BM25F Field-Aware Scorer for RBSP.

Extends Okapi BM25 to handle multiple document fields (filename, content,
comments, etc.) with per-field weights and length normalization.  Used by
Elasticsearch and Tantivy for structured document ranking.

BM25F formula:
    For each query term t:
        tf_hat(t) = SUM_f ( weight_f * tf(t,f) / (1 - b_f + b_f * dl_f / avgdl_f) )
        score(t)  = IDF(t) * tf_hat(t) / (k1 + tf_hat(t))

    Total score = SUM_t score(t)

where:
    IDF(t)   = ln((N - df(t) + 0.5) / (df(t) + 0.5) + 1)
    tf(t,f)  = frequency of term t in field f
    dl_f     = length of field f in the document
    avgdl_f  = average length of field f across the corpus
    weight_f = boost weight for field f
    b_f      = length normalization for field f (0=none, 1=full)
    k1       = term frequency saturation parameter
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Dict, List, Optional


@dataclass(frozen=True, slots=True)
class FieldConfig:
    """Configuration for a single document field.

    Attributes:
        name:   Field identifier (e.g. "filename", "content").
        weight: Boost multiplier for this field. Higher values make
                matches in this field more important.
        b:      Length normalization parameter for this field.
                0 = no normalization, 1 = full normalization.
    """
    name: str
    weight: float = 1.0
    b: float = 0.75


_DEFAULT_FIELDS: List[FieldConfig] = [
    FieldConfig("filename", weight=3.0, b=0.3),
    FieldConfig("content", weight=1.0, b=0.75),
]


class BM25FScorer:
    """BM25F multi-field relevance scorer.

    Computes BM25F scores where each document is treated as a collection
    of fields, each with its own weight and length normalization.  A match
    in a high-weight field (e.g. filename) scores higher than the same
    match in a low-weight field (e.g. content body).

    Args:
        k1:     Term frequency saturation parameter. Higher values give
                more weight to repeated terms. Default 1.2 (standard).
        fields: List of FieldConfig defining the fields to score over.
                Default: filename (weight=3.0, b=0.3) and content
                (weight=1.0, b=0.75).
    """

    def __init__(
        self,
        k1: float = 1.2,
        fields: Optional[List[FieldConfig]] = None,
    ) -> None:
        self.k1 = k1
        self.fields = list(fields) if fields is not None else list(_DEFAULT_FIELDS)
        self._field_map: Dict[str, FieldConfig] = {f.name: f for f in self.fields}
        self._num_docs: int = 0
        self._field_avg_lengths: Dict[str, float] = {}
        self._doc_freqs: Dict[str, int] = {}

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

    @property
    def field_avg_lengths(self) -> Dict[str, float]:
        return dict(self._field_avg_lengths)

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

        Args:
            num_docs:          Total documents in corpus.
            field_avg_lengths: {field_name: average_field_length} across corpus.
            doc_freqs:         {term: number_of_documents_containing_term}.
        """
        self._num_docs = num_docs
        self._field_avg_lengths = {
            k: v if v > 0 else 1.0 for k, v in field_avg_lengths.items()
        }
        self._doc_freqs = doc_freqs

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

        Uses the Robertson-Sparck Jones IDF variant with +1 smoothing:
            IDF(t) = ln((N - df + 0.5) / (df + 0.5) + 1)

        Args:
            term: The query term.

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

    def _weighted_tf(
        self,
        term: str,
        doc_fields: Dict[str, Dict[str, int]],
        doc_field_lengths: Dict[str, int],
    ) -> float:
        """Compute the BM25F weighted pseudo term-frequency.

        tf_hat = SUM_f ( weight_f * tf_f / (1 - b_f + b_f * dl_f / avgdl_f) )

        Args:
            term:              The query term.
            doc_fields:        {field_name: {term: tf}} for the document.
            doc_field_lengths: {field_name: length} for the document.

        Returns:
            Weighted pseudo-TF across all fields.
        """
        tf_hat = 0.0
        for fc in self.fields:
            field_tfs = doc_fields.get(fc.name)
            if field_tfs is None:
                continue
            tf = field_tfs.get(term, 0)
            if tf <= 0:
                continue
            dl = doc_field_lengths.get(fc.name, 0)
            avgdl = self._field_avg_lengths.get(fc.name, 1.0)
            norm = 1.0 - fc.b + fc.b * dl / avgdl
            tf_hat += fc.weight * tf / norm
        return tf_hat

    def score(
        self,
        query_terms: List[str],
        doc_fields: Dict[str, Dict[str, int]],
        doc_field_lengths: Dict[str, int],
    ) -> float:
        """Score a document against a query using BM25F.

        Args:
            query_terms:       List of query terms.
            doc_fields:        {field_name: {term: tf}} for the document.
            doc_field_lengths: {field_name: length} for the document.

        Returns:
            BM25F relevance score (non-negative, unbounded).
        """
        total = 0.0
        for term in query_terms:
            tf_hat = self._weighted_tf(term, doc_fields, doc_field_lengths)
            if tf_hat <= 0.0:
                continue
            idf_val = self.idf(term)
            total += idf_val * tf_hat / (self.k1 + tf_hat)
        return total

    def score_upper_bound(self, query_terms: List[str]) -> float:
        """Compute an upper bound on the BM25F score.

        As tf_hat -> infinity, each term's contribution approaches IDF(t).
        Used for WAND-style early termination.

        Args:
            query_terms: List of query terms.

        Returns:
            Upper bound BM25F score.
        """
        total = 0.0
        for term in query_terms:
            total += self.idf(term)
        return total

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

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

        Args:
            query_terms:       List of query terms.
            doc_fields:        {field_name: {term: tf}} for the document.
            doc_field_lengths: {field_name: length} for the document.

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