"""Benchmark harness — Themis side. Run all bench_queries through the eSCR pipeline
(loads the persisted index ONCE), save top-K per query for scoring against CaseMine.

Run on Thor:  python3 19_benchmark_themis.py   (expects bench_queries.json alongside)
Out: themis_bench_results.json  [{id,intent,query,results:[{neutral_citation,case_name,date,disposition,rr,passage}]}]
"""
import json, os, re, time
from collections import defaultdict
import numpy as np, torch
from sentence_transformers import SentenceTransformer, CrossEncoder
from rank_bm25 import BM25Okapi

TOPK = int(os.getenv("TOPK", "10")); K = 60; CAND = 40
BGE_Q = "Represent this sentence for searching relevant passages: "
def tok(s): return re.sub(r"[^a-z0-9 ]", " ", (s or "").lower()).split()

queries = json.load(open("bench_queries.json"))
print("loading index...", flush=True)
chunks = [json.loads(l) for l in open("escr_chunks.jsonl")]
texts = [c["text"] for c in chunks]; chunk_doc = [c["doc_id"] for c in chunks]
M = np.load("escr_vectors.npy")
meta = {}
for l in open("escr_meta.jsonl"):
    m = json.loads(l); meta[m["doc_id"]] = m
st = SentenceTransformer("BAAI/bge-small-en-v1.5", device="cuda", model_kwargs={"torch_dtype": torch.bfloat16})
ce = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", device="cuda")
t = time.time(); bm25 = BM25Okapi([tok(t_) for t_ in texts]); print(f"index+bm25 ready ({len(chunks)} chunks) in {time.time()-t:.0f}s", flush=True)

def search(q):
    qv = st.encode(BGE_Q + q, normalize_embeddings=True, convert_to_numpy=True).astype(np.float32)
    dense = M @ qv
    d_top = np.argpartition(-dense, 60)[:60]; d_top = d_top[np.argsort(-dense[d_top])]
    bs = bm25.get_scores(tok(q))
    b_top = sorted(range(len(bs)), key=lambda i: bs[i], reverse=True)[:60]
    sc = defaultdict(float)
    for r, ci in enumerate(d_top): sc[int(ci)] += 1.0 / (K + r)
    for r, ci in enumerate(b_top):
        if bs[ci] > 0: sc[ci] += 1.0 / (K + r)
    cand = [ci for ci, _ in sorted(sc.items(), key=lambda x: x[1], reverse=True)[:CAND]]
    rr = ce.predict([(q, texts[ci]) for ci in cand])
    best = {}
    for ci, s in zip(cand, rr):
        d = chunk_doc[ci]
        if d not in best or s > best[d][0]: best[d] = (float(s), ci)
    out = []
    for d, (s, ci) in sorted(best.items(), key=lambda x: x[1][0], reverse=True)[:TOPK]:
        m = meta.get(d, {})
        out.append({"neutral_citation": m.get("neutral_citation"), "case_name": m.get("case_name"),
                    "date": m.get("date"), "disposition": m.get("disposition"),
                    "rr": round(s, 2), "passage": re.sub(r"\s+", " ", texts[ci])[:350]})
    return out

results = []
for qd in queries:
    res = search(qd["query"])
    results.append({**qd, "results": res})
    print(f"  {qd['id']:12s} top1: {res[0]['case_name'][:50] if res else '-'}", flush=True)
json.dump(results, open("themis_bench_results.json", "w"), ensure_ascii=False, indent=1)
print("wrote themis_bench_results.json", flush=True)