"""Download toloka/u-math from HuggingFace and convert to the agentic-rl
training format.

U-MATH has only a test split (1100 problems, 900 text-only).
We filter out multimodal problems and split 700/200 for train/test.

Target format (one JSON object per line):
{
  "messages": [{"content": "<problem>", "role": "user"}],
  "rollout_extra": {
    "solution": "<golden_answer>",
    "dataset": "umath",
    "subject": "<subject>"
  }
}
"""

import json
import random
from pathlib import Path

from datasets import load_dataset

OUTPUT_DIR = Path(__file__).resolve().parent.parent / "data"
SEED = 42
DATASET_REVISION = "7210f97b3f21122c3a90126935e966ca7d8951b6"


def convert_example(example: dict) -> dict:
    return {
        "messages": [
            {"content": example["problem_statement"], "role": "user"}
        ],
        "rollout_extra": {
            "solution": example["golden_answer"],
            "dataset": "umath",
            "subject": example["subject"],
        },
    }


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print("Loading toloka/u-math ...")
    ds = load_dataset("toloka/u-math", split="test", revision=DATASET_REVISION)

    text_only = [ex for ex in ds if not ex["has_image"]]
    print(f"Total: {len(ds)}, text-only: {len(text_only)}")

    random.seed(SEED)
    random.shuffle(text_only)

    split_idx = 700
    train_raw = text_only[:split_idx]
    test_raw = text_only[split_idx:]

    train_items = [convert_example(ex) for ex in train_raw]
    test_items = [convert_example(ex) for ex in test_raw]

    train_path = OUTPUT_DIR / "umath_train.jsonl"
    test_path = OUTPUT_DIR / "umath_test.jsonl"

    for path, items in [(train_path, train_items), (test_path, test_items)]:
        with open(path, "w", encoding="utf-8") as f:
            for item in items:
                f.write(json.dumps(item, ensure_ascii=False) + "\n")

    from collections import Counter
    subjects = Counter(ex["subject"] for ex in text_only)
    print(f"\nDone!")
    print(f"  Train: {len(train_items)} examples -> {train_path}")
    print(f"  Test:  {len(test_items)} examples -> {test_path}")
    print(f"  Subjects: {dict(subjects)}")


if __name__ == "__main__":
    main()
