"""Hard-negative candidates with BAAI/bge-m3 dense retrieval (MIT).   python rerank/mine.py PASSAGES.parquet QDIR OUT.parquet [K]
For each generated query, the top-K passages (cosine, CLS-pooled normalised bge-m3 embeddings) from the passages in the
same language as the query's source passage; the source passage is always included. OUT rows: qid, pid (source),
query, qtype, qlang, cands (list of K pids, source first).
"""
import glob, os, sys
import numpy as np, pandas as pd, torch
from transformers import AutoModel, AutoTokenizer

P, QD, OUT = sys.argv[1:4]; K = int(sys.argv[4]) if len(sys.argv) > 4 else 24
ps = pd.read_parquet(P)
qs = pd.concat([pd.read_json(f, lines=True) for f in sorted(glob.glob(f"{QD}/*.jsonl"))], ignore_index=True)
qs = qs.drop_duplicates(["pid", "query"]).reset_index(drop=True); qs["qid"] = range(len(qs))
tok = AutoTokenizer.from_pretrained("BAAI/bge-m3"); m = AutoModel.from_pretrained("BAAI/bge-m3", torch_dtype=torch.float16).cuda().eval()


@torch.no_grad()
def emb(texts, L, bs=256):
    out = np.zeros((len(texts), 1024), dtype=np.float16); order = np.argsort([len(t) for t in texts])
    for s in range(0, len(texts), bs):
        b = order[s:s + bs]
        e = tok([texts[i] for i in b], truncation=True, max_length=L, padding=True, return_tensors="pt").to("cuda")
        h = m(**e).last_hidden_state[:, 0]; out[b] = torch.nn.functional.normalize(h.float(), dim=-1).half().cpu().numpy()
    return out


PE = emb(ps.text.tolist(), 512); QE = emb(qs["query"].tolist(), 64)
plang = dict(zip(ps.pid, ps.lang)); pidx = {p: i for i, p in enumerate(ps.pid)}
qs["plang"] = qs.pid.map(plang)
cands = [None] * len(qs)
for lang, g in qs.groupby("plang"):
    pool = np.where(ps.lang.values == lang)[0]
    A = torch.tensor(PE[pool]).cuda()
    for s in range(0, len(g), 4096):
        gi = g.index.values[s:s + 4096]
        sim = torch.tensor(QE[gi]).cuda() @ A.T
        top = sim.topk(min(K + 1, len(pool)), dim=1).indices.cpu().numpy()
        for row, t in zip(gi, top):
            src = qs.pid[row]; c = [int(ps.pid.values[pool[j]]) for j in t if ps.pid.values[pool[j]] != src][: K - 1]
            cands[row] = [int(src)] + c
    print(lang, len(g), flush=True)
qs["cands"] = cands
qs[["qid", "pid", "query", "qtype", "qlang", "cands"]].to_parquet(OUT); print("queries", len(qs), "pairs", sum(map(len, cands)), flush=True)