
"""
Entropy-RAG Retriever
Implements entropy-balanced retrieval selection (Reid, 2025).
"""

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from dataclasses import dataclass

@dataclass
class EntropySelectorConfig:
    Omega_mode: str = "median"     # How to compute Ω across topics
    RE_mode: str = "sqrt_count"    # How to compute recursive entropy weighting
    diversity_floor: float = 0.3   # Minimum entropy for small topic clusters
    alpha: float = 1.0             # Global temperature scaling for entropy weights

def topic_diversity(doc_vecs_topic):
    """Proxy for topic entropy H_inf: average pairwise dissimilarity within a topic."""
    if doc_vecs_topic.shape[0] <= 1:
        return 0.3
    S = cosine_similarity(doc_vecs_topic)
    mask = np.triu(np.ones(S.shape), 1) == 1
    return float(np.clip((1 - S[mask]).mean(), 0.1, 1.0))

class EntropyBalancedSelector:
    """
    Entropy-weighted retrieval system.
    Combines cosine similarity with recursive entropy balancing across topics.
    """

    def __init__(self, index, cfg: EntropySelectorConfig):
        self.index = index
        self.cfg = cfg

        # Precompute topic-level entropies
        self.topic_Hinf = {}
        for t, idxs in index.topics.items():
            self.topic_Hinf[t] = topic_diversity(index.doc_vecs[idxs])

        # Compute global Ω (entropy coupling)
        if cfg.Omega_mode == "median":
            self.Omega = np.median(list(self.topic_Hinf.values()))
        elif cfg.Omega_mode == "mean":
            self.Omega = np.mean(list(self.topic_Hinf.values()))
        else:
            self.Omega = 1.0

        # Compute per-topic λ scaling (entropy-based)
        self.lambda_t = {t: max(1e-6, self.topic_Hinf[t] / self.Omega)
                         for t in self.topic_Hinf}
        self.C_t = {t: (1 + np.log1p(self.lambda_t[t])) * self.cfg.alpha
                    for t in self.lambda_t}

    # ------------------------------------------------------------------
    def retrieve(self, query: str, k: int = 5):
        """
        Retrieve top-k documents with entropy-balanced weighting.
        """
        q_vec = self.index.embedder.encode([query])
        sims = cosine_similarity(q_vec, self.index.doc_vecs).ravel()

        weighted_scores = np.zeros_like(sims)
        for t, idxs in self.index.topics.items():
            weight = self.C_t.get(t, 1.0)
            weighted_scores[idxs] = sims[idxs] * weight

        topk_idx = np.argsort(weighted_scores)[::-1][:k]
        return [self.index.docs[i] for i in topk_idx]

# ----------------------------------------------------------------------
def vanilla_topk(index, query, k=5):
    """Baseline retrieval: plain cosine similarity."""
    q_vec = index.embedder.encode([query])
    sims = cosine_similarity(q_vec, index.doc_vecs).ravel()
    topk_idx = np.argsort(sims)[::-1][:k]
    return [index.docs[i] for i in topk_idx]
