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

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

- MMLU-Redux 2.0: 1,000 of the items its annotators marked `ok` (label verified), stratified over the
  57 subjects (largest-remainder allocation), seed 0.
- Global-MMLU-Lite: all 400 test items for each RS-fast language (the set is already fixed and
  human-verified; the file pins the ids and the parquet hash).
- BELEBELE: 200 of the 900 questions, seed 0, the same questions in every RS-fast language (BELEBELE
  is fully parallel), so cross-language differences are not item differences.
- GSM8K-250 already exists (configs/eval/subsets/gsm8k_250.json, made for the T1.4 FP8 check).

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

import argparse
import hashlib
import json
import os
import random
import urllib.request
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
DATA = Path(os.environ.get("MT_DATA", "/path/to/data"))
OUT = REPO / "configs" / "eval" / "subsets"
RS_FAST = ["zh", "ar", "ko", "hu", "bn"]  # T0.2 (D-12)
BELEBELE_CFG = {"zh": "zho_Hans", "ar": "arb_Arab", "ko": "kor_Hang", "hu": "hun_Latn", "bn": "ben_Beng",
                "en": "eng_Latn"}
REV = {"mmlu_redux": ("edinburgh-dawg/mmlu-redux-2.0", "372ea425445d51e1ba1188c56e5e893f8138621f"),
       "global_mmlu_lite": ("CohereLabs/Global-MMLU-Lite", "36c2fd756f19ccf13a9a96c8e53ccecc02192b8b"),
       "belebele": ("facebook/belebele", "7899cdfa4e1e0d733fd77c848e2c273cb1d32be2")}
PARQUET = "https://huggingface.co/api/datasets/{ds}/parquet/{cfg}/{split}/0.parquet"


def fetch(url: str, dest: Path) -> str:
    dest.parent.mkdir(parents=True, exist_ok=True)
    if not dest.exists():
        tmp = dest.with_suffix(".part")
        urllib.request.urlretrieve(url, tmp)
        tmp.rename(dest)
    return hashlib.sha256(dest.read_bytes()).hexdigest()


def rows(ds: str, cfg: str, split: str, dest: Path) -> tuple[list[dict], str]:
    import pyarrow.parquet as pq
    sha = fetch(PARQUET.format(ds=ds, cfg=cfg, split=split), dest)
    return pq.read_table(dest).to_pylist(), sha


def write(name: str, obj: dict) -> None:
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / name).write_text(json.dumps(obj, indent=1, ensure_ascii=False) + "\n")
    print(f"{name}: {obj.get('n')} items")


def allocate(sizes: dict[str, int], total: int) -> dict[str, int]:
    """Largest-remainder allocation of `total` proportional to `sizes`, capped at each size."""
    n = sum(sizes.values())
    quota = {k: total * v / n for k, v in sizes.items()}
    alloc = {k: min(sizes[k], int(q)) for k, q in quota.items()}
    for k in sorted(sizes, key=lambda k: (-(quota[k] - int(quota[k])), k)):
        if sum(alloc.values()) >= total:
            break
        if alloc[k] < sizes[k]:
            alloc[k] += 1
    return alloc


def mmlu_redux(seed: int, n: int) -> None:
    import urllib.request as ur
    ds, rev = REV["mmlu_redux"]
    subjects = sorted(json.loads(ur.urlopen(f"https://huggingface.co/api/datasets/{ds}/parquet").read()))
    ok, shas = {}, {}
    for s in subjects:
        rs, shas[s] = rows(ds, s, "test", DATA / f"evalsets/mmlu_redux/{s}_test.parquet")
        ok[s] = [i for i, r in enumerate(rs) if r.get("error_type") == "ok"]
    alloc = allocate({s: len(v) for s, v in ok.items()}, n)
    rng = random.Random(seed)
    items = {s: sorted(rng.sample(ok[s], alloc[s])) for s in subjects}
    write("mmlu_redux_1000.json", {"dataset": ds, "revision": rev, "split": "test", "filter": "error_type == ok",
                                   "seed": seed, "n": sum(alloc.values()), "ok_items": sum(map(len, ok.values())),
                                   "parquet_sha256": shas, "ids": items})


def global_mmlu_lite() -> None:
    ds, rev = REV["global_mmlu_lite"]
    out = {"dataset": ds, "revision": rev, "split": "test", "languages": RS_FAST, "parquet_sha256": {}, "ids": {}}
    for lang in RS_FAST:
        rs, out["parquet_sha256"][lang] = rows(ds, lang, "test", DATA / f"evalsets/global_mmlu_lite/{lang}_test.parquet")
        out["ids"][lang] = [r["sample_id"] for r in rs]
    out["n"] = sum(len(v) for v in out["ids"].values())
    write("global_mmlu_lite_rs5.json", out)


def belebele(seed: int, n: int) -> None:
    ds, rev = REV["belebele"]
    out = {"dataset": ds, "revision": rev, "split": "test", "languages": RS_FAST, "seed": seed, "parquet_sha256": {}}
    keys = None
    for lang in RS_FAST:
        cfg = BELEBELE_CFG[lang]
        rs, out["parquet_sha256"][lang] = rows(ds, cfg, "test", DATA / f"evalsets/belebele/{cfg}_test.parquet")
        k = {(r["link"], r["question_number"]) for r in rs}
        keys = k if keys is None else keys & k
    pick = sorted(random.Random(seed).sample(sorted(keys), n))
    out.update(n=n, parallel_questions=len(keys), ids=[{"link": l, "question_number": q} for l, q in pick])
    write("belebele_200.json", out)


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--seed", type=int, default=0)
    a = p.parse_args()
    mmlu_redux(a.seed, 1000)
    global_mmlu_lite()
    belebele(a.seed, 200)


if __name__ == "__main__":
    main()