"""Hybrid search over the openai/math catalogue.

Each family is indexed as one document (title + summary + aliases + plain-English
line) plus one document per paper (title + abstract). A family's dense score is
the max cosine similarity over its documents; its BM25 score is the max BM25
score over the same documents. The two family rankings are fused with
reciprocal-rank fusion. This is the configuration that scored best on the
eval (see eval/RESULTS.md: qwen3-0.6b+expanded, hybrid_rrf).
"""

import json
import threading
from collections import OrderedDict
from pathlib import Path

import bm25s
import numpy as np
import Stemmer

MODEL_ID = "Qwen/Qwen3-Embedding-0.6B"
# Document vectors in index/ were computed with this exact revision; queries must use it too.
MODEL_REVISION = "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3"
QUERY_PROMPT = ("Instruct: Given a search query about mathematics, retrieve research results "
                "whose statement answers or matches the query\nQuery: ")
RRF_K = 60


def family_doc_text(f: dict) -> str:
    return f"{f['title']}. {f['summary']} Also known as: {', '.join(f['aliases'])}. {f['plain']}"


def paper_doc_text(p: dict) -> str:
    return f"{p['title']}. {p['abstract']}"


def build_docs(families: list[dict]) -> list[dict]:
    """Documents in family order: each family's summary doc, then its papers."""
    docs = []
    for fi, f in enumerate(families):
        docs.append({"family": fi, "paper": None, "text": family_doc_text(f)})
        for pi, p in enumerate(f["papers"]):
            docs.append({"family": fi, "paper": pi, "text": paper_doc_text(p)})
    return docs


class SearchEngine:
    def __init__(self, index_dir: Path, load_model: bool = True):
        self.catalog = json.loads((index_dir / "catalog.json").read_text())
        self.families = self.catalog["families"]
        self.docs = build_docs(self.families)
        self.emb = np.load(index_dir / "embeddings.npy")
        assert self.emb.shape[0] == len(self.docs), "index out of date: rebuild with pipeline/build_index.py"

        doc_family = np.array([d["family"] for d in self.docs])
        self.starts = np.flatnonzero(np.r_[True, doc_family[1:] != doc_family[:-1]])
        self.ends = np.r_[self.starts[1:], len(self.docs)]

        self.stemmer = Stemmer.Stemmer("english")
        self.bm25 = bm25s.BM25()
        self.bm25.index(bm25s.tokenize([d["text"] for d in self.docs], stopwords="en",
                                       stemmer=self.stemmer, show_progress=False), show_progress=False)

        self.model = None
        self._lock = threading.Lock()
        self._cache: OrderedDict[str, np.ndarray] = OrderedDict()
        if load_model:
            from sentence_transformers import SentenceTransformer
            self.model = SentenceTransformer(MODEL_ID, revision=MODEL_REVISION, device="cpu")
            self.model.max_seq_length = 512

    def embed_query(self, query: str) -> np.ndarray:
        if query in self._cache:
            self._cache.move_to_end(query)
            return self._cache[query]
        with self._lock:
            vec = self.model.encode([query], prompt=QUERY_PROMPT, normalize_embeddings=True,
                                    convert_to_numpy=True)[0].astype(np.float32)
        self._cache[query] = vec
        if len(self._cache) > 2048:
            self._cache.popitem(last=False)
        return vec

    def _family_max(self, doc_scores: np.ndarray) -> np.ndarray:
        return np.maximum.reduceat(doc_scores, self.starts)

    def search(self, query: str) -> list[dict]:
        """All families ranked for `query`, best first."""
        sims = self.emb @ self.embed_query(query)
        dense = self._family_max(sims)

        tokens = bm25s.tokenize([query], stopwords="en", stemmer=self.stemmer,
                                return_ids=False, show_progress=False)[0]
        lexical = self._family_max(self.bm25.get_scores(tokens)) if tokens else np.zeros_like(dense)

        dense_rank = np.empty(len(dense), dtype=int)
        dense_rank[np.argsort(-dense, kind="stable")] = np.arange(len(dense))
        lex_rank = np.empty(len(lexical), dtype=int)
        lex_rank[np.argsort(-lexical, kind="stable")] = np.arange(len(lexical))
        fused = 1 / (RRF_K + dense_rank + 1) + (lexical > 0) / (RRF_K + lex_rank + 1)

        results = []
        for fi in np.argsort(-fused, kind="stable"):
            best = self.starts[fi] + int(np.argmax(sims[self.starts[fi]:self.ends[fi]]))
            results.append({
                "id": self.families[fi]["id"],
                "score": round(float(fused[fi]), 6),
                "similarity": round(float(dense[fi]), 4),
                "keyword_match": bool(lexical[fi] > 0),
                "matched_paper": self.docs[best]["paper"],
            })
        return results