"""Hybrid retrieval over the private unit index: BM25 + cosine (local embeddings), metadata filters, provenance on every hit."""
from __future__ import annotations
import json, re
from functools import lru_cache
import numpy as np, pandas as pd
from rank_bm25 import BM25Okapi
from .embeddings import INDEX_DIR
from . import config

def _tok(s: str) -> list[str]: return re.findall(r"[a-z0-9]+", s.lower())

@lru_cache(maxsize=1)
def _load():
    df = pd.read_parquet(INDEX_DIR / "units.parquet"); emb = np.load(INDEX_DIR / "embeddings.npy")
    bm = BM25Okapi([_tok(t) for t in df["text"]])
    from sentence_transformers import SentenceTransformer
    model = SentenceTransformer(json.loads((INDEX_DIR / "index_meta.json").read_text())["embedding_model"], device="cpu")
    return df, emb, bm, model

def search(query: str, k: int = 8, section: str | None = None, year: int | None = None, unit: str | None = "annotation", genre: str | None = None, subfield: str | None = None,
           split: str | None = "train", alpha: float = 0.5, mode: str = "hybrid") -> list[dict]:
    """unit: 'annotation' = structural units only (summary, identification, intro_move, claim_evidence); 'raw' = section openings/table notes/captions
    (conventional RAG over short excerpts); None = everything. mode: hybrid | bm25 | dense."""
    df, emb, bm, model = _load()
    mask = np.ones(len(df), bool)
    if split: mask &= (df["split"] == split).values
    if year: mask &= (df["year"] == year).values
    if genre: mask &= (df["genre"] == genre).values
    if subfield: mask &= (df["subfield"] == subfield).values
    if section: mask &= (df["role"] == section).values | df["section"].fillna("").str.contains(section, case=False).values
    if unit == "annotation": mask &= df["unit"].isin(["document_summary", "identification", "intro_move", "claim_evidence"]).values
    elif unit == "raw": mask &= df["unit"].isin(["section_opening", "table_notes", "figure_caption"]).values
    elif unit: mask &= (df["unit"] == unit).values
    idx = np.where(mask)[0]
    if len(idx) == 0: return []
    b = np.array(bm.get_scores(_tok(query)))[idx]; b = (b - b.min()) / (b.max() - b.min() + 1e-9)
    qv = model.encode([query], normalize_embeddings=True)[0]; c = emb[idx] @ qv; c = (c - c.min()) / (c.max() - c.min() + 1e-9)
    score = b if mode == "bm25" else c if mode == "dense" else alpha * b + (1 - alpha) * c
    order = idx[np.argsort(-score)][:k * 3]
    out, seen = [], {}
    for j in order:
        r = df.iloc[j]
        if seen.get(r.doc_id, 0) >= 2: continue  # at most 2 units per document
        seen[r.doc_id] = seen.get(r.doc_id, 0) + 1
        out.append({"score": round(float(score[np.where(idx == j)[0][0]]), 4), "unit": r.unit, "text": r.text, "provenance": {"doc_id": r.doc_id, "doi": r.doi, "title": r.title, "year": int(r.year) if pd.notna(r.year) else None, "page": int(r.page) if pd.notna(r.page) else None, "section": r.section, "sha256": r.sha256[:16], "role": r.role, "genre": r.genre, "subfield": r.subfield, "confidence": float(r.confidence) if pd.notna(r.confidence) else None, "split": r.split}})
        if len(out) >= k: break
    return out