"""Reranking evaluation sets from MTEB reranking datasets (evaluation only).   python rerank/build_evals.py OUT_DIR
Each output parquet = one (dataset, language): rows qid, query, docs (candidate texts, in the dataset's candidate order),
rels (graded relevance per candidate). Queries are sampled (seed 0) up to CAP per set; all candidates of a query are kept.
Sets: MIRACL (18 languages, dev, CC-BY-SA-4.0), WikipediaRerankingMultilingual (CC-BY-SA-3.0), ESCI (product search,
Apache-2.0), RuBQ (ru), T2Reranking (zh), VoyageMMarco (ja), AskUbuntuDupQuestions and StackOverflowDupQuestions (en).
"""
import os, sys
import pandas as pd
from huggingface_hub import HfApi, hf_hub_download

OUT = sys.argv[1]
os.makedirs(OUT, exist_ok=True)
api = HfApi()
SETS = [("mteb/MIRACLReranking", "miracl", 60), ("mteb/WikipediaRerankingMultilingual", "wiki", 60), ("mteb/ESCIReranking", "esci", 150),
        ("mteb/RuBQReranking", "rubq", 150), ("mteb/T2Reranking", "t2", 150), ("mteb/VoyageMMarcoReranking", "mmarco", 150),
        ("mteb/AskUbuntuDupQuestions", "askubuntu", 361), ("mteb/stackoverflowdupquestions-reranking", "stackoverflow", 300)]


def load(repo, files, part, split):
    cand = [f for f in files if (f.startswith(part + "/") or f.split("/")[0].endswith("-" + part)) and f.split("/")[-1].startswith(split)]
    return cand


for repo, name, cap in SETS:
    files = [s.rfilename for s in api.dataset_info(repo).siblings if s.rfilename.endswith(".parquet")]
    split = "test" if any("/test-" in f for f in files) else "dev"
    prefixes = sorted({f.split("/")[0].rsplit("-", 1)[0] for f in files if "-" in f.split("/")[0]}) or [""]
    for pre in prefixes:
        def part(p):
            alts = [f"{pre}-{p}/{split}" if pre else f"{p}/{split}"] + ([f"{pre}-qrels/{split}"] if p == "data" else [])
            for f in files:
                if any(f.startswith(a) for a in alts):
                    return pd.read_parquet(hf_hub_download(repo, f, repo_type="dataset"))
            return None
        corpus, queries, top = part("corpus"), part("queries"), part("top_ranked")
        qrels = part("qrels")
        if qrels is None:
            qrels = part("data")
        if corpus is None or queries is None or top is None or qrels is None:
            print(name, pre, "incomplete; skipped"); continue
        ctext = dict(zip(corpus["_id"], (corpus.get("title", "").fillna("") + " " + corpus.text).str.strip().str.slice(0, 2000)))
        qtext = dict(zip(queries["_id"], queries.text))
        rel = {(q, c): s for q, c, s in zip(qrels["query-id"], qrels["corpus-id"], qrels["score"])}
        top = top[top["query-id"].isin(qtext)].sample(frac=1.0, random_state=0)
        rows = []
        for q, cids in zip(top["query-id"], top["corpus-ids"]):
            cids = [c for c in cids if c in ctext]
            rels = [int(rel.get((q, c), 0)) for c in cids]
            if not any(r > 0 for r in rels) or all(r > 0 for r in rels):
                continue
            rows.append((q, qtext[q], [ctext[c] for c in cids], rels))
            if len(rows) >= cap:
                break
        tag = f"{name}_{pre}" if pre else name
        pd.DataFrame(rows, columns=["qid", "query", "docs", "rels"]).to_parquet(f"{OUT}/{tag}.parquet")
        print(tag, len(rows), "queries,", sum(len(r[2]) for r in rows), "pairs", flush=True)