#!/usr/bin/env python3
"""hybrid-search-bm25-dense-fusion.py - BM25, two embedding models and two ways to fuse them, measured on BEIR SciFact.

What it does
  Ranks the 5,183 SciFact abstracts for all 1,109 SciFact claims (809 train, 300 test) with
    bm25     Okapi BM25, k1 = 0.9, b = 0.4 (the Lucene defaults BEIR and Bruch et al. used), on title + abstract,
             lowercased [a-z0-9]+ tokens, Porter2 stemming; bm25_nostem is the same without stemming
    minilm   sentence-transformers/all-MiniLM-L6-v2: mean pooling, 256-token truncation, cosine similarity
    bge      BAAI/bge-small-en-v1.5: [CLS] pooling, 512-token truncation, cosine similarity, and the query prefix
             "Represent this sentence for searching relevant passages: " that its model card gives for retrieval
  and fuses BM25 with each embedding model in two ways (Bruch, Gai and Ingber, arXiv 2210.11934v2):
    rrf      reciprocal rank fusion, 1/(eta + rank_bm25) + 1/(eta + rank_dense); eta = 60 is the value
             Cormack, Clarke and Buettcher fixed in their 2009 SIGIR paper
    cc       convex combination alpha * phi(dense) + (1 - alpha) * phi(bm25). phi = tmm divides by the query's
             best score after subtracting the score's theoretical minimum (0 for BM25, -1 for cosine), which is
             Bruch et al.'s TM2C2; phi = mm is plain min-max over the candidate set
  Candidates for fusion are the union of each retriever's top 1,000 (as in Bruch et al.); ranks are exact,
  because every query is scored against the whole corpus. Ties are broken by corpus order.
  Fixed settings (eta 60, alpha 0.8 from Bruch et al.) are compared with settings tuned on the 809 TRAIN claims
  only; the 300 test claims are never used to choose anything. It reports nDCG@10 (BEIR's metric), nDCG@100 and
  Recall@100 (Bruch et al.'s SciFact metrics), paired bootstrap intervals on test deltas, per-query wins and
  losses, and how many training claims it takes to tune each fusion.

Inputs
  Downloads once into --cache (default ./hybrid-search-cache) from huggingface.co, at pinned revisions:
    BeIR/scifact (corpus + queries parquet) and BeIR/scifact-qrels (train.tsv, test.tsv); SciFact's original
    licence is CC BY-NC 2.0 (BEIR paper, Appendix E), so use it for research only
    the two embedding models above
  huggingface.co's robots.txt is read first with this script's own user agent; the script stops if it is not
  allowed. Downloads go through huggingface_hub with that user agent appended. Embeddings are cached as .npy.
  --csv FILE writes one row per (split, query, system).

Packages (versions it was run with): numpy 2.3.5, torch 2.12.0 (CPU), transformers 4.56.2, huggingface_hub 0.36.2,
pyarrow 24.0.0, snowballstemmer 3.0.1; matplotlib 3.10.9 for --figure only. Python 3.13.12.
Run: python hybrid-search-bm25-dense-fusion.py --csv hybrid-search-bm25-dense-fusion.csv [--figure fusion.png]
Time on a 32-thread CPU without a GPU: about 6 minutes the first time (mostly encoding), under 1 minute after.
"""
import argparse, csv, math, os, random, re, sys, time, urllib.request, urllib.robotparser
from collections import Counter, defaultdict

import numpy as np

UA = "hybrid-search-bm25-dense-fusion/1.0 (+https://matrixcognition.com/posts/hybrid-search-bm25-dense-fusion.html)"
DATA = ("BeIR/scifact", "b3b5335604bf5ee3c4447671af975ea25143d4f5")
QRELS = ("BeIR/scifact-qrels", "2938d17dc3b09882fdb8c12bbbe2e2dc0e75a029")
MODELS = {
    "minilm": dict(repo="sentence-transformers/all-MiniLM-L6-v2", rev="1110a243fdf4706b3f48f1d95db1a4f5529b4d41",
                   pooling="mean", max_len=256, query_prefix=""),
    "bge": dict(repo="BAAI/bge-small-en-v1.5", rev="5c38ec7c405ec4b44b94cc5a9bb96e735b38267a",
                pooling="cls", max_len=512, query_prefix="Represent this sentence for searching relevant passages: "),
}
K1, B = 0.9, 0.4
POOL = 1000                                    # candidates per retriever before fusion
ALPHAS = [round(0.05 * i, 2) for i in range(21)]
ETAS = [1, 2, 5, 10, 20, 30, 40, 60, 80, 100, 150, 200, 300, 500, 1000]
SEED = 20261013


def robots_ok(url="https://huggingface.co/robots.txt"):
    req = urllib.request.Request(url, headers={"User-Agent": UA, "Accept": "text/plain"})
    rp = urllib.robotparser.RobotFileParser()
    rp.parse(urllib.request.urlopen(req, timeout=30).read().decode("utf-8", "replace").splitlines())
    return all(rp.can_fetch(a, "https://huggingface.co/datasets/BeIR/scifact") for a in (UA, "*"))


def fetch(repo, rev, filename, cache, kind):
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo, filename, revision=rev, repo_type=kind, cache_dir=cache, user_agent=UA)


def load_scifact(cache):
    import pyarrow.parquet as pq
    corpus = pq.read_table(fetch(*DATA, "corpus/corpus-00000-of-00001.parquet", cache, "dataset")).to_pylist()
    queries = pq.read_table(fetch(*DATA, "queries/queries-00000-of-00001.parquet", cache, "dataset")).to_pylist()
    qrels = {}
    for split in ("train", "test"):
        with open(fetch(*QRELS, split + ".tsv", cache, "dataset"), encoding="utf-8") as f:
            rows = list(csv.reader(f, delimiter="\t"))[1:]
        for qid, did, score in rows:
            if int(score) > 0:
                qrels.setdefault(split, defaultdict(set))[qid].add(did)
    return corpus, {q["_id"]: q["text"] for q in queries}, qrels


# ---------------------------------------------------------------- BM25 (inverted index, whole-corpus scores)
def tokenize(text, stem=None):
    toks = re.findall(r"[a-z0-9]+", text.lower())
    return stem.stemWords(toks) if stem else toks


class BM25:
    def __init__(self, docs, k1=K1, b=B):
        self.k1, self.b, self.n = k1, b, len(docs)
        self.dl = np.array([len(d) for d in docs], dtype=np.float64)
        self.avgdl = self.dl.mean()
        self.post = defaultdict(list)                       # term -> [(doc index, tf)]
        for i, d in enumerate(docs):
            for t, tf in Counter(d).items():
                self.post[t].append((i, tf))
        self.norm = k1 * (1 - b + b * self.dl / self.avgdl)

    def idf(self, t):                                       # Lucene's BM25 idf, never negative
        df = len(self.post.get(t, ()))
        return math.log(1 + (self.n - df + 0.5) / (df + 0.5))

    def scores(self, query_tokens):
        s = np.zeros(self.n)
        for t in query_tokens:                              # a repeated query term counts each time
            if t not in self.post:
                continue
            ids, tfs = zip(*self.post[t])
            ids, tfs = np.array(ids), np.array(tfs, dtype=np.float64)
            s[ids] += self.idf(t) * tfs * (self.k1 + 1) / (tfs + self.norm[ids])
        return s


# ---------------------------------------------------------------- dense encoders (plain transformers)
def encode(texts, spec, cache, batch=32):
    import torch
    from transformers import AutoModel, AutoTokenizer
    path = os.path.dirname(fetch(spec["repo"], spec["rev"], "config.json", cache, "model"))
    for f in ("model.safetensors", "tokenizer.json", "tokenizer_config.json", "vocab.txt", "special_tokens_map.json"):
        fetch(spec["repo"], spec["rev"], f, cache, "model")
    tok, model = AutoTokenizer.from_pretrained(path), AutoModel.from_pretrained(path).eval()
    order = sorted(range(len(texts)), key=lambda i: len(texts[i]))    # similar lengths per batch = less padding
    out = np.zeros((len(texts), model.config.hidden_size), dtype=np.float32)
    with torch.inference_mode():
        for s in range(0, len(order), batch):
            idx = order[s:s + batch]
            enc = tok([texts[i] for i in idx], padding=True, truncation=True, max_length=spec["max_len"],
                      return_tensors="pt")
            h = model(**enc).last_hidden_state
            if spec["pooling"] == "cls":
                v = h[:, 0]
            else:
                m = enc["attention_mask"].unsqueeze(-1).to(h.dtype)
                v = (h * m).sum(1) / m.sum(1).clamp(min=1e-9)
            out[idx] = torch.nn.functional.normalize(v, dim=-1).numpy()
    return out


def embeddings(name, corpus, qids, qtext, cache):
    spec, f = MODELS[name], os.path.join(cache, f"emb_{name}.npz")
    if os.path.exists(f):
        z = np.load(f)
        return z["docs"], z["queries"]
    docs = encode([(d["title"] + " " + d["text"]).strip() for d in corpus], spec, cache)
    qs = encode([spec["query_prefix"] + qtext[q] for q in qids], spec, cache)
    np.savez(f, docs=docs, queries=qs)
    return docs, qs


# ---------------------------------------------------------------- metrics and fusion
def ranking(scores):
    return np.argsort(-scores, kind="stable")


def metrics(order, rel_idx):
    """binary nDCG@10, nDCG@100, Recall@100 for one ranked list of doc indices"""
    out = []
    for k in (10, 100):
        dcg = sum(1 / math.log2(r + 2) for r, d in enumerate(order[:k]) if d in rel_idx)
        idcg = sum(1 / math.log2(r + 2) for r in range(min(len(rel_idx), k)))
        out.append(dcg / idcg)
    out.append(len(rel_idx.intersection(order[:100].tolist())) / len(rel_idx))
    return out


def prepare(lex, sem):
    """the candidate union of both top-POOL lists, and each retriever's exact 1-based rank of every document"""
    lo, so = ranking(lex), ranking(sem)
    rl, rs = np.empty_like(lo), np.empty_like(so)
    rl[lo], rs[so] = np.arange(1, len(lo) + 1), np.arange(1, len(so) + 1)
    return np.union1d(lo[:POOL], so[:POOL]), rl, rs


def fuse(lex, sem, method, param, norm="tmm", prep=None):
    """lex, sem: whole-corpus scores for one query. Returns fused scores (-inf outside the candidate union)."""
    cand, rl, rs = prep or prepare(lex, sem)
    fused = np.full(lex.shape, -np.inf)
    if method == "rrf":
        fused[cand] = 1 / (param + rl[cand]) + 1 / (param + rs[cand])
    else:
        l, s = lex[cand], sem[cand]
        if norm == "tmm":
            ln, sn = l / l.max(), (s + 1) / (s.max() + 1)
        else:
            ln, sn = (l - l.min()) / (l.max() - l.min()), (s - s.min()) / (s.max() - s.min())
        fused[cand] = param * sn + (1 - param) * ln
    return fused


def boot(a, b, n=10000, seed=SEED):
    """paired bootstrap of mean(a - b): (delta, 2.5th, 97.5th percentile, share of resamples <= 0)"""
    d = np.asarray(a) - np.asarray(b)
    idx = np.random.default_rng(seed).integers(0, len(d), size=(n, len(d)))
    m = d[idx].mean(1)
    return d.mean(), np.percentile(m, 2.5), np.percentile(m, 97.5), float((m <= 0).mean())


def draw_figure(path, grid, chosen):
    """the article's figure: mean nDCG@10 against alpha for both embedding models, train and test, with RRF eta=60"""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    surface, ink, ink2, gridc, base = "#fcfcfb", "#0b0b0b", "#52514e", "#e1e0d9", "#c3c2b7"
    blue, orange, aqua = "#2a78d6", "#eb6834", "#1baf7a"        # categorical slots 1-3, fixed order
    mk = dict(ms=4.5, mec=surface, mew=1.0, lw=1.6, solid_capstyle="round", solid_joinstyle="round")
    fig, axes = plt.subplots(1, 2, figsize=(10.4, 4.1), dpi=110, facecolor=surface, sharey=True)
    for ax, dense, title in zip(axes, MODELS, ("BM25 + all-MiniLM-L6-v2", "BM25 + bge-small-en-v1.5")):
        test = [grid[(dense, "cc", a, "tmm")]["test"].mean() for a in ALPHAS]
        train = [grid[(dense, "cc", a, "tmm")]["train"].mean() for a in ALPHAS]
        rrf = grid[(dense, "rrf", 60, "-")]["test"].mean()
        ax.plot(ALPHAS, train, color=orange, marker="o", label="convex blend, 809 train claims", **mk)
        ax.plot(ALPHAS, test, color=blue, marker="o", label="convex blend, 300 test claims", **mk)
        ax.axhline(rrf, color=aqua, lw=1.6, ls=(0, (4, 3)), label="RRF, eta = 60, test claims")
        a = chosen[dense]["cc_tmm"]
        ax.axvline(a, color=base, lw=1.0, zorder=1)
        ax.annotate(f"chosen on train: {a:.2f}", (a, 0.640), xytext=(-6, 0), textcoords="offset points",
                    ha="right", va="bottom", fontsize=8.2, color=ink2)
        ax.axvline(0.8, color=base, lw=1.0, ls=(0, (1, 2)), zorder=1)
        ax.annotate("0.8 (Bruch et al.)", (0.8, 0.781), xytext=(5, 0), textcoords="offset points",
                    ha="left", va="top", fontsize=8.2, color=ink2)
        ax.set_title(title, color=ink, fontsize=10.5, loc="left")
        ax.set_facecolor(surface)
        ax.grid(True, color=gridc, lw=0.6)
        ax.set_axisbelow(True)
        ax.set_xlim(-0.03, 1.03)
        ax.set_ylim(0.635, 0.785)
        ax.set_xlabel("alpha (0 = BM25 only, 1 = embeddings only)", color=ink2)
        for sp in ("top", "right"):
            ax.spines[sp].set_visible(False)
        for sp in ("left", "bottom"):
            ax.spines[sp].set_color(base)
        ax.tick_params(colors=ink2, labelsize=8.5)
    axes[0].set_ylabel("mean nDCG@10", color=ink2)
    axes[0].legend(frameon=False, fontsize=8.2, loc="upper left", labelcolor=ink, bbox_to_anchor=(0, 0.93))
    fig.tight_layout()
    fig.savefig(path, dpi=110, facecolor=surface)
    plt.close(fig)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cache", default="hybrid-search-cache")
    ap.add_argument("--csv")
    ap.add_argument("--figure", help="also draw the article's figure to this PNG (needs matplotlib)")
    args = ap.parse_args()
    os.makedirs(args.cache, exist_ok=True)
    if not robots_ok():
        sys.exit("huggingface.co robots.txt does not allow this fetch; stopping")
    t0 = time.time()
    corpus, qtext, qrels = load_scifact(args.cache)
    did_index = {d["_id"]: i for i, d in enumerate(corpus)}
    splits = {s: sorted(qrels[s], key=int) for s in ("train", "test")}
    qids = splits["train"] + splits["test"]
    rel = {q: {did_index[d] for d in qrels[s][q]} for s in splits for q in splits[s]}
    print(f"SciFact: {len(corpus)} abstracts, {len(splits['train'])} train + {len(splits['test'])} test claims, "
          f"{sum(len(qrels['test'][q]) for q in splits['test'])} test judgments")

    import snowballstemmer
    stem = snowballstemmer.stemmer("english")
    raw = {}
    for name, st in (("bm25", stem), ("bm25_nostem", None)):
        idx = BM25([tokenize(d["title"] + " " + d["text"], st) for d in corpus])
        raw[name] = np.array([idx.scores(tokenize(qtext[q], st)) for q in qids])
    for name in MODELS:
        docs, qs = embeddings(name, corpus, qids, qtext, args.cache)
        raw[name] = qs @ docs.T
    print(f"scored in {time.time() - t0:.0f} s")

    per = {}                                          # (system, qid) -> [ndcg10, ndcg100, recall100]

    def run(system, score_fn):
        for i, q in enumerate(qids):
            per[(system, q)] = metrics(ranking(score_fn(i)), rel[q])

    def mean(system, split, m=0):
        return float(np.mean([per[(system, q)][m] for q in splits[split]]))

    for name in raw:
        run(name, lambda i, n=name: raw[n][i])
    grid = {}                                         # (dense, method, param, norm) -> per-query nDCG@10 by split
    settings = [(m, p, n) for m, ps, n in (("cc", ALPHAS, "tmm"), ("cc", ALPHAS, "mm"), ("rrf", ETAS, "-")) for p in ps]
    ntr = len(splits["train"])
    for dense in MODELS:
        v = np.zeros((len(settings), len(qids)))
        for i, q in enumerate(qids):
            prep = prepare(raw["bm25"][i], raw[dense][i])
            for j, (m, p, n) in enumerate(settings):
                v[j, i] = metrics(ranking(fuse(raw["bm25"][i], raw[dense][i], m, p, n, prep)), rel[q])[0]
        for j, (m, p, n) in enumerate(settings):
            grid[(dense, m, p, n)] = {"train": v[j, :ntr], "test": v[j, ntr:]}

    def tuned(dense, method, norm, sub=None):
        """the grid value with the best mean train nDCG@10 (first in grid order on a tie); sub = train positions"""
        params = ALPHAS if method == "cc" else ETAS
        best = [grid[(dense, method, p, norm)]["train"] if sub is None else grid[(dense, method, p, norm)]["train"][sub]
                for p in params]
        return params[int(np.argmax([b.mean() for b in best]))]

    chosen = {}
    for dense in MODELS:
        chosen[dense] = dict(cc_tmm=tuned(dense, "cc", "tmm"), cc_mm=tuned(dense, "cc", "mm"),
                             rrf=tuned(dense, "rrf", "-"))
        c = chosen[dense]
        for system, method, p, norm in ((f"rrf60_{dense}", "rrf", 60, "-"), (f"tm2c2_a0.8_{dense}", "cc", 0.8, "tmm"),
                                        (f"tm2c2_tuned_{dense}", "cc", c["cc_tmm"], "tmm"),
                                        (f"minmax_tuned_{dense}", "cc", c["cc_mm"], "mm"),
                                        (f"rrf_tuned_{dense}", "rrf", c["rrf"], "-")):
            run(system, lambda i, m=method, p=p, n=norm, d=dense: fuse(raw["bm25"][i], raw[d][i], m, p, n))

    systems = list(raw) + [s for d in MODELS for s in (f"rrf60_{d}", f"tm2c2_a0.8_{d}", f"tm2c2_tuned_{d}",
                                                         f"minmax_tuned_{d}", f"rrf_tuned_{d}")]
    print("\nTEST (300 claims)        nDCG@10  nDCG@100  R@100   | train nDCG@10")
    for s in systems:
        print(f"{s:24s} {mean(s, 'test'):.4f}   {mean(s, 'test', 1):.4f}   {mean(s, 'test', 2):.4f}  | "
              f"{mean(s, 'train'):.4f}")
    print("\ntuned on train only:", chosen)

    for dense in MODELS:
        print(f"\n--- BM25 + {dense}, alpha sweep (TM2C2) mean nDCG@10   train / test")
        for a in ALPHAS:
            g = grid[(dense, "cc", a, "tmm")]
            print(f"  alpha {a:.2f}  {g['train'].mean():.4f} / {g['test'].mean():.4f}")
        print(f"--- eta sweep (RRF) mean nDCG@10   train / test")
        for e in ETAS:
            g = grid[(dense, "rrf", e, "-")]
            print(f"  eta {e:4d}  {g['train'].mean():.4f} / {g['test'].mean():.4f}")
        best_single = max(("bm25", dense), key=lambda s: mean(s, "train"))
        print(f"--- paired bootstrap on test nDCG@10 (10,000 resamples): delta [95% interval] share<=0")
        for a, b in ((f"rrf60_{dense}", best_single), (f"tm2c2_a0.8_{dense}", best_single),
                     (f"tm2c2_tuned_{dense}", best_single), (f"tm2c2_tuned_{dense}", f"rrf60_{dense}"),
                     (f"tm2c2_a0.8_{dense}", f"rrf60_{dense}"), (f"rrf_tuned_{dense}", f"rrf60_{dense}")):
            d, lo, hi, p = boot([per[(a, q)][0] for q in splits["test"]], [per[(b, q)][0] for q in splits["test"]])
            print(f"  {a} - {b}: {d:+.4f} [{lo:+.4f}, {hi:+.4f}] {p:.4f}")
        h = np.array([per[(f"tm2c2_tuned_{dense}", q)][0] for q in splits["test"]])
        bs = np.array([per[(best_single, q)][0] for q in splits["test"]])
        print(f"  tm2c2_tuned vs {best_single}, per test claim: better {int((h > bs + 1e-12).sum())}, "
              f"same {int((abs(h - bs) <= 1e-12).sum())}, worse {int((h < bs - 1e-12).sum())}")
        orc = np.max([grid[(dense, "cc", a, "tmm")]["test"] for a in ALPHAS], axis=0)
        print(f"  oracle (best alpha per test claim): {orc.mean():.4f}")
        top = [(set(ranking(raw["bm25"][qids.index(q)])[:10]), set(ranking(raw[dense][qids.index(q)])[:10]))
               for q in splits["test"]]
        print(f"  mean top-10 overlap BM25 vs {dense}: {np.mean([len(a & b) for a, b in top]):.2f} of 10")
        hits = [(bool(a & rel[q]), bool(b & rel[q])) for (a, b), q in zip(top, splits["test"])]
        print(f"  test claims with a relevant abstract in the top 10: both {sum(a and b for a, b in hits)}, "
              f"BM25 only {sum(a and not b for a, b in hits)}, {dense} only {sum(b and not a for a, b in hits)}, "
              f"neither {sum(not a and not b for a, b in hits)}")
        print(f"--- tuning from a few train claims (200 draws each): test nDCG@10 median / 5th percentile")
        rng = random.Random(SEED)
        for n in (5, 10, 20, 50, 100, 200):
            res = {"cc": [], "rrf": []}
            for _ in range(200):
                sub = np.array(rng.sample(range(len(splits["train"])), n))
                res["cc"].append(grid[(dense, "cc", tuned(dense, "cc", "tmm", sub), "tmm")]["test"].mean())
                res["rrf"].append(grid[(dense, "rrf", tuned(dense, "rrf", "-", sub), "-")]["test"].mean())
            print(f"  n={n:3d}  TM2C2 {np.median(res['cc']):.4f} / {np.percentile(res['cc'], 5):.4f}   "
                  f"RRF {np.median(res['rrf']):.4f} / {np.percentile(res['rrf'], 5):.4f}")

    if args.csv:
        with open(args.csv, "w", newline="", encoding="utf-8") as f:
            w = csv.writer(f, lineterminator="\n")
            w.writerow(["split", "query_id", "n_relevant", "system", "ndcg_at_10", "ndcg_at_100", "recall_at_100"])
            for s in ("train", "test"):
                for q in splits[s]:
                    for system in systems:
                        m = per[(system, q)]
                        w.writerow([s, q, len(rel[q]), system, f"{m[0]:.6f}", f"{m[1]:.6f}", f"{m[2]:.6f}"])
        print(f"\nwrote {args.csv}")
    if args.figure:
        draw_figure(args.figure, grid, chosen)
        print(f"drew {args.figure}")
    print(f"done in {time.time() - t0:.0f} s")


if __name__ == "__main__":
    main()