"""Language Model Scorers for RBSP.

Statistical language model approaches to document ranking.  Instead of
BM25's term-frequency saturation, these scorers estimate the probability
that a query was generated by a document's language model, smoothed with
corpus-level statistics to handle unseen terms.

Three smoothing variants are provided:

    DirichletLMScorer        -- Bayesian smoothing with Dirichlet prior (default)
    JelinekMercerLMScorer    -- Fixed interpolation between document and corpus
    AbsoluteDiscountLMScorer -- Subtracts a fixed discount from observed counts

All methods use ``math.log`` (natural logarithm) from the standard library.
Zero external dependencies.

References:
    Zhai & Lafferty, "A Study of Smoothing Methods for Language Models
    Applied to Information Retrieval" (ACM TOIS, 2004).
"""

from __future__ import annotations

import math
from typing import Dict, List


class DirichletLMScorer:
    """Dirichlet-smoothed language model scorer.

    Smooths document term probabilities with a Dirichlet prior based on
    corpus term frequencies.  The smoothing parameter ``mu`` controls how
    much weight is given to the corpus model vs. the document model.
    Larger ``mu`` values trust the corpus more (useful for short documents).

    Score formula for query q = (q1, ..., qn) against document d::

        score(q, d) = SUM_i log(1 + tf(qi, d) / (mu * cf(qi) / T))
                    + |q| * log(mu / (|d| + mu))

    where:
        tf(qi, d) = frequency of term qi in document d
        cf(qi)    = collection frequency of term qi (total occurrences)
        T         = total tokens in the corpus
        |d|       = document length (tokens)
        |q|       = number of query terms
        mu        = Dirichlet smoothing parameter

    Args:
        mu: Dirichlet smoothing parameter (default 2000.0).
    """

    def __init__(self, mu: float = 2000.0) -> None:
        self.mu = mu
        self._total_corpus_tokens: int = 0
        self._corpus_term_freqs: Dict[str, int] = {}

    @property
    def total_corpus_tokens(self) -> int:
        return self._total_corpus_tokens

    def update_stats(
        self,
        total_corpus_tokens: int,
        corpus_term_freqs: Dict[str, int],
    ) -> None:
        """Set collection-level statistics.

        Args:
            total_corpus_tokens: Total number of tokens across all documents.
            corpus_term_freqs:   {term: total_occurrences_in_corpus}.
        """
        self._total_corpus_tokens = total_corpus_tokens
        self._corpus_term_freqs = corpus_term_freqs

    def score(
        self,
        query_terms: List[str],
        doc_term_freqs: Dict[str, int],
        doc_len: int,
    ) -> float:
        """Score a document using Dirichlet-smoothed language model.

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

        Returns:
            Log-probability score (higher is better, typically negative).
        """
        t = self._total_corpus_tokens
        if t <= 0:
            return 0.0

        mu = self.mu
        total = 0.0

        for term in query_terms:
            tf = doc_term_freqs.get(term, 0)
            cf = self._corpus_term_freqs.get(term, 0)
            if cf <= 0:
                continue
            # log(1 + tf / (mu * cf / T))
            total += math.log(1.0 + tf / (mu * cf / t))

        # |q| * log(mu / (doc_len + mu))
        total += len(query_terms) * math.log(mu / (doc_len + mu))

        return total

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

        Maps the raw log-probability score into [0, 1] using a sigmoid-style
        transformation.  A score of 0.5 corresponds to the baseline (all
        terms at corpus frequency).

        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 [0, 1].
        """
        raw = self.score(query_terms, doc_term_freqs, doc_len)
        # Sigmoid mapping: 1 / (1 + e^(-raw))
        # Clamp input to avoid overflow
        if raw > 500.0:
            return 1.0
        if raw < -500.0:
            return 0.0
        return 1.0 / (1.0 + math.exp(-raw))


class JelinekMercerLMScorer:
    """Jelinek-Mercer smoothed language model scorer.

    Uses fixed linear interpolation between the document model and the
    corpus model.  The parameter ``lambda_param`` controls the mix:
    higher values trust the document model more.

    Score formula for query q against document d::

        score(q, d) = SUM_i log(lambda * tf(qi, d) / |d|
                                + (1 - lambda) * cf(qi) / T)

    where:
        lambda = interpolation parameter (weight on document model)

    Args:
        lambda_param: Interpolation weight on the document model.
                      Must be in (0, 1). Default 0.1.
    """

    def __init__(self, lambda_param: float = 0.1) -> None:
        if not (0.0 < lambda_param < 1.0):
            raise ValueError("lambda_param must be in (0, 1)")
        self.lambda_param = lambda_param
        self._total_corpus_tokens: int = 0
        self._corpus_term_freqs: Dict[str, int] = {}

    @property
    def total_corpus_tokens(self) -> int:
        return self._total_corpus_tokens

    def update_stats(
        self,
        total_corpus_tokens: int,
        corpus_term_freqs: Dict[str, int],
    ) -> None:
        """Set collection-level statistics.

        Args:
            total_corpus_tokens: Total number of tokens across all documents.
            corpus_term_freqs:   {term: total_occurrences_in_corpus}.
        """
        self._total_corpus_tokens = total_corpus_tokens
        self._corpus_term_freqs = corpus_term_freqs

    def score(
        self,
        query_terms: List[str],
        doc_term_freqs: Dict[str, int],
        doc_len: int,
    ) -> float:
        """Score a document using Jelinek-Mercer smoothed language model.

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

        Returns:
            Log-probability score (sum of log-probabilities per term).
        """
        t = self._total_corpus_tokens
        if t <= 0 or doc_len <= 0:
            return 0.0

        lam = self.lambda_param
        total = 0.0

        for term in query_terms:
            tf = doc_term_freqs.get(term, 0)
            cf = self._corpus_term_freqs.get(term, 0)

            # P(term | document) interpolated with P(term | corpus)
            p_doc = tf / doc_len
            p_corpus = cf / t

            prob = lam * p_doc + (1.0 - lam) * p_corpus
            if prob <= 0.0:
                continue
            total += math.log(prob)

        return total

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

        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 [0, 1].
        """
        raw = self.score(query_terms, doc_term_freqs, doc_len)
        if raw > 500.0:
            return 1.0
        if raw < -500.0:
            return 0.0
        return 1.0 / (1.0 + math.exp(-raw))


class AbsoluteDiscountLMScorer:
    """Absolute discount smoothed language model scorer.

    Subtracts a fixed discount ``delta`` from each observed term count and
    redistributes the probability mass to unseen terms proportional to their
    corpus frequency.

    Score formula for query q against document d::

        score(q, d) = SUM_i log(max(tf(qi, d) - delta, 0) / |d|
                                + delta * |d_unique| / |d| * cf(qi) / T)

    where:
        delta      = absolute discount value
        |d_unique| = number of unique terms in document d

    Args:
        delta: Absolute discount parameter. Must be in (0, 1).
               Default 0.7.
    """

    def __init__(self, delta: float = 0.7) -> None:
        if not (0.0 < delta < 1.0):
            raise ValueError("delta must be in (0, 1)")
        self.delta = delta
        self._total_corpus_tokens: int = 0
        self._corpus_term_freqs: Dict[str, int] = {}

    @property
    def total_corpus_tokens(self) -> int:
        return self._total_corpus_tokens

    def update_stats(
        self,
        total_corpus_tokens: int,
        corpus_term_freqs: Dict[str, int],
    ) -> None:
        """Set collection-level statistics.

        Args:
            total_corpus_tokens: Total number of tokens across all documents.
            corpus_term_freqs:   {term: total_occurrences_in_corpus}.
        """
        self._total_corpus_tokens = total_corpus_tokens
        self._corpus_term_freqs = corpus_term_freqs

    def score(
        self,
        query_terms: List[str],
        doc_term_freqs: Dict[str, int],
        doc_len: int,
    ) -> float:
        """Score a document using absolute discount smoothed language model.

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

        Returns:
            Log-probability score (sum of log-probabilities per term).
        """
        t = self._total_corpus_tokens
        if t <= 0 or doc_len <= 0:
            return 0.0

        delta = self.delta
        unique_terms = len(doc_term_freqs)
        total = 0.0

        for term in query_terms:
            tf = doc_term_freqs.get(term, 0)
            cf = self._corpus_term_freqs.get(term, 0)

            # Discounted document probability + backoff
            discounted = max(tf - delta, 0.0) / doc_len
            backoff = delta * unique_terms / doc_len * cf / t

            prob = discounted + backoff
            if prob <= 0.0:
                continue
            total += math.log(prob)

        return total

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

        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 [0, 1].
        """
        raw = self.score(query_terms, doc_term_freqs, doc_len)
        if raw > 500.0:
            return 1.0
        if raw < -500.0:
            return 0.0
        return 1.0 / (1.0 + math.exp(-raw))
