import json
import random
from pathlib import Path

from datasets import load_dataset

SEED = 42
rng = random.Random(SEED)

LETTERS = "ABCDEFGHIJ"

OUTPUT = Path("data/tasks.real-300.jsonl")


def answer_letter(value):
    if isinstance(value, int):
        return LETTERS[value]

    text = str(value).strip().upper()

    if text in LETTERS:
        return text

    if text.isdigit():
        return LETTERS[int(text)]

    raise ValueError(f"Unknown answer format: {value!r}")


def make_prompt(question, choices):
    valid_letters = LETTERS[:len(choices)]

    options = "\n".join(
        f"{LETTERS[i]}. {choice}"
        for i, choice in enumerate(choices)
    )

    return (
        f"{question.strip()}\n\n"
        f"{options}\n\n"
        f"Answer with only one letter ({', '.join(valid_letters)})."
    )


def domain_for(subject):
    s = subject.lower()

    if any(x in s for x in [
        "math",
        "algebra",
        "statistics",
        "physics",
    ]):
        return "math"

    if any(x in s for x in [
        "computer",
        "machine_learning",
        "security",
    ]):
        return "code"

    return "reasoning"


records = []


# --------------------------------------------------
# MMLU
# --------------------------------------------------

print("Loading MMLU...")

mmlu = list(
    load_dataset(
        "cais/mmlu",
        "all",
        split="validation",
    )
)

rng.shuffle(mmlu)

for i, row in enumerate(mmlu[:150]):
    subject = row.get("subject", "general")

    records.append({
        "task_id": f"mmlu-{i:04d}",
        "prompt": make_prompt(
            row["question"],
            row["choices"],
        ),
        "domain": domain_for(subject),
        "reference_answer": answer_letter(row["answer"]),
        "grader": {
            "type": "exact",
            "case_sensitive": False,
        },
        "split": "train",
        "tags": [
            "mmlu",
            subject,
        ],
        "metadata": {
            "source": "cais/mmlu",
            "subject": subject,
        },
    })


# --------------------------------------------------
# ARC Challenge
# --------------------------------------------------

print("Loading ARC-Challenge...")

arc = list(
    load_dataset(
        "allenai/ai2_arc",
        "ARC-Challenge",
        split="train",
    )
)

rng.shuffle(arc)

for i, row in enumerate(arc[:100]):
    labels = list(row["choices"]["label"])
    choices = list(row["choices"]["text"])

    answer_key = str(row["answerKey"]).strip()

    answer_index = labels.index(answer_key)

    records.append({
        "task_id": f"arc-{i:04d}",
        "prompt": make_prompt(
            row["question"],
            choices,
        ),
        "domain": "reasoning",
        "reference_answer": LETTERS[answer_index],
        "grader": {
            "type": "exact",
            "case_sensitive": False,
        },
        "split": "train",
        "tags": [
            "arc",
            "science",
            "reasoning",
        ],
        "metadata": {
            "source": "allenai/ai2_arc",
            "source_id": row["id"],
        },
    })


# --------------------------------------------------
# MMLU-Pro
# --------------------------------------------------

print("Loading MMLU-Pro...")

mmlu_pro = list(
    load_dataset(
        "TIGER-Lab/MMLU-Pro",
        split="validation",
    )
)

rng.shuffle(mmlu_pro)

for i, row in enumerate(mmlu_pro[:50]):
    category = row.get("category", "general")

    records.append({
        "task_id": f"mmlupro-{i:04d}",
        "prompt": make_prompt(
            row["question"],
            row["options"],
        ),
        "domain": domain_for(category),
        "reference_answer": answer_letter(row["answer"]),
        "grader": {
            "type": "exact",
            "case_sensitive": False,
        },
        "split": "train",
        "tags": [
            "mmlu-pro",
            category,
        ],
        "metadata": {
            "source": "TIGER-Lab/MMLU-Pro",
            "category": category,
            "source_id": row.get("question_id"),
        },
    })


# --------------------------------------------------
# Shuffle + train/validation/test
# --------------------------------------------------

rng.shuffle(records)

for i, record in enumerate(records):
    if i < 240:
        record["split"] = "train"
    elif i < 270:
        record["split"] = "validation"
    else:
        record["split"] = "test"


OUTPUT.parent.mkdir(parents=True, exist_ok=True)

with OUTPUT.open("w", encoding="utf-8") as handle:
    for record in records:
        handle.write(
            json.dumps(record, ensure_ascii=False) + "\n"
        )


print()
print(f"Created: {OUTPUT}")
print(f"Total: {len(records)}")
print("Train: 240")
print("Validation: 30")
print("Test: 30")