"""Run ONE query through the eSCR retrieval pipeline (loads the persisted index on Thor).
dense (BGE bf16, numpy cosine) + BM25 -> RRF -> cross-encoder rerank -> top-K docs.
Prints top-K as JSON: case_name, neutral_citation, date, disposition, bench_strength, passage.

Run on Thor:  Q="your query" python3 18_query_escr.py
"""
import json, os, re, time
import numpy as np, torch
from sentence_transformers import SentenceTransformer, CrossEncoder
from rank_bm25 import BM25Okapi

QUERY = os.environ["Q"]; 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()

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
print(f"{len(chunks)} chunks, {len(meta)} docs", flush=True)

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"bm25 built in {time.time()-t:.0f}s", flush=True)

qv = st.encode(BGE_Q + QUERY, 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(QUERY))
b_top = sorted(range(len(bs)), key=lambda i: bs[i], reverse=True)[:60]
from collections import defaultdict
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([(QUERY, 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"),
                "bench_strength": m.get("bench_strength"), "rr": round(s, 2),
                "passage": re.sub(r"\s+", " ", texts[ci])[:300]})
print("RESULTS_JSON:" + json.dumps(out, ensure_ascii=False))
for i, o in enumerate(out):
    print(f"{i+1}. [{o['neutral_citation']}] {o['case_name']}  ({o['date']}, {o['disposition']}) rr={o['rr']}")