"""Fetch raw pretraining shards into $TA_DATA/raw. Idempotent; re-run to resume.

Sources (all permissive / openly licensed):
  english: HuggingFaceFW/fineweb-edu sample/10BT        (ODC-By)   ~0.7B tok/file
  math:    HuggingFaceTB/finemath finemath-4plus         (ODC-By)   ~0.15B tok/file
  code:    common-pile/stackv2_edu_filtered              (openly licensed code only)
"""
import argparse, os
from huggingface_hub import hf_hub_download

DATA = os.environ.get("TA_DATA", "data")
SOURCES = {
    "english": ("HuggingFaceFW/fineweb-edu", [f"sample/10BT/{i:03d}_00000.parquet" for i in range(14)]),
    "math": ("HuggingFaceTB/finemath", [f"finemath-4plus/train-{i:05d}-of-00064.parquet" for i in range(64)]),
    "code": ("common-pile/stackv2_edu_filtered", None),
}
# stack-edu shards are sorted by language (checked 2026-10-06 by reading each shard's first rows):
# 0-17 Markdown, 18-19 C, 20-27 C++, 28-33 C#, 34-37 Go, 38-49 Java, 50-64 JavaScript, 65-72 PHP,
# 73-84 Python, 85 Ruby, 86 Rust, 87-88 Shell, 89-90 Swift, 91-94 TypeScript.
CODE_SHARDS = {"markdown": [0, 1], "python": [73, 74, 75, 76], "shell": [87, 88], "javascript": [50],
               "typescript": [91], "go": [34], "rust": [86], "c": [18]}

# Small supervised sets rendered into the synthetic sources (scripts/gen_synth.py):
#   OpenMathInstruct-2 (CC-BY-4.0, solutions written by Llama-3.1-405B) and GSM8K (MIT) -> synth_reasoning
#   SQuAD v2 and HotpotQA (CC-BY-SA-4.0) -> synth_grounded (copy-the-answer reading comprehension)
EXTRA = [("nvidia/OpenMathInstruct-2", f"data/train-{i:05d}-of-00032.parquet") for i in range(2)] + [
    ("openai/gsm8k", "main/train-00000-of-00001.parquet"),
    ("rajpurkar/squad_v2", "squad_v2/train-00000-of-00001.parquet"),
    ("hotpotqa/hotpot_qa", "distractor/train-00000-of-00002.parquet"),
]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--english", type=int, default=4)
    ap.add_argument("--math", type=int, default=14)

    a = ap.parse_args()
    for lang, shards in CODE_SHARDS.items():
        for i in shards:
            p = hf_hub_download(SOURCES["code"][0], f"stack-edu-{i:04d}.json.gz", repo_type="dataset",
                                local_dir=f"{DATA}/raw/code")
            print("code", lang, p, os.path.getsize(p) // 2**20, "MiB", flush=True)
    for repo, f in EXTRA:
        p = hf_hub_download(repo, f, repo_type="dataset", local_dir=f"{DATA}/raw/extra/{repo.split('/')[1]}")
        print("extra", p, os.path.getsize(p) // 2**20, "MiB", flush=True)
    for name, n in [("math", a.math), ("english", a.english)]:
        repo, files = SOURCES[name]
        for f in files[:n]:
            p = hf_hub_download(repo, f, repo_type="dataset", local_dir=f"{DATA}/raw/{name}")
            print(name, p, os.path.getsize(p) // 2**20, "MiB", flush=True)

if __name__ == "__main__":
    main()