"""One retriever, three modes — the configurations compared in SQ1.

    from retrieval import Retriever
    r = Retriever(model="bge", mode="hybrid")      # mode: dense | bm25 | hybrid ; model: any key in embed_models.py
    hits = r.search("ကျောင်းဆရာ လစာ", k=5)            # -> list of (score, meta)

dense   cosine similarity between the query vector and every passage vector (the embedding model's view of meaning)
bm25    keyword overlap on Burmese syllables / syllable bigrams (exact words, names, numbers)
hybrid  both lists merged by reciprocal-rank fusion: score = 1/(60 + dense rank) + w · 1/(60 + keyword rank)
        (Cormack et al., 2009; w = bm25_weight, default 0.5).
Defaults (model="sealion", mode="hybrid") are the SQ1 winner of 2 Nov 2026: SEA-LION-E5 + hybrid, R@5 0.84, MRR 0.74 on 57 questions.
        The keyword side is used only for queries that contain Burmese words.

per_speech=N caps how many passages from one speech may appear in the top-k (the pilot's "one long speech
crowded out the others"). The dense index and the BM25 index hold the same passages in the same order, so
positions are shared.
"""
import numpy as np
from search import Index, load_meta
from bm25_index import BM25

RRF_K = 60
POOL = 100      # candidates taken from each ranking before fusion

class Retriever:
    def __init__(self, model="sealion", mode="hybrid", device=None, meta=None, dense=None, bm25=None, bm25_weight=0.5):
        self.mode = mode; self.model_key = model; self.bm25_weight = bm25_weight
        self.meta = meta if meta is not None else load_meta()
        self.dense = dense if dense is not None else (Index(model=model, device=device, meta=self.meta) if mode != "bm25" else None)
        self.bm25 = bm25 if bm25 is not None else (BM25() if mode != "dense" else None)
        if self.dense is not None and len(self.dense) != len(self.meta): raise RuntimeError("dense index / metadata size mismatch")

    @property
    def name(self): return f"{self.model_key}-{self.mode}" if self.mode != "bm25" else "bm25"

    def ranking(self, query, pool=POOL):
        """Fused (or single) ranking: list of (score, passage position), best first."""
        if self.mode == "dense":
            s = self.dense.scores(query); top = np.argpartition(-s, pool)[:pool]
            return sorted(((float(s[j]), int(j)) for j in top), reverse=True)
        if self.mode == "bm25":
            return self.bm25.search(query, pool)
        d = self.dense.scores(query)
        # keyword side only when the query has Burmese words to match: an English question would otherwise match
        # only the Latin header tokens every passage carries (house name, section type), which is noise
        use_bm25 = any("\u1000" <= ch <= "\u109f" for ch in query)
        fused = {}
        top = np.argpartition(-d, pool)[:pool]
        for rank, j in enumerate(sorted(top, key=lambda j: -d[j]), 1):
            fused[int(j)] = fused.get(int(j), 0.0) + 1.0 / (RRF_K + rank)
        if use_bm25:
            b = self.bm25.scores(query); top = np.argpartition(-b, pool)[:pool]
            for rank, j in enumerate(sorted(top, key=lambda j: -b[j]), 1):
                if b[j] <= 0: break
                fused[int(j)] = fused.get(int(j), 0.0) + self.bm25_weight / (RRF_K + rank)
        return sorted(((v, j) for j, v in fused.items()), reverse=True)

    def ranking_multi(self, queries, pool=POOL):
        """Fuse the rankings of several phrasings of one question (e.g. the English original and its Burmese
        translation) by reciprocal rank fusion, so the keyword index can contribute for a translated query."""
        if len(queries) == 1: return self.ranking(queries[0], pool)
        fused = {}
        for q in queries:
            for rank, (_, j) in enumerate(self.ranking(q, pool), 1):
                fused[j] = fused.get(j, 0.0) + 1.0 / (RRF_K + rank)
        return sorted(((v, j) for j, v in fused.items()), reverse=True)

    def search(self, query, k=5, per_speech=None, filter_fn=None, extra_queries=()):
        out, seen = [], {}
        queries = [query] + [q for q in extra_queries if q and q.strip() and q.strip() != query.strip()]
        for score, j in self.ranking_multi(queries):
            m = self.meta[j]
            if filter_fn and not filter_fn(m): continue
            if per_speech:
                c = seen.get(m["speech_id"], 0)
                if c >= per_speech: continue
                seen[m["speech_id"]] = c + 1
            out.append((score, m))
            if len(out) >= k: break
        return out

def all_configs(models=("bge",), modes=("dense", "hybrid"), device=None):
    """Build every (model × mode) retriever once, sharing metadata and the BM25 index."""
    meta = load_meta(); bm = BM25() if any(m != "dense" for m in modes) else None
    out = {}
    for mk in models:
        idx = Index(model=mk, device=device, meta=meta) if any(m != "bm25" for m in modes) else None
        for mode in modes:
            r = Retriever(model=mk, mode=mode, meta=meta, dense=idx, bm25=bm)
            out[r.name] = r
    return out