#!/usr/bin/env python3
"""T0.3 · Fixed item-id files for RS-full (spec §7), so every G4 run scores the same items.

    python -I scripts/prep_rs_full.py               # downloads ~250 MB of parquet into $MT_DATA/evalsets

- MMLU, 5-shot: all 14,042 test items; the shots are each subject's 5 dev items in order (the standard
  MMLU prompt), written into the prompt so lm-eval itself runs 0-shot.
- MMLU-Redux 2.0: every item its annotators marked `ok`, 0-shot.
- MMLU-Pro: 2,000 of the 12,032 test items, stratified over the 14 categories (largest remainder), seed 0.
- GSM8K: all 1,319 test items. MATH-500: all 500. IFEval: all 541 prompts.
- Global-MMLU: 1,000 items per Marco language it covers (24 of 29; th, ur, kk, hu and az are absent),
  stratified over the 57 subjects, seed 0, the same sample ids in every language (it is parallel).
- BELEBELE: all 900 questions in all 29 Marco languages.

Parquet comes from the Hub's parquet conversion; each file's sha256 is recorded with the dataset's
revision at download time. Downloaded files are data: read with pyarrow only.
"""
from __future__ import annotations

import argparse
import json
import random
import sys
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))  # run with python -I, which drops the script's dir
from prep_rs_subsets import DATA, OUT as OUT_SUBSETS, REV, allocate, rows, write  # noqa: E402

MARCO = "en zh ar de es fr ko ja pt tr id it nl pl ru vi th he uk ms bn cs ur kk el ro hu ne az".split()
FLORES = dict(en="eng_Latn", zh="zho_Hans", ar="arb_Arab", de="deu_Latn", es="spa_Latn", fr="fra_Latn",
              ko="kor_Hang", ja="jpn_Jpan", pt="por_Latn", tr="tur_Latn", id="ind_Latn", it="ita_Latn",
              nl="nld_Latn", pl="pol_Latn", ru="rus_Cyrl", vi="vie_Latn", th="tha_Thai", he="heb_Hebr",
              uk="ukr_Cyrl", ms="zsm_Latn", bn="ben_Beng", cs="ces_Latn", ur="urd_Arab", kk="kaz_Cyrl",
              el="ell_Grek", ro="ron_Latn", hu="hun_Latn", ne="npi_Deva", az="azj_Latn")
GMMLU_ABSENT = {"th", "ur", "kk", "hu", "az"}
REV_FULL = {"mmlu": ("cais/mmlu", "c30699e8356da336a370243923dbaf21066bb9fe"),
            "mmlu_pro": ("TIGER-Lab/MMLU-Pro", "b189ec765aa7ed75c8acfea42df31fdae71f97be"),
            "math500": ("HuggingFaceH4/MATH-500", "6e4ed1a2a79af7d8630a6b768ec859cb5af4d3be"),
            "ifeval": ("google/IFEval", "966cd89545d6b6acfd7638bc708b98261ca58e84"),
            "global_mmlu": ("CohereLabs/Global-MMLU", "0e619dbeb34206cd48705a1a0ea7fb21cae09993")}


def stratified(groups: dict[str, list], n: int, seed: int) -> dict[str, list]:
    alloc = allocate({k: len(v) for k, v in groups.items()}, n)
    rng = random.Random(seed)
    return {k: sorted(rng.sample(groups[k], alloc[k])) for k in sorted(groups)}


def mmlu() -> None:
    ds, rev = REV_FULL["mmlu"]
    test, sha_t = rows(ds, "all", "test", DATA / "evalsets/mmlu/all_test.parquet")
    dev, sha_d = rows(ds, "all", "dev", DATA / "evalsets/mmlu/all_dev.parquet")
    shots = defaultdict(list)
    for i, r in enumerate(dev):
        shots[r["subject"]].append(i)
    assert all(len(v) == 5 for v in shots.values()) and len(shots) == 57
    write("mmlu_5shot.json", {"dataset": ds, "revision": rev, "split": "test", "n": len(test), "shots": "dev, 5 per subject, in order",
                              "parquet_sha256": {"test": sha_t, "dev": sha_d}})


def mmlu_redux_ok() -> None:
    ds, rev = REV["mmlu_redux"]
    shas = json.loads((OUT_SUBSETS / "mmlu_redux_1000.json").read_text())["parquet_sha256"]
    ids = {}
    for s in sorted(shas):
        rs, sha = rows(ds, s, "test", DATA / f"evalsets/mmlu_redux/{s}_test.parquet")
        assert sha == shas[s], s
        ids[s] = [i for i, r in enumerate(rs) if r.get("error_type") == "ok"]
    write("mmlu_redux_ok.json", {"dataset": ds, "revision": rev, "split": "test", "filter": "error_type == ok",
                                 "n": sum(map(len, ids.values())), "parquet_sha256": shas, "ids": ids})


def mmlu_pro(seed: int, n: int) -> None:
    ds, rev = REV_FULL["mmlu_pro"]
    rs, sha = rows(ds, "default", "test", DATA / "evalsets/mmlu_pro/test.parquet")
    by_cat = defaultdict(list)
    for r in rs:
        by_cat[r["category"]].append(r["question_id"])
    ids = stratified(by_cat, n, seed)
    write("mmlu_pro_2000.json", {"dataset": ds, "revision": rev, "split": "test", "seed": seed, "n": sum(map(len, ids.values())),
                                 "population": len(rs), "parquet_sha256": sha, "ids": ids})


def whole(key: str, cfg: str, split: str, dest: str, name: str, id_field: str | None = None) -> None:
    ds, rev = REV_FULL[key]
    rs, sha = rows(ds, cfg, split, DATA / dest)
    obj = {"dataset": ds, "revision": rev, "split": split, "n": len(rs), "parquet_sha256": sha}
    if id_field:
        obj["ids"] = [r[id_field] for r in rs]
    write(name, obj)


def gsm8k_full() -> None:
    g = json.loads((OUT_SUBSETS / "gsm8k_250.json").read_text())
    rs, sha = rows("openai/gsm8k", "main", "test", DATA / "evalsets/gsm8k/main_test.parquet")
    assert sha == g["parquet_sha256"]
    write("gsm8k_full.json", {"dataset": g["dataset"], "revision": g["dataset_sha"], "split": "test", "n": len(rs), "parquet_sha256": sha})


def global_mmlu(seed: int, n: int) -> None:
    ds, rev = REV_FULL["global_mmlu"]
    langs = [l for l in MARCO if l not in GMMLU_ABSENT]
    out = {"dataset": ds, "revision": rev, "split": "test", "seed": seed, "languages": langs,
           "absent_marco_languages": sorted(GMMLU_ABSENT), "parquet_sha256": {}}
    present = None
    subj = {}
    for lang in langs:
        rs, out["parquet_sha256"][lang] = rows(ds, lang, "test", DATA / f"evalsets/global_mmlu/{lang}_test.parquet")
        sids = {r["sample_id"] for r in rs}
        present = sids if present is None else present & sids
        if lang == "en":
            subj = {r["sample_id"]: r["subject"] for r in rs}
    by_subj = defaultdict(list)
    for sid in sorted(present):
        by_subj[subj[sid]].append(sid)
    pick = stratified(by_subj, n, seed)
    out.update(n_per_language=n, parallel_items=len(present), ids=sorted(s for v in pick.values() for s in v))
    out["n"] = n * len(langs)
    write("global_mmlu_1000.json", out)


def belebele_full() -> None:
    ds, rev = REV["belebele"]
    out = {"dataset": ds, "revision": rev, "split": "test", "languages": MARCO, "flores": FLORES, "parquet_sha256": {}}
    keys = None
    for lang in MARCO:
        rs, out["parquet_sha256"][lang] = rows(ds, FLORES[lang], "test", DATA / f"evalsets/belebele/{FLORES[lang]}_test.parquet")
        k = {(r["link"], r["question_number"]) for r in rs}
        keys = k if keys is None else keys & k
    out.update(n_per_language=len(keys), n=len(keys) * len(MARCO))
    write("belebele_full.json", out)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--seed", type=int, default=0)
    a = p.parse_args()
    mmlu()
    mmlu_redux_ok()
    mmlu_pro(a.seed, 2000)
    gsm8k_full()
    whole("math500", "default", "test", "evalsets/math500/test.parquet", "math500.json", "unique_id")
    whole("ifeval", "default", "train", "evalsets/ifeval/train.parquet", "ifeval.json", "key")
    global_mmlu(a.seed, 1000)
    belebele_full()


if __name__ == "__main__":
    main()