"""Pick queries for teacher scoring.   python rerank/subsample.py MINED.parquet PASSAGES.parquet OUT.parquet [PER_LANG] [ENG] [K]
Stratified by the source passage's language: up to PER_LANG queries per language (ENG for English), keeping the
question/keywords/english mix; candidates cut to the top K (source first, then the hardest retrieved negatives)."""
import sys
import pandas as pd
m, p, out = sys.argv[1:4]
per, eng, K = (int(sys.argv[4]) if len(sys.argv) > 4 else 3200), (int(sys.argv[5]) if len(sys.argv) > 5 else 20000), (int(sys.argv[6]) if len(sys.argv) > 6 else 16)
df = pd.read_parquet(m); lang = dict(pd.read_parquet(p)[["pid", "lang"]].values); df["plang"] = df.pid.map(lang)
df = pd.concat([g.sample(min(len(g), eng if l == "eng_Latn" else per), random_state=0) for l, g in df.groupby("plang")])
df["cands"] = df.cands.map(lambda c: list(c)[:K])
df.sample(frac=1.0, random_state=0).reset_index(drop=True).to_parquet(out)
print(len(df), "queries", len(df) * K, "pairs", df.qtype.value_counts().to_dict())