"""
RBSP Ensemble Score Fusion.

Combines scores from multiple rankers (BM25, language models, semantic
similarity, etc.) into a single unified ranking.  Standard technique in
TREC-winning information retrieval systems.

Supported fusion methods:
    - LINEAR:           Weighted sum of min-max normalised scores.
    - RECIPROCAL_RANK:  RRF -- score = sum 1/(k + rank_i).  State of the art
                        for combining heterogeneous rankings.
    - COMBSCORE:        Sum of raw scores across all rankings.
    - COMBMAX:          Maximum score per document across rankings.
    - BORDA:            Borda count positional voting.
    - MIN:              Minimum score per document across rankings.
"""

from __future__ import annotations

import heapq
import math
from enum import Enum, auto
from typing import Dict, List, Optional, Tuple

Ranking = List[Tuple[str, float]]


class FusionMethod(Enum):
    """Available ensemble score fusion strategies."""

    LINEAR = auto()
    RECIPROCAL_RANK = auto()
    COMBSCORE = auto()
    COMBMAX = auto()
    BORDA = auto()
    MIN = auto()


def normalize_scores(
    ranking: Ranking,
    method: str = "minmax",
) -> Ranking:
    """Normalize scores in a ranking to [0, 1].

    Args:
        ranking: List of (doc_id, score) tuples.
        method:  ``"minmax"`` for min-max scaling, ``"zscore"`` for
                 z-score normalisation (then clamped to [0, 1]).

    Returns:
        New ranking with normalised scores, same order.
    """
    if not ranking:
        return []

    scores = [s for _, s in ranking]

    if method == "minmax":
        lo = min(scores)
        hi = max(scores)
        span = hi - lo
        if span == 0.0:
            return [(doc_id, 0.0) for doc_id, _ in ranking]
        return [(doc_id, (s - lo) / span) for doc_id, s in ranking]

    if method == "zscore":
        n = len(scores)
        mean = sum(scores) / n
        variance = sum((s - mean) ** 2 for s in scores) / n
        std = math.sqrt(variance) if variance > 0 else 0.0
        if std == 0.0:
            return [(doc_id, 0.0) for doc_id, _ in ranking]
        normed: Ranking = []
        for doc_id, s in ranking:
            z = (s - mean) / std
            # Clamp to [0, 1] via sigmoid-like min/max
            normed.append((doc_id, max(0.0, min(1.0, (z + 3.0) / 6.0))))
        return normed

    raise ValueError(f"Unknown normalisation method: {method!r}")


class ScoreFusion:
    """Ensemble score fusion for combining multiple ranking lists.

    Each ranking is a list of ``(doc_id, score)`` tuples.  Results
    include every document appearing in *any* input ranking.
    """

    @staticmethod
    def linear(
        rankings: List[Ranking],
        weights: List[float],
    ) -> Ranking:
        """Weighted linear combination of normalised scores.

        Scores are first min-max normalised to [0, 1], then combined as
        ``fused = sum(w_i * norm_score_i)`` for each document.

        Args:
            rankings: One ranking per ranker.
            weights:  Per-ranker weights (must match length of *rankings*).

        Returns:
            Fused ranking sorted by descending score.
        """
        if len(rankings) != len(weights):
            raise ValueError(
                f"rankings ({len(rankings)}) and weights ({len(weights)}) "
                "must have the same length"
            )
        if not rankings:
            return []

        w_total = sum(weights)
        if w_total <= 0:
            raise ValueError("Weights must sum to a positive value")

        fused: Dict[str, float] = {}
        for ranking, w in zip(rankings, weights):
            normed = normalize_scores(ranking, method="minmax")
            for doc_id, score in normed:
                fused[doc_id] = fused.get(doc_id, 0.0) + w * score

        # Normalise by weight total
        result = [(doc_id, score / w_total) for doc_id, score in fused.items()]
        return heapq.nlargest(len(result), result, key=lambda x: x[1])

    @staticmethod
    def reciprocal_rank_fusion(
        rankings: List[Ranking],
        k: int = 60,
    ) -> Ranking:
        """Reciprocal Rank Fusion (RRF).

        For each document, ``score = sum 1 / (k + rank_i)`` across all
        rankings where it appears.  Higher-ranked documents contribute
        more.  ``k`` controls how much rank position matters (higher k
        compresses rank differences).

        Args:
            rankings: One ranking per ranker (order matters -- first
                      entry is rank 1).
            k:        Smoothing constant (default 60, standard in the
                      literature).

        Returns:
            Fused ranking sorted by descending RRF score.
        """
        if not rankings:
            return []

        fused: Dict[str, float] = {}
        for ranking in rankings:
            # Sort by score descending to assign rank positions
            sorted_ranking = sorted(ranking, key=lambda x: -x[1])
            for rank_pos, (doc_id, _score) in enumerate(sorted_ranking, start=1):
                fused[doc_id] = fused.get(doc_id, 0.0) + 1.0 / (k + rank_pos)

        result = list(fused.items())
        return heapq.nlargest(len(result), result, key=lambda x: x[1])

    @staticmethod
    def combscore(rankings: List[Ranking]) -> Ranking:
        """CombSUM / CombScore -- sum of raw scores across rankings.

        Args:
            rankings: One ranking per ranker.

        Returns:
            Fused ranking sorted by descending summed score.
        """
        if not rankings:
            return []

        fused: Dict[str, float] = {}
        for ranking in rankings:
            for doc_id, score in ranking:
                fused[doc_id] = fused.get(doc_id, 0.0) + score

        result = list(fused.items())
        return heapq.nlargest(len(result), result, key=lambda x: x[1])

    @staticmethod
    def combmax(rankings: List[Ranking]) -> Ranking:
        """CombMAX -- maximum score per document across rankings.

        Args:
            rankings: One ranking per ranker.

        Returns:
            Fused ranking sorted by descending max score.
        """
        if not rankings:
            return []

        fused: Dict[str, float] = {}
        for ranking in rankings:
            for doc_id, score in ranking:
                if doc_id not in fused or score > fused[doc_id]:
                    fused[doc_id] = score

        result = list(fused.items())
        return heapq.nlargest(len(result), result, key=lambda x: x[1])

    @staticmethod
    def borda_count(rankings: List[Ranking]) -> Ranking:
        """Borda count positional voting.

        Each document gets ``N - rank_i`` points from each ranking of
        size *N* where it appears.  The document ranked first gets
        ``N - 1`` points, the last gets 0.

        Args:
            rankings: One ranking per ranker.

        Returns:
            Fused ranking sorted by descending Borda score.
        """
        if not rankings:
            return []

        fused: Dict[str, float] = {}
        for ranking in rankings:
            sorted_ranking = sorted(ranking, key=lambda x: -x[1])
            n = len(sorted_ranking)
            for rank_pos, (doc_id, _score) in enumerate(sorted_ranking):
                # rank_pos 0 (best) gets n-1 points, last gets 0
                fused[doc_id] = fused.get(doc_id, 0.0) + (n - 1 - rank_pos)

        result = list(fused.items())
        return heapq.nlargest(len(result), result, key=lambda x: x[1])

    @staticmethod
    def combmin(rankings: List[Ranking]) -> Ranking:
        """CombMIN -- minimum score per document across rankings.

        Only meaningful when a document appears in multiple rankings;
        documents appearing in a single ranking use that score.

        Args:
            rankings: One ranking per ranker.

        Returns:
            Fused ranking sorted by descending min score.
        """
        if not rankings:
            return []

        fused: Dict[str, float] = {}
        for ranking in rankings:
            for doc_id, score in ranking:
                if doc_id not in fused or score < fused[doc_id]:
                    fused[doc_id] = score

        result = list(fused.items())
        return heapq.nlargest(len(result), result, key=lambda x: x[1])

    @classmethod
    def fuse(
        cls,
        method: FusionMethod,
        rankings: List[Ranking],
        *,
        weights: Optional[List[float]] = None,
        k: int = 60,
    ) -> Ranking:
        """Dispatch fusion by method enum.

        Args:
            method:   The ``FusionMethod`` to apply.
            rankings: One ranking per ranker.
            weights:  Required for ``LINEAR``; ignored otherwise.
            k:        Smoothing constant for ``RECIPROCAL_RANK``.

        Returns:
            Fused ranking sorted by descending score.
        """
        if method is FusionMethod.LINEAR:
            if weights is None:
                raise ValueError("weights required for LINEAR fusion")
            return cls.linear(rankings, weights)
        if method is FusionMethod.RECIPROCAL_RANK:
            return cls.reciprocal_rank_fusion(rankings, k=k)
        if method is FusionMethod.COMBSCORE:
            return cls.combscore(rankings)
        if method is FusionMethod.COMBMAX:
            return cls.combmax(rankings)
        if method is FusionMethod.BORDA:
            return cls.borda_count(rankings)
        if method is FusionMethod.MIN:
            return cls.combmin(rankings)
        raise ValueError(f"Unknown fusion method: {method!r}")
