#!/usr/bin/env python3
import argparse
import json
import os
import random
from pathlib import Path

from datasets import get_dataset_config_names, load_dataset


MMLU_DATASET = "cais/mmlu"
GPQA_DATASET = "Idavidrein/gpqa"
CEVAL_DATASET = "ceval/ceval-exam"
CMMLU_DATASET = "haonan-li/cmmlu"

MMLU_CONFIGS = [
    "abstract_algebra",
    "anatomy",
    "astronomy",
    "business_ethics",
    "clinical_knowledge",
    "college_biology",
    "college_chemistry",
    "college_computer_science",
    "college_mathematics",
    "college_physics",
    "computer_security",
    "conceptual_physics",
    "econometrics",
    "electrical_engineering",
    "elementary_mathematics",
    "formal_logic",
    "global_facts",
    "high_school_biology",
    "high_school_chemistry",
    "high_school_computer_science",
    "high_school_government_and_politics",
    "high_school_macroeconomics",
    "high_school_mathematics",
    "high_school_physics",
    "high_school_statistics",
    "international_law",
    "jurisprudence",
    "machine_learning",
    "management",
    "marketing",
    "medical_genetics",
    "moral_disputes",
    "nutrition",
    "philosophy",
    "professional_law",
    "professional_medicine",
    "professional_psychology",
    "public_relations",
    "security_studies",
    "sociology",
    "us_foreign_policy",
    "world_religions",
]

CEVAL_CONFIGS = [
    "computer_network",
    "operating_system",
    "computer_architecture",
    "college_programming",
    "college_physics",
    "college_chemistry",
    "advanced_mathematics",
    "probability_and_statistics",
    "discrete_mathematics",
    "electrical_engineer",
    "metrology_engineer",
    "high_school_mathematics",
    "high_school_physics",
    "high_school_chemistry",
    "high_school_biology",
    "legal_professional",
    "business_administration",
    "marxism",
    "mao_zedong_thought",
    "education_science",
    "teacher_qualification",
    "modern_chinese_history",
    "chinese_language_and_literature",
    "logic",
]

CMMLU_CONFIGS = [
    "agronomy",
    "anatomy",
    "ancient_chinese",
    "arts",
    "astronomy",
    "business_ethics",
    "chinese_civil_service_exam",
    "chinese_driving_rule",
    "chinese_food_culture",
    "chinese_foreign_policy",
    "chinese_history",
    "college_actuarial_science",
    "college_education",
    "college_engineering_hydrology",
    "college_law",
    "college_mathematics",
    "college_medical_statistics",
    "college_medicine",
    "computer_science",
    "conceptual_physics",
    "econometrics",
    "education",
    "electrical_engineering",
    "elementary_chinese",
    "elementary_commonsense",
    "elementary_information_and_technology",
    "elementary_mathematics",
    "ethnology",
    "food_science",
    "genetics",
    "global_facts",
    "high_school_biology",
    "high_school_chemistry",
    "high_school_geography",
    "high_school_mathematics",
    "high_school_physics",
    "human_sexuality",
    "international_law",
    "journalism",
    "jurisprudence",
    "legal_and_moral_basis",
    "logical",
    "machine_learning",
    "management",
    "marketing",
    "marxist_theory",
    "modern_chinese",
    "nutrition",
    "philosophy",
    "professional_accounting",
    "professional_law",
    "professional_medicine",
    "professional_psychology",
    "public_relations",
    "security_study",
    "sociology",
    "sports_science",
    "traditional_chinese_medicine",
    "virology",
    "world_history",
    "world_religions",
]


def norm(text):
    return " ".join(str(text or "").replace("\x00", " ").split()).strip()


def stable_sample(rng, rows, n):
    if len(rows) <= n:
        return list(rows)
    return rng.sample(rows, n)


def add_item(items, prefix, idx, source, subset, question, choices, answer_idx, metadata=None):
    question = norm(question)
    choices = [norm(x) for x in choices]
    if not question or len(choices) < 2:
        return
    if answer_idx < 0 or answer_idx >= len(choices):
        return
    answer_text = choices[answer_idx]
    if not answer_text:
        return
    mcq_prompt = (
        "Answer the following multiple-choice question. Choose the single best option.\n\n"
        f"Question:\n{question}\n\nAnswer:"
    )
    ppl_text = f"Question: {question}\nCorrect answer: {answer_text}"
    items.append(
        {
            "id": f"{prefix}_{idx:05d}",
            "category": prefix,
            "source": source,
            "subset": subset,
            "ppl_text": ppl_text,
            "mcq_prompt": mcq_prompt,
            "choices": choices,
            "answer_idx": int(answer_idx),
            "answer_text": answer_text,
            "metadata": metadata or {},
        }
    )


def build_mmlu(rng, quota):
    configs = list(MMLU_CONFIGS)
    rng.shuffle(configs)
    rows = []
    per_config = max(20, quota // 24)
    for cfg in configs:
        print(f"[info] load mmlu {cfg}", flush=True)
        try:
            ds = load_dataset(MMLU_DATASET, cfg, split="test")
        except Exception as exc:
            print(f"[warn] skip MMLU {cfg}: {type(exc).__name__}: {str(exc)[:160]}")
            continue
        sampled = stable_sample(rng, list(ds), min(per_config, len(ds)))
        for r in sampled:
            rows.append((cfg, r))
        if len(rows) >= quota * 2:
            break
    rng.shuffle(rows)
    items = []
    for i, (cfg, r) in enumerate(rows[:quota]):
        add_item(
            items,
            "mmlu",
            i,
            MMLU_DATASET,
            cfg,
            r.get("question"),
            r.get("choices", []),
            r.get("answer"),
            {"subject": r.get("subject", cfg)},
        )
    return items


def build_gpqa(rng, quota, config="gpqa_main"):
    ds = load_dataset(GPQA_DATASET, config, split="train", token=True)
    rows = stable_sample(rng, list(ds), min(quota, len(ds)))
    items = []
    for i, r in enumerate(rows):
        choices = [
            r.get("Correct Answer"),
            r.get("Incorrect Answer 1"),
            r.get("Incorrect Answer 2"),
            r.get("Incorrect Answer 3"),
        ]
        order = list(range(4))
        rng.shuffle(order)
        shuffled = [choices[j] for j in order]
        answer_idx = order.index(0)
        add_item(
            items,
            "gpqa",
            i,
            GPQA_DATASET,
            config,
            r.get("Question"),
            shuffled,
            answer_idx,
            {
                "high_level_domain": r.get("High-level domain"),
                "subdomain": r.get("Subdomain"),
                "record_id": r.get("Record ID"),
            },
        )
    return items


def letter_answer_idx(ans):
    if ans is None:
        return -1
    s = str(ans).strip()
    if s in {"0", "1", "2", "3"}:
        return int(s)
    if s:
        ch = s[0].upper()
        if ch in "ABCD":
            return ord(ch) - ord("A")
    return -1


def build_abcd_dataset(rng, dataset_name, prefix, quota):
    if prefix == "ceval":
        configs = list(CEVAL_CONFIGS)
    elif prefix == "cmmlu":
        configs = list(CMMLU_CONFIGS)
    else:
        configs = get_dataset_config_names(dataset_name)
    rng.shuffle(configs)
    rows = []
    per_config = max(24, quota // 18)
    for cfg in configs:
        print(f"[info] load {prefix} {cfg}", flush=True)
        split = "test"
        try:
            ds = load_dataset(dataset_name, cfg, split=split)
        except Exception:
            try:
                split = "val"
                ds = load_dataset(dataset_name, cfg, split=split)
            except Exception as exc:
                print(f"[warn] skip {prefix} {cfg}: {type(exc).__name__}: {str(exc)[:160]}")
                continue
        sampled = stable_sample(rng, list(ds), min(per_config, len(ds)))
        for r in sampled:
            rows.append((cfg, split, r))
        if len(rows) >= quota * 2:
            break
    rng.shuffle(rows)
    items = []
    for i, (cfg, split, r) in enumerate(rows[:quota]):
        choices = [r.get("A"), r.get("B"), r.get("C"), r.get("D")]
        add_item(
            items,
            prefix,
            i,
            dataset_name,
            cfg,
            r.get("question"),
            choices,
            letter_answer_idx(r.get("answer")),
            {"split": split},
        )
    return items


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--seed", type=int, default=20260607)
    ap.add_argument("--mmlu", type=int, default=700)
    ap.add_argument("--gpqa", type=int, default=400)
    ap.add_argument("--ceval", type=int, default=450)
    ap.add_argument("--cmmlu", type=int, default=450)
    args = ap.parse_args()

    rng = random.Random(args.seed)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    print("[info] HF_ENDPOINT", os.environ.get("HF_ENDPOINT", ""))
    builders = [
        ("mmlu", lambda: build_mmlu(rng, args.mmlu)),
        ("gpqa", lambda: build_gpqa(rng, args.gpqa)),
        ("ceval", lambda: build_abcd_dataset(rng, CEVAL_DATASET, "ceval", args.ceval)),
        ("cmmlu", lambda: build_abcd_dataset(rng, CMMLU_DATASET, "cmmlu", args.cmmlu)),
    ]
    all_items = []
    errors = {}
    for name, fn in builders:
        if getattr(args, name) <= 0:
            print(f"[info] skip {name}: quota=0")
            continue
        try:
            items = fn()
            print(f"[info] built {name}: {len(items)}")
            all_items.extend(items)
        except Exception as exc:
            errors[name] = f"{type(exc).__name__}: {str(exc)[:500]}"
            print(f"[error] {name}: {errors[name]}")

    rng.shuffle(all_items)
    bench_path = out_dir / "heldout_public_mcq_2k.jsonl"
    with bench_path.open("w", encoding="utf-8") as f:
        for item in all_items:
            f.write(json.dumps(item, ensure_ascii=False) + "\n")

    counts = {}
    for item in all_items:
        counts[item["category"]] = counts.get(item["category"], 0) + 1
    stats = {
        "seed": args.seed,
        "requested": {"mmlu": args.mmlu, "gpqa": args.gpqa, "ceval": args.ceval, "cmmlu": args.cmmlu},
        "total_items": len(all_items),
        "counts": counts,
        "errors": errors,
        "output": str(bench_path),
    }
    with (out_dir / "heldout_public_mcq_2k_stats.json").open("w", encoding="utf-8") as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)
    print(json.dumps(stats, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()