"""
MMLU data construction for Claim 5 (paper Section 5.2 / Appendix B.2).

Builds, deterministically (seed 0):
  * a 1,000-example Humanities fine-tuning set,
  * a 4x250 = 1,000-example evaluation set (Humanities / Social Sciences /
    STEM / Other) that is DISJOINT from the fine-tuning set,
  * 7 few-shot exemplars per subject drawn from the MMLU `validation` split
    (topped up from `dev` when a subject has fewer than 7 validation rows),
    so no exemplar can ever appear in the fine-tuning set or the eval set.

The paper fine-tunes Qwen2.5-3B-Instruct on "the Humanities subset of MMLU,
which contains 1,000 training examples" and evaluates on "the MMLU test set
across four categories ... consisting of 1,000 examples in total".  MMLU ships
no dedicated Humanities train split, so we carve the 1,000 fine-tuning rows out
of the Humanities *test* pool and hold the 250 evaluation rows out of it.
"""
import random

HUMANITIES = [
    "formal_logic", "high_school_european_history", "high_school_us_history",
    "high_school_world_history", "international_law", "jurisprudence",
    "logical_fallacies", "moral_disputes", "moral_scenarios", "philosophy",
    "prehistory", "professional_law", "world_religions",
]
STEM = [
    "abstract_algebra", "astronomy", "college_biology", "college_chemistry",
    "college_computer_science", "college_mathematics", "college_physics",
    "computer_security", "conceptual_physics", "electrical_engineering",
    "elementary_mathematics", "high_school_biology", "high_school_chemistry",
    "high_school_computer_science", "high_school_mathematics",
    "high_school_physics", "high_school_statistics", "machine_learning",
]
SOCIAL = [
    "econometrics", "high_school_geography",
    "high_school_government_and_politics", "high_school_macroeconomics",
    "high_school_microeconomics", "high_school_psychology", "human_sexuality",
    "professional_psychology", "public_relations", "security_studies",
    "sociology", "us_foreign_policy",
]
OTHER = [
    "anatomy", "business_ethics", "clinical_knowledge", "college_medicine",
    "global_facts", "human_aging", "management", "marketing",
    "medical_genetics", "miscellaneous", "nutrition", "professional_accounting",
    "professional_medicine", "virology",
]
CATEGORIES = {
    "Humanities": HUMANITIES,
    "Social Sciences": SOCIAL,
    "STEM": STEM,
    "Other": OTHER,
}
SUBJ2CAT = {s: c for c, ss in CATEGORIES.items() for s in ss}
LETTERS = ["A", "B", "C", "D"]

N_TRAIN = 1000
N_EVAL_PER_CAT = 250
N_SHOT = 7
DATA_SEED = 0


def _pretty(subject):
    return subject.replace("_", " ")


def question_text(row):
    """Zero-shot user turn for one MMLU row."""
    lines = [
        f"The following is a multiple choice question about {_pretty(row['subject'])}.",
        "",
        row["question"].strip(),
    ]
    for letter, choice in zip(LETTERS, row["choices"]):
        lines.append(f"{letter}. {str(choice).strip()}")
    lines.append("")
    lines.append("Answer with the letter of the correct option (A, B, C, or D).")
    return "\n".join(lines)


def answer_text(row):
    return LETTERS[int(row["answer"])]


def _rows(ds_split, subjects):
    out = []
    for r in ds_split:
        if r["subject"] in subjects:
            out.append({
                "subject": r["subject"],
                "question": r["question"],
                "choices": list(r["choices"]),
                "answer": int(r["answer"]),
            })
    return out


def build(cache_dir=None):
    from datasets import load_dataset
    dd = load_dataset("cais/mmlu", "all", cache_dir=cache_dir)
    test, val, dev = dd["test"], dd["validation"], dd["dev"]

    all_subjects = sorted(SUBJ2CAT)

    # ---- few-shot exemplars: validation split first, dev as top-up -------
    val_rows, dev_rows = {}, {}
    for r in _rows(val, set(all_subjects)):
        val_rows.setdefault(r["subject"], []).append(r)
    for r in _rows(dev, set(all_subjects)):
        dev_rows.setdefault(r["subject"], []).append(r)
    shots = {}
    for s in all_subjects:
        pool = list(val_rows.get(s, []))
        random.Random(DATA_SEED).shuffle(pool)
        picked = pool[:N_SHOT]
        if len(picked) < N_SHOT:
            picked = picked + dev_rows.get(s, [])[: N_SHOT - len(picked)]
        assert len(picked) == N_SHOT, (s, len(picked))
        shots[s] = picked

    # ---- eval sets (250 per category, from `test`) -----------------------
    eval_sets, used_hum = {}, set()
    for ci, (cat, subjects) in enumerate(CATEGORIES.items()):
        pool = _rows(test, set(subjects))
        idx = list(range(len(pool)))
        random.Random(DATA_SEED + 100 * (ci + 1)).shuffle(idx)
        chosen = idx[:N_EVAL_PER_CAT]
        eval_sets[cat] = [pool[i] for i in chosen]
        if cat == "Humanities":
            hum_pool = pool
            used_hum = set(chosen)

    # ---- 1,000 Humanities fine-tuning rows, disjoint from eval ----------
    remaining = [i for i in range(len(hum_pool)) if i not in used_hum]
    random.Random(DATA_SEED + 7).shuffle(remaining)
    train_rows = [hum_pool[i] for i in remaining[:N_TRAIN]]
    assert len(train_rows) == N_TRAIN

    return {"train": train_rows, "eval": eval_sets, "shots": shots}


def chat_messages(row, shots=None):
    """Multi-turn chat prompt.  `shots` = list of exemplar rows (same subject)."""
    msgs = []
    for ex in (shots or []):
        msgs.append({"role": "user", "content": question_text(ex)})
        msgs.append({"role": "assistant", "content": answer_text(ex)})
    msgs.append({"role": "user", "content": question_text(row)})
    return msgs