"""
RBSP Freshness Boosting.

Time-decay freshness boosting for search results.  Newer documents
receive a higher score contribution, modulated by a configurable
decay function (exponential, Gaussian, or linear).

Classes:
    - DecayFunction:    Enum of supported decay curves.
    - FreshnessBooster: Computes and applies time-decay boost factors.
    - RecencyRanker:    Re-ranks scored results with freshness boosting.
    - DateRangeFilter:  Filters document sets by timestamp bounds.
"""

from __future__ import annotations

import math
import time
from enum import Enum
from typing import Dict, List, Optional, Set, Tuple


class DecayFunction(Enum):
    """Supported time-decay curves for freshness boosting."""

    EXPONENTIAL = "exponential"
    GAUSSIAN = "gaussian"
    LINEAR = "linear"


class FreshnessBooster:
    """Computes time-decay boost factors for search results.

    The booster maps a document timestamp to a factor in [0, 1] that
    indicates how "fresh" the document is relative to *now*.  Three
    decay curves are supported:

    - **Exponential**: ``exp(-0.693 * age / halflife)``
    - **Gaussian**: ``exp(-0.5 * (age / scale)^2)``
    - **Linear**: ``max(0, 1 - age / scale)``

    The ``apply`` method blends the boost into an existing relevance
    score:  ``score * (1 - weight + weight * boost)``.  With the
    default weight of 0.1, freshness accounts for at most 10 % of the
    final score.

    Args:
        decay_fn:  Decay curve to use (default EXPONENTIAL).
        halflife:  Seconds until the exponential boost drops to 0.5
                   (default 86400 = 24 h).
        origin:    Reference timestamp treated as "now".  When *None*,
                   ``time.time()`` is used at call time.
        scale:     Scale parameter in seconds for Gaussian and linear
                   decay (default 86400).
        weight:    Blending weight in [0, 1] (default 0.1).
    """

    def __init__(
        self,
        decay_fn: DecayFunction = DecayFunction.EXPONENTIAL,
        halflife: float = 86400.0,
        origin: Optional[float] = None,
        scale: float = 86400.0,
        weight: float = 0.1,
    ) -> None:
        self.decay_fn = decay_fn
        self.halflife = halflife
        self.origin = origin
        self.scale = scale
        self.weight = weight

    def boost(self, timestamp: float, now: Optional[float] = None) -> float:
        """Compute the freshness boost factor for *timestamp*.

        Args:
            timestamp: POSIX timestamp of the document.
            now:       Reference "current" time.  Falls back to
                       ``self.origin``, then ``time.time()``.

        Returns:
            A float in [0.0, 1.0] where 1.0 means perfectly fresh.
        """
        if now is None:
            now = self.origin if self.origin is not None else time.time()
        age = max(now - timestamp, 0.0)

        if self.decay_fn is DecayFunction.EXPONENTIAL:
            if self.halflife <= 0:
                return 1.0 if age == 0.0 else 0.0
            return math.exp(-0.693 * age / self.halflife)

        if self.decay_fn is DecayFunction.GAUSSIAN:
            if self.scale <= 0:
                return 1.0 if age == 0.0 else 0.0
            return math.exp(-0.5 * (age / self.scale) ** 2)

        # LINEAR
        if self.scale <= 0:
            return 1.0 if age == 0.0 else 0.0
        return max(0.0, 1.0 - age / self.scale)

    def apply(self, score: float, timestamp: float, now: Optional[float] = None) -> float:
        """Apply freshness boosting to *score*.

        The blended score is::

            score * (1 - weight + weight * boost)

        This ensures the original score is scaled by a factor in
        ``[1 - weight, 1]``, so freshness never contributes more than
        ``weight`` of the final result.

        Args:
            score:     Original relevance score.
            timestamp: Document timestamp.
            now:       Reference time (see :meth:`boost`).

        Returns:
            The boosted score.
        """
        b = self.boost(timestamp, now=now)
        return score * (1.0 - self.weight + self.weight * b)


class RecencyRanker:
    """Re-ranks scored results by applying freshness boosting.

    Args:
        booster: A configured :class:`FreshnessBooster` instance.
    """

    def __init__(self, booster: FreshnessBooster) -> None:
        self.booster = booster

    def rerank(
        self,
        results: List[Tuple[str, float, float]],
    ) -> List[Tuple[str, float]]:
        """Re-rank results with freshness boosting.

        Args:
            results: A list of ``(doc_id, score, timestamp)`` tuples.

        Returns:
            A list of ``(doc_id, boosted_score)`` tuples sorted by
            descending boosted score.
        """
        boosted: List[Tuple[float, int, str]] = []
        for idx, (doc_id, score, timestamp) in enumerate(results):
            new_score = self.booster.apply(score, timestamp)
            boosted.append((-new_score, idx, doc_id))

        boosted.sort()
        return [(doc_id, -neg_score) for neg_score, _idx, doc_id in boosted]


class DateRangeFilter:
    """Filters document sets by timestamp bounds.

    Documents whose timestamp falls outside ``[min_date, max_date]``
    are excluded.  Either bound may be *None* to leave that side open.

    Args:
        min_date: Minimum POSIX timestamp (inclusive).  *None* means
                  no lower bound.
        max_date: Maximum POSIX timestamp (inclusive).  *None* means
                  no upper bound.
    """

    def __init__(
        self,
        min_date: Optional[float] = None,
        max_date: Optional[float] = None,
    ) -> None:
        self.min_date = min_date
        self.max_date = max_date

    def filter(
        self,
        doc_ids: Set[str],
        timestamps: Dict[str, float],
    ) -> Set[str]:
        """Return the subset of *doc_ids* within the date range.

        Documents missing from *timestamps* are excluded.

        Args:
            doc_ids:    Candidate document IDs.
            timestamps: Mapping of document ID to POSIX timestamp.

        Returns:
            A new set containing only the IDs that pass the filter.
        """
        result: Set[str] = set()
        for doc_id in doc_ids:
            ts = timestamps.get(doc_id)
            if ts is None:
                continue
            if self.min_date is not None and ts < self.min_date:
                continue
            if self.max_date is not None and ts > self.max_date:
                continue
            result.add(doc_id)
        return result
