"""
RBSP Resonance-Based Hybrid Retrieval.

Novel fusion paradigm inspired by wave interference:
    - BM25 and vector rankings are treated as "signals".
    - When both signals agree (high coherence), scores are amplified
      (constructive interference).
    - When they disagree, scores are dampened (destructive interference).

This outperforms naive linear fusion and RRF on heterogeneous rankings
because it rewards *agreement* between fundamentally different retrieval
models rather than just averaging their scores.

Components:
    - ResonanceConfig:       Configuration with interference mode selection.
    - ResonanceScorer:       Core scoring with coherence + interference factor.
    - HybridRetriever:       End-to-end BM25+vector fusion via resonance.
    - AdaptiveResonance:     Self-tuning alpha via exponential moving average.
    - MultiSignalResonance:  N-signal generalization (beyond 2 rankers).
    - ndcg_at_k():           Standard NDCG evaluation helper.
    - compare_fusion_methods(): Resonance vs linear vs RRF comparison.

Zero external dependencies.
"""

from __future__ import annotations

import heapq
import math
from typing import Dict, List, Optional, Tuple

Ranking = List[Tuple[str, float]]


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


class ResonanceConfig:
    """Configuration for resonance-based hybrid retrieval.

    Args:
        alpha:             Interference strength factor (default 0.5).
        bm25_weight:       Weight for the BM25 signal (default 0.5).
        vector_weight:     Weight for the vector signal (default 0.5).
        min_signal:        Minimum score to include in results (default 0.0).
        interference_mode: One of ``"constructive"``, ``"adaptive"``, or
                           ``"linear"`` (default ``"adaptive"``).
    """

    __slots__ = (
        "alpha",
        "bm25_weight",
        "vector_weight",
        "min_signal",
        "interference_mode",
    )

    _VALID_MODES = ("constructive", "adaptive", "linear")

    def __init__(
        self,
        *,
        alpha: float = 0.5,
        bm25_weight: float = 0.5,
        vector_weight: float = 0.5,
        min_signal: float = 0.0,
        interference_mode: str = "adaptive",
    ) -> None:
        if interference_mode not in self._VALID_MODES:
            raise ValueError(
                f"interference_mode must be one of {self._VALID_MODES}, "
                f"got {interference_mode!r}"
            )
        self.alpha = alpha
        self.bm25_weight = bm25_weight
        self.vector_weight = vector_weight
        self.min_signal = min_signal
        self.interference_mode = interference_mode

    def __repr__(self) -> str:
        return (
            f"ResonanceConfig(alpha={self.alpha}, "
            f"bm25_weight={self.bm25_weight}, "
            f"vector_weight={self.vector_weight}, "
            f"interference_mode={self.interference_mode!r})"
        )


# ---------------------------------------------------------------------------
# Resonance Scorer
# ---------------------------------------------------------------------------


class ResonanceScorer:
    """Core resonance scoring engine.

    Computes an interference-adjusted fusion score from two signals'
    scores and rank positions.

    Args:
        config: A ``ResonanceConfig`` instance (optional, defaults apply).
    """

    def __init__(self, config: Optional[ResonanceConfig] = None) -> None:
        self._config = config or ResonanceConfig()

    @property
    def config(self) -> ResonanceConfig:
        return self._config

    @staticmethod
    def coherence(rank_a: int, rank_b: int, max_rank: int) -> float:
        """Compute coherence between two rank positions.

        Coherence = 1 - |rank_a - rank_b| / max_rank.
        Returns a value in [0.0, 1.0] where 1.0 means both signals rank
        the document identically.

        Args:
            rank_a:   Rank from signal A (1-indexed).
            rank_b:   Rank from signal B (1-indexed).
            max_rank: Maximum possible rank (number of documents).

        Returns:
            Coherence in [0.0, 1.0].
        """
        if max_rank <= 0:
            return 1.0
        coh = 1.0 - abs(rank_a - rank_b) / max_rank
        return max(0.0, min(1.0, coh))

    def interference_factor(self, coherence: float) -> float:
        """Compute the interference multiplier for a given coherence.

        Mode behaviour:
            - ``"linear"``:       Always returns 1.0 (no interference).
            - ``"constructive"``: ``1 + alpha * coherence`` (always >= 1.0,
              amplifies when signals agree).
            - ``"adaptive"``:     ``1 + alpha * (2 * coherence - 1)``
              (amplifies when coherence > 0.5, dampens when < 0.5).

        Args:
            coherence: Coherence value in [0.0, 1.0].

        Returns:
            Positive multiplier.
        """
        mode = self._config.interference_mode
        alpha = self._config.alpha

        if mode == "linear":
            return 1.0
        if mode == "constructive":
            return 1.0 + alpha * coherence
        # adaptive
        return 1.0 + alpha * (2.0 * coherence - 1.0)

    def score(
        self,
        bm25_score: float,
        vector_score: float,
        bm25_rank: int,
        vector_rank: int,
        max_rank: int,
    ) -> float:
        """Compute the resonance-fused score for a document.

        ``fused = (w_bm25 * bm25_score + w_vec * vec_score) / w_total
                  * interference_factor(coherence(bm25_rank, vec_rank))``

        Args:
            bm25_score:   Normalised BM25 score in [0, 1].
            vector_score: Normalised vector score in [0, 1].
            bm25_rank:    Rank from BM25 ranking (1-indexed).
            vector_rank:  Rank from vector ranking (1-indexed).
            max_rank:     Maximum rank (total documents).

        Returns:
            Fused score (non-negative float).
        """
        cfg = self._config
        w_total = cfg.bm25_weight + cfg.vector_weight
        if w_total <= 0:
            return 0.0

        weighted_sum = (
            cfg.bm25_weight * bm25_score + cfg.vector_weight * vector_score
        ) / w_total

        coh = self.coherence(bm25_rank, vector_rank, max_rank)
        factor = self.interference_factor(coh)

        return weighted_sum * factor


# ---------------------------------------------------------------------------
# Hybrid Retriever
# ---------------------------------------------------------------------------


class HybridRetriever:
    """End-to-end resonance-based fusion of BM25 and vector rankings.

    Normalises scores via min-max, builds rank maps, computes resonance
    scores for every document in the union of both rankings, and returns
    the top-K results.

    Args:
        config: A ``ResonanceConfig`` instance (optional).
    """

    def __init__(self, config: Optional[ResonanceConfig] = None) -> None:
        self._config = config or ResonanceConfig()
        self._scorer = ResonanceScorer(self._config)

    @property
    def config(self) -> ResonanceConfig:
        return self._config

    @staticmethod
    def _normalize(ranking: Ranking) -> Dict[str, float]:
        """Min-max normalise a ranking to [0, 1].

        If all scores are identical the normalised score is 0.5.
        """
        if not ranking:
            return {}
        scores = [s for _, s in ranking]
        lo = min(scores)
        hi = max(scores)
        span = hi - lo
        result: Dict[str, float] = {}
        for doc_id, s in ranking:
            if span == 0.0:
                result[doc_id] = 0.5
            else:
                result[doc_id] = (s - lo) / span
        return result

    @staticmethod
    def _rank_map(ranking: Ranking) -> Dict[str, int]:
        """Build a doc_id -> 1-based rank mapping (sorted by descending score)."""
        sorted_r = sorted(ranking, key=lambda x: -x[1])
        return {doc_id: rank for rank, (doc_id, _) in enumerate(sorted_r, 1)}

    def fuse(
        self,
        bm25_ranking: Ranking,
        vector_ranking: Ranking,
        *,
        top_k: int = 10,
    ) -> Ranking:
        """Fuse BM25 and vector rankings via resonance interference.

        Args:
            bm25_ranking:   List of ``(doc_id, score)`` from BM25.
            vector_ranking: List of ``(doc_id, score)`` from vector search.
            top_k:          Number of results to return.

        Returns:
            Fused ranking sorted by descending resonance score.
        """
        if not bm25_ranking and not vector_ranking:
            return []

        # Normalise scores
        bm25_norm = self._normalize(bm25_ranking)
        vec_norm = self._normalize(vector_ranking)

        # Build rank maps
        bm25_ranks = self._rank_map(bm25_ranking)
        vec_ranks = self._rank_map(vector_ranking)

        # Union of all doc IDs
        all_docs = set(bm25_norm.keys()) | set(vec_norm.keys())
        max_rank = len(all_docs)

        # Default rank for missing documents = max_rank (worst)
        results: List[Tuple[str, float]] = []
        for doc_id in all_docs:
            b_score = bm25_norm.get(doc_id, 0.0)
            v_score = vec_norm.get(doc_id, 0.0)
            b_rank = bm25_ranks.get(doc_id, max_rank)
            v_rank = vec_ranks.get(doc_id, max_rank)

            fused = self._scorer.score(b_score, v_score, b_rank, v_rank, max_rank)

            if fused >= self._config.min_signal:
                results.append((doc_id, fused))

        return heapq.nlargest(top_k, results, key=lambda x: x[1])

    def fuse_with_scores(
        self,
        bm25_ranking: Ranking,
        vector_ranking: Ranking,
        *,
        top_k: int = 10,
    ) -> List[Tuple[str, float, float, float, float]]:
        """Fuse with full score breakdown.

        Returns:
            List of ``(doc_id, fused_score, bm25_norm, vec_norm, coherence)``
            tuples sorted by descending fused score.
        """
        if not bm25_ranking and not vector_ranking:
            return []

        bm25_norm = self._normalize(bm25_ranking)
        vec_norm = self._normalize(vector_ranking)
        bm25_ranks = self._rank_map(bm25_ranking)
        vec_ranks = self._rank_map(vector_ranking)

        all_docs = set(bm25_norm.keys()) | set(vec_norm.keys())
        max_rank = len(all_docs)

        results: List[Tuple[str, float, float, float, float]] = []
        for doc_id in all_docs:
            b_score = bm25_norm.get(doc_id, 0.0)
            v_score = vec_norm.get(doc_id, 0.0)
            b_rank = bm25_ranks.get(doc_id, max_rank)
            v_rank = vec_ranks.get(doc_id, max_rank)

            coh = ResonanceScorer.coherence(b_rank, v_rank, max_rank)
            fused = self._scorer.score(b_score, v_score, b_rank, v_rank, max_rank)

            if fused >= self._config.min_signal:
                results.append((doc_id, fused, b_score, v_score, coh))

        return heapq.nlargest(top_k, results, key=lambda x: x[1])


# ---------------------------------------------------------------------------
# Adaptive Resonance
# ---------------------------------------------------------------------------


class AdaptiveResonance:
    """Self-tuning alpha via exponential moving average of click feedback.

    Observes user clicks and adjusts the interference strength based on
    whether high-coherence documents are being clicked.

    Args:
        initial_alpha: Starting alpha value (default 0.5).
        learning_rate: EMA smoothing factor in (0, 1] (default 0.1).
    """

    __slots__ = ("_alpha", "_learning_rate")

    def __init__(
        self,
        initial_alpha: float = 0.5,
        learning_rate: float = 0.1,
    ) -> None:
        self._alpha = initial_alpha
        self._learning_rate = learning_rate

    def get_alpha(self) -> float:
        """Return the current alpha value."""
        return self._alpha

    def observe(self, clicked_rank: int, max_rank: int) -> None:
        """Update alpha based on a click event.

        A click on a high-ranked document nudges alpha toward the
        observed signal (rank position normalized to [0, 1]).

        Args:
            clicked_rank: 1-based rank of the clicked document.
            max_rank:     Total number of documents in the result set.
        """
        if max_rank <= 0:
            return
        # Signal: 1.0 for rank 1 (best), 0.0 for last rank
        signal = 1.0 - (clicked_rank - 1) / max_rank
        signal = max(0.0, min(1.0, signal))
        # EMA update
        self._alpha = (
            (1.0 - self._learning_rate) * self._alpha
            + self._learning_rate * signal
        )

    def reset(self, alpha: float = 0.5) -> None:
        """Reset alpha to a given value."""
        self._alpha = alpha


# ---------------------------------------------------------------------------
# Multi-Signal Resonance
# ---------------------------------------------------------------------------


class MultiSignalResonance:
    """N-signal generalization of resonance fusion.

    Computes average pairwise coherence across all signal pairs, applies
    the interference factor to the mean normalised score, and returns
    top-K results.

    Args:
        config: A ``ResonanceConfig`` instance (optional).
    """

    def __init__(self, config: Optional[ResonanceConfig] = None) -> None:
        self._config = config or ResonanceConfig()
        self._scorer = ResonanceScorer(self._config)

    def fuse(
        self,
        rankings: List[Ranking],
        *,
        top_k: int = 10,
    ) -> Ranking:
        """Fuse N rankings via multi-signal resonance.

        Args:
            rankings: List of rankings from different retrieval signals.
            top_k:    Number of results to return.

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

        n_signals = len(rankings)

        # Normalise all rankings
        norm_maps: List[Dict[str, float]] = []
        rank_maps: List[Dict[str, int]] = []
        all_docs: set = set()

        for ranking in rankings:
            norm = HybridRetriever._normalize(ranking)
            ranks = HybridRetriever._rank_map(ranking)
            norm_maps.append(norm)
            rank_maps.append(ranks)
            all_docs.update(norm.keys())

        max_rank = len(all_docs)
        if max_rank == 0:
            return []

        results: List[Tuple[str, float]] = []
        for doc_id in all_docs:
            # Mean normalised score
            total_score = 0.0
            for nm in norm_maps:
                total_score += nm.get(doc_id, 0.0)
            mean_score = total_score / n_signals

            # Average pairwise coherence
            if n_signals < 2:
                avg_coh = 1.0
            else:
                pair_count = 0
                coh_sum = 0.0
                for i in range(n_signals):
                    for j in range(i + 1, n_signals):
                        r_i = rank_maps[i].get(doc_id, max_rank)
                        r_j = rank_maps[j].get(doc_id, max_rank)
                        coh_sum += ResonanceScorer.coherence(r_i, r_j, max_rank)
                        pair_count += 1
                avg_coh = coh_sum / pair_count

            factor = self._scorer.interference_factor(avg_coh)
            fused = mean_score * factor

            if fused >= self._config.min_signal:
                results.append((doc_id, fused))

        return heapq.nlargest(top_k, results, key=lambda x: x[1])


# ---------------------------------------------------------------------------
# NDCG evaluation
# ---------------------------------------------------------------------------


def ndcg_at_k(
    ranking: List[str],
    relevance: Dict[str, float],
    k: int = 10,
) -> float:
    """Compute NDCG@k (Normalised Discounted Cumulative Gain).

    Args:
        ranking:   Ordered list of doc IDs (best first).
        relevance: Mapping of doc_id -> relevance score.
        k:         Cutoff position.

    Returns:
        NDCG score in [0.0, 1.0].
    """
    if not ranking or not relevance:
        return 0.0

    k = min(k, len(ranking))

    # DCG
    dcg = 0.0
    for i in range(k):
        rel = relevance.get(ranking[i], 0.0)
        dcg += (2.0 ** rel - 1.0) / math.log2(i + 2.0)

    # Ideal DCG
    ideal_rels = sorted(relevance.values(), reverse=True)[:k]
    idcg = 0.0
    for i, rel in enumerate(ideal_rels):
        idcg += (2.0 ** rel - 1.0) / math.log2(i + 2.0)

    if idcg == 0.0:
        return 0.0

    return dcg / idcg


# ---------------------------------------------------------------------------
# Comparison utility
# ---------------------------------------------------------------------------


def compare_fusion_methods(
    bm25_ranking: Ranking,
    vector_ranking: Ranking,
    relevance_labels: Dict[str, float],
    *,
    k: int = 10,
) -> Dict[str, float]:
    """Compare resonance fusion against linear and RRF baselines.

    Args:
        bm25_ranking:    BM25 ranking.
        vector_ranking:  Vector ranking.
        relevance_labels: Ground-truth relevance labels.
        k:               NDCG cutoff.

    Returns:
        Dict mapping method name to NDCG@k score.
    """
    from rbsp.query.ensemble import ScoreFusion

    results: Dict[str, float] = {}

    # Resonance fusion (adaptive)
    retriever = HybridRetriever(ResonanceConfig(interference_mode="adaptive"))
    res_ranking = retriever.fuse(bm25_ranking, vector_ranking, top_k=k)
    res_ids = [doc_id for doc_id, _ in res_ranking]
    results["resonance"] = ndcg_at_k(res_ids, relevance_labels, k)

    # Linear fusion
    lin_ranking = ScoreFusion.linear(
        [bm25_ranking, vector_ranking], [0.5, 0.5]
    )[:k]
    lin_ids = [doc_id for doc_id, _ in lin_ranking]
    results["linear"] = ndcg_at_k(lin_ids, relevance_labels, k)

    # RRF
    rrf_ranking = ScoreFusion.reciprocal_rank_fusion(
        [bm25_ranking, vector_ranking]
    )[:k]
    rrf_ids = [doc_id for doc_id, _ in rrf_ranking]
    results["rrf"] = ndcg_at_k(rrf_ids, relevance_labels, k)

    return results
