"""Generate the unified QA dataset consumed by GraphRouter experiments."""

from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Dict, Iterable, Iterator, Mapping, Optional

import pandas as pd
import wandb
import yaml
from tqdm.auto import tqdm

from ember.api.data import stream
from ember.api.datasets.catalog import get_entry
from ember.api.exceptions import DatasetNotFoundError

Example = Dict[str, str]
Loader = Callable[[int], Iterator[Example]]

# Base directory for data files (embergraphrouter/)
_MODULE_DIR = Path(__file__).resolve().parent.parent


@dataclass(frozen=True)
class TaskSpec:
    """Contract describing how a task contributes rows to the unified table."""

    task_id: str
    metric: str
    description: str
    loader: Loader


def _example(query: str, ground_truth: str) -> Example:
    """Return a normalized question/answer pair."""

    return {"query": str(query), "ground_truth": str(ground_truth)}


def _load_json_rows(
    path: str,
    sample_size: int,
    *,
    query_fields: Iterable[str],
    ground_truth_field: str,
) -> Iterator[Example]:
    """Yield JSON rows while respecting ``sample_size``."""
    resolved_path = _MODULE_DIR / path
    data = json.loads(resolved_path.read_text(encoding="utf-8"))
    produced = 0
    for index, item in enumerate(data):
        if index >= sample_size:
            break
        query = "".join(item[field] for field in query_fields)
        yield _example(query, item[ground_truth_field])
        produced += 1

    if produced < sample_size:
        raise ValueError(f"File '{path}' provides {produced} examples; expected {sample_size}.")


def _load_parquet_rows(
    path: str,
    sample_size: int,
    *,
    query_field: str,
    ground_truth_field: str,
    ground_truth_subfield: str | None = None,
    ground_truth_index: int = 0,
) -> Iterator[Example]:
    """Yield parquet rows with optional nested-answer selection."""
    resolved_path = _MODULE_DIR / path
    frame = pd.read_parquet(resolved_path)
    produced = 0
    for _, row in zip(range(sample_size), frame.itertuples(index=False), strict=False):
        query = getattr(row, query_field)
        answer = getattr(row, ground_truth_field)
        if ground_truth_subfield is not None:
            answer = answer[ground_truth_subfield][ground_truth_index]
        yield _example(query, answer)
        produced += 1

    if produced < sample_size:
        raise ValueError(f"File '{path}' provides {produced} examples; expected {sample_size}.")


def _load_squad(sample_size: int) -> Iterator[Example]:
    """Load SQuAD with context passages included in queries.

    SQuAD is a reading comprehension task: models receive a passage and
    answer a question about it. The context is prepended to the question
    so that LLMs perform reading comprehension, not closed-book trivia.
    """
    resolved_path = _MODULE_DIR / "data" / "SQUAD" / "SQUAD.parquet"
    frame = pd.read_parquet(resolved_path)
    produced = 0
    for _, row in zip(range(sample_size), frame.itertuples(index=False), strict=False):
        query = (
            f"Read the following passage and answer the question.\n\n"
            f"Passage: {row.context}\n\n"
            f"Question: {row.question}"
        )
        answer = row.answers["text"][0]
        yield _example(query, answer)
        produced += 1

    if produced < sample_size:
        raise ValueError(
            f"SQUAD provides {produced} examples; expected {sample_size}."
        )


def _parse_officeqa_source_files(source_files: str) -> list[str]:
    """Split OfficeQA source_files field (delimited by \\r\\n)."""
    raw = source_files.replace("\r\n", "\n").replace("\r", "\n")
    return [
        name.strip()
        for segment in raw.split("\n")
        for name in segment.split(",")
        if name.strip()
    ]


def _load_officeqa(sample_size: int) -> Iterator[Example]:
    """Load OfficeQA questions with oracle document context.

    OfficeQA: 246 questions requiring grounded reasoning over U.S.
    Treasury Bulletin documents.  Each question includes a
    ``source_files`` field listing the relevant documents for oracle
    retrieval.  Answers are primarily numerical.
    """
    data_dir = _MODULE_DIR / "data" / "OfficeQA"
    csv_path = data_dir / "officeqa.csv"
    transformed_dir = data_dir / "transformed"

    if not csv_path.exists():
        raise FileNotFoundError(
            f"OfficeQA data not found at {csv_path}. "
            "Copy officeqa.csv and transformed/ from the OfficeQA corpus."
        )

    frame = pd.read_csv(csv_path)
    produced = 0
    for _, row in zip(range(sample_size), frame.itertuples(index=False), strict=False):
        # Load oracle source documents.
        filenames = _parse_officeqa_source_files(row.source_files)
        doc_parts: list[str] = []
        for name in filenames:
            path = transformed_dir / name
            if path.exists():
                doc_parts.append(f"--- {name} ---\n{path.read_text(encoding='utf-8')}")
        documents = "\n\n".join(doc_parts)

        query = (
            f"Answer the following question using the provided Treasury "
            f"Bulletin documents.\n\n"
            f"DOCUMENTS:\n{documents}\n\n"
            f"QUESTION: {row.question}\n\n"
            f"Provide a precise answer."
        )
        yield _example(query, str(row.answer))
        produced += 1

    if produced < sample_size:
        raise ValueError(
            f"OfficeQA provides {produced} examples; expected {sample_size}."
        )


def _load_livecodebench(sample_size: int) -> Iterator[Example]:
    """Stream LiveCodeBench rows."""

    for item in stream("livecodebench.execution").as_dicts().limit(sample_size):
        yield _example(item["question"], item["answer"])


def _load_mbpp(sample_size: int) -> Iterator[Example]:
    """Stream MBPP (Mostly Basic Python Problems) rows."""
    emitted = 0
    for item in stream("mbpp").as_dicts():
        yield _example(item["question"], item["answer"])
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"mbpp returned {emitted} examples; expected {sample_size}.")


def _load_agieval(sample_size: int) -> Iterator[Example]:
    """Stream AGIEval rows with formatted multiple-choice options.

    Note: Ember's agieval dataset returns `answer` as the answer text (e.g., "12.5"),
    not the letter key (e.g., "C"). We extract the letter key by matching
    the answer text against the choices dictionary.
    """
    emitted = 0
    for item in stream("agieval").as_dicts():
        question = str(item.get("question") or "")
        choices = item.get("choices") or {}
        answer_text = str(item.get("answer") or "").strip()

        # Extract the letter key for the answer text
        answer_key = None
        if isinstance(choices, dict):
            for key, value in choices.items():
                if str(value).strip() == answer_text:
                    answer_key = key
                    break

        if isinstance(choices, dict) and choices:
            formatted_query = f"{question}\n\nChoices:\n" + "\n".join(
                f"{lab}. {txt}" for lab, txt in sorted(choices.items())
            )
        else:
            formatted_query = question

        # Use letter key for em_mc evaluation; fall back to text if no match
        answer = answer_key if answer_key else answer_text
        yield _example(formatted_query, answer)
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"agieval returned {emitted} examples; expected {sample_size}.")


def _load_cruxeval_output(sample_size: int) -> Iterator[Example]:
    """Stream CRUXEval output prediction rows."""
    emitted = 0
    for item in stream("cruxeval.output").as_dicts():
        yield _example(item["question"], item["answer"])
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"cruxeval.output returned {emitted} examples; expected {sample_size}.")


def _load_cruxeval_x(sample_size: int) -> Iterator[Example]:
    """Stream CRUXEval-X multilingual output prediction rows."""
    emitted = 0
    for item in stream("cruxeval_x").as_dicts():
        yield _example(item["question"], item["answer"])
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"cruxeval_x returned {emitted} examples; expected {sample_size}.")


def _load_codexglue_defect(sample_size: int) -> Iterator[Example]:
    """Stream CodeXGLUE defect detection rows."""
    emitted = 0
    for item in stream("codexglue_defect").as_dicts():
        yield _example(item["question"], item["answer"])
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(
            f"codexglue_defect returned {emitted} examples; expected {sample_size}."
        )


def _format_choices(choices: Mapping[str, object]) -> str:
    """Render multiple-choice options into prompt-friendly prose."""

    ordered = sorted(choices.items())
    return "\n".join(f"{label}. {value}" for label, value in ordered)


def _mmlu_subject_loader(subject: str) -> Loader:
    """Return a loader for a specific MMLU-Pro subject."""
    # MMLU-Pro uses spaces in categories ("computer science"), but task IDs
    # use underscores ("computer_science").  Normalize for comparison.
    category = subject.replace("_", " ")

    def loader(sample_size: int) -> Iterator[Example]:
        emitted = 0
        for item in stream("mmlu_pro").as_dicts():
            if item["metadata"]["category"] == category:
                question = f"{item['question']}\n\nChoices:\n{_format_choices(item['choices'])}"
                yield _example(question, item["answer"])
                emitted += 1
                if emitted >= sample_size:
                    break

        if emitted < sample_size:
            raise ValueError(
                f"Subject '{category}' in mmlu_pro returned {emitted} examples; expected {sample_size}."
            )

    return loader


def _bbeh_subject_loader(subject: str) -> Loader:
    """Return a loader for a specific BBEH subject."""

    def loader(sample_size: int) -> Iterator[Example]:
        emitted = 0
        for item in stream("bbeh").as_dicts():
            if item["metadata"]["task"] == subject:
                yield _example(item["question"], item["answer"])
                emitted += 1
                if emitted >= sample_size:
                    break

        if emitted < sample_size:
            raise ValueError(
                f"Subject '{subject}' in bbeh returned {emitted} examples; expected {sample_size}."
            )

    return loader


def _load_mmlu(sample_size: int) -> Iterator[Example]:
    """Stream MMLU rows with formatted multiple-choice options."""
    emitted = 0
    for item in stream("mmlu").as_dicts():
        question = str(item.get("question") or "")
        raw_choices = item.get("choices") or item.get("options") or ()

        if isinstance(raw_choices, Mapping):
            labels = sorted(raw_choices.keys())
            options = [str(raw_choices[lab]) for lab in labels]
        else:
            options = [str(opt) for opt in raw_choices]
            labels = [chr(ord("A") + i) for i in range(len(options))]

        answer_raw = item.get("answer")
        if isinstance(answer_raw, int) and 0 <= answer_raw < len(labels):
            answer = labels[answer_raw]
        else:
            answer = str(answer_raw or "").strip().upper()

        choices_dict = dict(zip(labels, options, strict=False))
        formatted_query = f"{question}\n\nChoices:\n{_format_choices(choices_dict)}"
        yield _example(formatted_query, answer)
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"mmlu returned {emitted} examples; expected {sample_size}.")


def _load_arc(sample_size: int) -> Iterator[Example]:
    """Stream ARC (AI2 Reasoning Challenge) rows."""
    emitted = 0
    for item in stream("arc").as_dicts():
        question = str(item.get("question") or "")
        choices = item.get("choices") or {}
        raw_labels = [str(x) for x in (choices.get("label") or [])]
        texts = [str(x) for x in (choices.get("text") or [])]

        if not raw_labels or not texts:
            continue

        # Normalize labels to letters (some ARC items use 1,2,3,4 instead of A,B,C,D)
        letter_labels = [chr(ord("A") + i) for i in range(len(raw_labels))]
        label_map = dict(zip(raw_labels, letter_labels, strict=False))

        # answerKey is in metadata for Ember's normalized schema
        meta = item.get("metadata") or {}
        raw_answer = str(meta.get("answerKey") or item.get("answerKey") or "").strip()
        # Map numeric answer to letter if needed
        answer = label_map.get(raw_answer, raw_answer.upper())

        formatted_query = f"{question}\n\nChoices:\n" + "\n".join(
            f"{lab}. {txt}" for lab, txt in zip(letter_labels, texts, strict=False)
        )
        yield _example(formatted_query, answer)
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"arc returned {emitted} examples; expected {sample_size}.")


def _load_hellaswag(sample_size: int) -> Iterator[Example]:
    """Stream HellaSwag commonsense inference rows."""
    emitted = 0
    for item in stream("hellaswag", split="validation").as_dicts():
        meta = item.get("metadata") or {}
        data = {**meta, **{k: v for k, v in item.items() if v is not None}}

        context_parts = []
        for key in ("ctx", "ctx_a", "activity_label"):
            val = data.get(key)
            if val and str(val).strip():
                context_parts.append(str(val).strip())
        context = " ".join(context_parts)

        raw_endings = data.get("endings") or data.get("choices") or ()
        if isinstance(raw_endings, Mapping):
            endings = list(raw_endings.get("text") or [])
        else:
            endings = [str(e) for e in raw_endings]

        if not context or not endings:
            continue

        labels = [chr(ord("A") + i) for i in range(len(endings))]
        # Ember schema uses 'answer' at root, which contains the index as string
        label_raw = item.get("answer") or data.get("label")
        if isinstance(label_raw, int) and 0 <= label_raw < len(labels):
            answer = labels[label_raw]
        elif isinstance(label_raw, str) and label_raw.strip().isdigit():
            idx = int(label_raw.strip())
            answer = labels[idx] if 0 <= idx < len(labels) else ""
        else:
            answer = str(label_raw or "").strip().upper()

        formatted_query = f"{context}\n\nChoices:\n" + "\n".join(
            f"{lab}. {end}" for lab, end in zip(labels, endings, strict=False)
        )
        yield _example(formatted_query, answer)
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"hellaswag returned {emitted} examples; expected {sample_size}.")


def _load_truthfulqa(sample_size: int) -> Iterator[Example]:
    """Stream TruthfulQA multiple-choice rows."""
    emitted = 0
    for item in stream("truthfulqa").as_dicts():
        meta = item.get("metadata") or {}
        data = {**meta, **{k: v for k, v in item.items() if v is not None}}

        question = str(data.get("question") or "")
        mc1 = data.get("mc1_targets") or {}
        mc_choices = mc1.get("choices") or []
        mc_labels = mc1.get("labels") or []

        if not question or not mc_choices:
            continue

        labels = [chr(ord("A") + i) for i in range(len(mc_choices))]
        correct_idx = next((idx for idx, label_val in enumerate(mc_labels) if label_val == 1), None)
        answer = labels[correct_idx] if correct_idx is not None else ""

        formatted_query = f"{question}\n\nChoices:\n" + "\n".join(
            f"{lab}. {opt}" for lab, opt in zip(labels, mc_choices, strict=False)
        )
        yield _example(formatted_query, answer)
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"truthfulqa returned {emitted} examples; expected {sample_size}.")


def _load_winogrande(sample_size: int) -> Iterator[Example]:
    """Stream Winogrande pronoun resolution rows."""
    emitted = 0
    for item in stream("winogrande").as_dicts():
        meta = item.get("metadata") or {}
        data = {**meta, **{k: v for k, v in item.items() if v}}

        sentence = str(data.get("sentence") or "")
        option1 = str(data.get("option1") or "")
        option2 = str(data.get("option2") or "")
        answer_raw = data.get("answer")

        if not sentence or not option1 or not option2:
            continue

        if answer_raw in {"1", 1}:
            answer = "A"
        elif answer_raw in {"2", 2}:
            answer = "B"
        else:
            answer = str(answer_raw or "").strip().upper()

        formatted_query = (
            f"{sentence.replace('_', '_____')}\n\nChoices:\nA. {option1}\nB. {option2}"
        )
        yield _example(formatted_query, answer)
        emitted += 1
        if emitted >= sample_size:
            break

    if emitted < sample_size:
        raise ValueError(f"winogrande returned {emitted} examples; expected {sample_size}.")


def _catalog_description(catalog_id: str, fallback: str) -> str:
    """Fetch the canonical dataset description when available."""

    try:
        entry = get_entry(catalog_id)
        if entry.description:
            return entry.description
    except DatasetNotFoundError:
        pass
    return fallback


# BBEH subjects with (name, description, metric).
# All BBEH tasks use the unified 'bbeh' metric per official evaluation:
# https://github.com/google-deepmind/bbeh
# The bbeh metric extracts answers and uses fuzzy matching to handle:
# - Format variations: "(A)" vs "a", "proved" vs "Proved"
# - Numeric equivalence: "42" vs "42.0"
# - Answer prefixes: "The final answer is: X" -> "X"
_BBEH_SUBJECTS = (
    (
        "boardgame qa",
        "BBEH boardgame question answering that requires tracking board state and rules.",
        "bbeh",  # Answers: "proved", "disproved", "unknown"
    ),
    (
        "causal understanding",
        "BBEH causal understanding tasks focused on cause-and-effect reasoning.",
        "bbeh",  # Answers: "Yes", "No", "Ambiguous"
    ),
    (
        "dyck languages",
        "BBEH Dyck language prompts stressing stack-based formal grammar reasoning.",
        "bbeh",  # Mixed: numbers and "No"
    ),
    (
        "linguini",
        "BBEH linguini logic games blending language and abstraction.",
        "bbeh",  # Foreign language text
    ),
    (
        "boolean expressions",
        "BBEH logical expression evaluation tasks that probe symbolic reasoning.",
        "bbeh",  # Answers: "(A)", "(B)", "(C)", "(D)", "(E)"
    ),
    (
        "disambiguation qa",
        "BBEH disambiguation questions challenging pragmatic interpretation.",
        "bbeh",  # Answers: "(A)", "(B)", etc.
    ),
    (
        "geometric shapes",
        "BBEH geometric shape tasks assessing spatial reasoning skills.",
        "bbeh",  # Answers: "(A)" through "(J)"
    ),
    (
        "hyperbaton",
        "BBEH hyperbaton linguistic puzzles targeting syntactic flexibility.",
        "bbeh",  # Answers: "A", "B", etc. (no parentheses)
    ),
    (
        "movie recommendation",
        "BBEH movie recommendation scenarios demanding preference reasoning.",
        "bbeh",  # Answers: "(A)" through "(J)"
    ),
    (
        "buggy tables",
        "BBEH buggy tables reconstruction where models must repair malformed tables.",
        "bbeh",  # Answers: decimal numbers like "14.47", "4666.0"
    ),
    (
        "multistep arithmetic",
        "BBEH multistep arithmetic problems emphasizing chained calculation accuracy.",
        "bbeh",  # Answers: integers like "8359", "-18311"
    ),
    (
        "nycc",
        "BBEH New Yorker Cartoon Caption contest requiring humor and pragmatic inference.",
        "bbeh",  # Answers: "(A)", "(B)", etc.
    ),
    (
        "object counting",
        "BBEH object counting tasks demanding precise enumeration from complex descriptions.",
        "bbeh",  # Answers: integers like "131", "1486"
    ),
    (
        "object properties",
        "BBEH object properties reasoning about attributes in structured scenarios.",
        "bbeh",  # Answers: integers like "43", "15"
    ),
    (
        "sarc triples",
        "BBEH sarcasm detection via triple classification of statement intent.",
        "bbeh",  # Answers: comma-separated triples like "1,1,0"
    ),
    (
        "shuffled objects",
        "BBEH tracking shuffled objects through sequences of swaps and moves.",
        "bbeh",  # Answers: "(A)", "(B)", etc.
    ),
    (
        "spatial reasoning",
        "BBEH spatial reasoning about object positions and relationships.",
        "bbeh",  # Answers: free text (object names, lists)
    ),
    (
        "sportqa",
        "BBEH sports understanding requiring domain knowledge and rule reasoning.",
        "bbeh",  # Answers: comma-separated letters like "A, A, A"
    ),
    (
        "temporal sequence",
        "BBEH temporal sequence ordering events and computing durations.",
        "bbeh",  # Answers: comma-separated numbers like "150, 2"
    ),
    (
        "time arithmetic",
        "BBEH time arithmetic computing schedules and detecting unsolvable constraints.",
        "bbeh",  # Answers: bracketed lists or "unanswerable"
    ),
    (
        "web of lies",
        "BBEH web of lies truth-value propagation through chains of assertions.",
        "bbeh",  # Answers: comma-separated yes/no/unknown
    ),
    (
        "word sorting",
        "BBEH word sorting requiring alphabetical ordering of token sequences.",
        "bbeh",  # Answers: integers or sorted word lists
    ),
    (
        "zebra puzzles",
        "BBEH zebra logic puzzles requiring constraint satisfaction and deduction.",
        "bbeh",  # Answers: integers like "6", "5"
    ),
)

_MMLU_SUBJECTS = (
    ("biology", "MMLU-Pro biology discipline grounded in advanced life sciences."),
    ("business", "MMLU-Pro business discipline covering management, finance, and operations."),
    (
        "chemistry",
        "MMLU-Pro chemistry discipline requiring mastery of molecular and reaction theory.",
    ),
    ("history", "MMLU-Pro history discipline spanning world events and historiography."),
    ("law", "MMLU-Pro law discipline focused on legal doctrine and reasoning."),
    (
        "other",
        "MMLU-Pro miscellaneous discipline aggregating specialized subjects beyond the core tracks.",
    ),
    ("psychology", "MMLU-Pro psychology discipline emphasizing cognitive and behavioral science."),
    (
        "computer_science",
        "MMLU-Pro computer science discipline covering algorithms, data structures, "
        "and computational theory.",
    ),
)

# Tasks disabled due to evaluation limitations.
# - livecodebench: Requires code execution in sandbox, not suitable for text-based eval
# - mbpp: Code generation requires execution-based testing, exact match fails
DISABLED_TASKS: frozenset[str] = frozenset({
    "livecodebench",
    "mbpp",
})


TASK_SPECS = (
    TaskSpec(
        task_id="alpaca_data",
        metric="f1_score",
        description=(
            "The Alpaca dataset is designed for instruction-following tasks, where the model is required to "
            "generate coherent and contextually appropriate responses to given instructions or prompts. It focuses "
            "on understanding diverse user requests and providing informative and accurate outputs based on those "
            "instructions."
        ),
        loader=lambda sample_size: _load_json_rows(
            "data/alpaca_data/alpaca_data.json",
            sample_size,
            query_fields=("instruction", "input"),
            ground_truth_field="output",
        ),
    ),
    TaskSpec(
        task_id="GSM8K",
        metric="GSM8K",
        description=(
            "The GSM8K dataset is tailored for mathematical problem-solving tasks. It consists of natural language math "
            "problems that require the model to comprehend the problem statement, apply the correct mathematical "
            "operations, and provide the solution. The primary challenge lies in both parsing complex language and "
            "performing accurate calculations."
        ),
        loader=lambda sample_size: _load_json_rows(
            "data/GSM8K/GSM8K.json",
            sample_size,
            query_fields=("instruction", "input"),
            ground_truth_field="answer",
        ),
    ),
    TaskSpec(
        task_id="mmlu",
        metric="em_mc",
        description=(
            "The MMLU (Massive Multitask Language Understanding) dataset evaluates broad factual knowledge "
            "across 57 subjects spanning STEM, humanities, social sciences, and more. Models must select "
            "the correct answer from multiple choices, testing both domain expertise and reasoning."
        ),
        loader=_load_mmlu,
    ),
    TaskSpec(
        task_id="arc",
        metric="em_mc",
        description=(
            "The ARC (AI2 Reasoning Challenge) dataset contains science exam questions requiring complex "
            "reasoning. The challenge set focuses on questions that simple retrieval or word co-occurrence "
            "methods cannot answer, demanding genuine scientific understanding."
        ),
        loader=_load_arc,
    ),
    TaskSpec(
        task_id="hellaswag",
        metric="em_mc",
        description=(
            "HellaSwag tests commonsense natural language inference by asking models to select the most "
            "plausible continuation of a scenario. The adversarially-filtered distractors require nuanced "
            "understanding of real-world situations and physical dynamics."
        ),
        loader=_load_hellaswag,
    ),
    TaskSpec(
        task_id="truthfulqa",
        metric="em_mc",
        description=(
            "TruthfulQA measures whether models generate truthful answers to questions where humans often "
            "hold misconceptions. It tests resistance to common falsehoods, conspiracy theories, and "
            "superstitions that appear frequently in training data."
        ),
        loader=_load_truthfulqa,
    ),
    TaskSpec(
        task_id="winogrande",
        metric="em_mc",
        description=(
            "Winogrande is a large-scale dataset for commonsense reasoning via pronoun resolution. Models "
            "must determine which entity a pronoun refers to based on contextual cues, requiring "
            "understanding of real-world knowledge and situational logic."
        ),
        loader=_load_winogrande,
    ),
    TaskSpec(
        task_id="multi_news",
        metric="f1_score",
        description=(
            "The Multi-News dataset is aimed at text summarization tasks. It contains multiple news articles on the "
            "same topic, and the model's objective is to generate a concise and comprehensive summary that integrates "
            "information from all the articles. The challenge is to distill key points while maintaining coherence and "
            "avoiding redundancy."
        ),
        loader=lambda sample_size: _load_json_rows(
            "data/multi_news/multi_news.json",
            sample_size,
            query_fields=("instruction", "input"),
            ground_truth_field="output",
        ),
    ),
    TaskSpec(
        task_id="SQUAD",
        metric="squad_em",
        description=(
            "The SQuAD dataset is focused on reading comprehension, where the model is given "
            "a passage of text and must extract a precise answer to a question based on the "
            "content of that passage. Tests comprehension, information retrieval, and concise "
            "answer generation."
        ),
        loader=_load_squad,
    ),
    TaskSpec(
        task_id="officeqa",
        metric="officeqa_numeric",
        description=(
            "OfficeQA: Grounded reasoning over U.S. Treasury Bulletin documents. "
            "Questions require extracting precise numerical data from dense "
            "financial tables, cross-referencing across documents, and performing "
            "analytical calculations. Answers are typically numbers, dates, or "
            "short financial terms."
        ),
        loader=_load_officeqa,
    ),
    TaskSpec(
        task_id="livecodebench",
        metric="em",
        description=_catalog_description(
            "livecodebench",
            "LiveCodeBench evaluates code generation by asking models to predict program outputs under provided "
            "inputs, emphasizing precise execution reasoning and determinism.",
        ),
        loader=_load_livecodebench,
    ),
    TaskSpec(
        task_id="mbpp",
        metric="em",
        description=(
            "MBPP (Mostly Basic Python Problems) contains 500 crowd-sourced Python programming problems "
            "designed to be solvable by entry-level programmers. Each problem includes a task description, "
            "code solution, and automated test cases for verification."
        ),
        loader=_load_mbpp,
    ),
    TaskSpec(
        task_id="agieval",
        metric="em_mc",
        description=(
            "AGIEval is a human-centric benchmark derived from standardized exams like SAT, LSAT, GRE, and "
            "professional certifications. It tests advanced reasoning, reading comprehension, and domain "
            "knowledge across law, math, and logic."
        ),
        loader=_load_agieval,
    ),
    TaskSpec(
        task_id="cruxeval_output",
        metric="em",
        description=(
            "CRUXEval output prediction: given a Python function and input, predict the exact "
            "output. Tests code reasoning and execution tracing without running code."
        ),
        loader=_load_cruxeval_output,
    ),
    # NOTE: cruxeval_x removed -- HF dataset gated/private (401).
    # Re-add when CRUXEVAL-X/cruxeval-x is publicly accessible.
    TaskSpec(
        task_id="codexglue_defect",
        metric="em",
        description=(
            "CodeXGLUE defect detection: classify C functions as vulnerable (1) or secure (0). "
            "Tests security reasoning and code comprehension."
        ),
        loader=_load_codexglue_defect,
    ),
    *(
        TaskSpec(
            task_id=f"bbeh_{subject.replace(' ', '_')}",
            metric=metric,
            description=description,
            loader=_bbeh_subject_loader(subject),
        )
        for subject, description, metric in _BBEH_SUBJECTS
    ),
    *(
        TaskSpec(
            task_id=f"mmlu_pro_{subject}",
            metric="em_mc",
            description=description,
            loader=_mmlu_subject_loader(subject),
        )
        for subject, description in _MMLU_SUBJECTS
    ),
)

# Per-task sample size multipliers for difficulty-weighted sampling.
# Hard tasks get more samples to ensure the router learns when to escalate.
# Format: prefix -> multiplier (applied to base sample_size)
TASK_DIFFICULTY_MULTIPLIERS: dict[str, float] = {
    # HARD: Complex reasoning, code, math - where routing matters most
    "bbeh_": 3.0,              # BBEH reasoning tasks
    "mmlu_pro_chemistry": 3.0,  # Hardest MMLU-Pro subject
    "mmlu_pro_law": 2.5,        # Complex legal reasoning
    "mmlu_pro_biology": 2.0,    # Scientific reasoning
    "mmlu_pro_business": 2.0,   # Business reasoning
    "mmlu_pro_": 1.5,           # Other MMLU-Pro (fallback)
    "mbpp": 3.0,                # Code generation
    "cruxeval_output": 3.0,     # Code execution prediction
    "codexglue_defect": 2.5,    # Security defect detection
    "mmlu_pro_computer_science": 3.0,  # CS reasoning
    "agieval": 3.0,             # Advanced reasoning (SAT/LSAT/GRE)
    "livecodebench": 2.5,       # Code execution prediction
    "GSM8K": 2.0,               # Math word problems
    # MEDIUM: Moderate difficulty
    "arc": 1.5,                 # Science reasoning
    "truthfulqa": 1.5,          # Truthfulness (tricky)
    "mmlu": 1.0,                # Standard MMLU
    # EASY: High accuracy across models - less routing value
    "hellaswag": 0.5,           # Near-ceiling performance
    "winogrande": 0.5,          # Near-ceiling performance
    "alpaca_data": 0.7,         # Instruction following
    "multi_news": 0.5,          # Summarization
    "SQUAD": 0.5,               # Reading comprehension
    "officeqa": 0.25,           # 246 total questions (250 at base_size=1000)
}


def _get_task_sample_size(task_id: str, base_size: int) -> int:
    """Return the sample size for a task based on difficulty multipliers."""
    # Check exact match first, then prefix matches (longest prefix wins)
    if task_id in TASK_DIFFICULTY_MULTIPLIERS:
        return int(base_size * TASK_DIFFICULTY_MULTIPLIERS[task_id])

    # Find longest matching prefix
    best_match = ""
    best_mult = 1.0
    for prefix, mult in TASK_DIFFICULTY_MULTIPLIERS.items():
        if task_id.startswith(prefix) and len(prefix) > len(best_match):
            best_match = prefix
            best_mult = mult

    return int(base_size * best_mult)


def generate_unified_qa_dataset(
    output_path: str = "data/unified_qa_data.csv",
    sample_size: int = 60,
    *,
    show_progress: bool = True,
    wandb_run: Optional["wandb.sdk.wandb_run.Run"] = None,
) -> pd.DataFrame:
    """Generate the unified dataset and persist it to ``output_path``.

    Args:
        output_path: Destination CSV path.
        sample_size: Base number of examples per task (scaled by difficulty).
        show_progress: Whether to render tqdm progress bars.
        wandb_run: Active W&B run for metric logging.

    Returns:
        pandas.DataFrame: Materialized dataset.
    """

    if sample_size <= 0:
        raise ValueError("sample_size must be positive")

    rows: list[dict[str, str]] = []
    task_stats: list[dict[str, object]] = []
    pipeline_start = time.perf_counter()

    # Filter out disabled tasks
    active_specs = [spec for spec in TASK_SPECS if spec.task_id not in DISABLED_TASKS]

    task_bar = tqdm(
        total=len(active_specs),
        desc="Tasks",
        unit="task",
        disable=not show_progress,
    )

    for spec in active_specs:
        task_start = time.perf_counter()
        task_sample_size = _get_task_sample_size(spec.task_id, sample_size)
        example_bar = tqdm(
            total=task_sample_size,
            desc=f"{spec.task_id} (n={task_sample_size})",
            unit="ex",
            leave=False,
            disable=not show_progress,
        )

        task_rows: list[dict[str, str]] = []
        produced = 0
        error_message: Optional[str] = None

        try:
            for example in spec.loader(task_sample_size):
                task_rows.append(
                    {
                        "task_id": spec.task_id,
                        "query": example["query"],
                        "ground_truth": example["ground_truth"],
                        "metric": spec.metric,
                        "task_description": spec.description,
                    }
                )
                produced += 1
                example_bar.update(1)
        except ValueError as exc:
            error_message = str(exc)
        except Exception:
            example_bar.close()
            raise
        finally:
            example_bar.n = produced
            example_bar.refresh()
            example_bar.close()

        status = "complete"
        if produced == 0:
            status = "empty"
        elif produced < task_sample_size:
            status = "partial"  # Fewer examples than requested, but no duplication

        # NOTE: We intentionally do NOT upsample/duplicate examples.
        # Duplication creates artificial correlations that harm router training.
        # If fewer examples are available, we use only the unique examples.
        duplicates_added = 0  # Always zero - no duplication

        if produced == 0:
            raise ValueError(
                f"Task '{spec.task_id}' produced no examples (requested={task_sample_size})."
            )

        rows.extend(task_rows)

        duration = time.perf_counter() - task_start
        if wandb_run is not None:
            wandb_run.log(
                {
                    "dataset/task": spec.task_id,
                    "dataset/task_examples": len(task_rows),
                    "dataset/task_unique_examples": produced,
                    "dataset/task_duplicates": duplicates_added,
                    "dataset/task_duration_sec": duration,
                    "dataset/task_status": status,
                },
                commit=False,
            )

        task_stats.append(
            {
                "task_id": spec.task_id,
                "requested_examples": task_sample_size,
                "produced_examples": produced,
                "duplicates_added": duplicates_added,
                "final_examples": len(task_rows),
                "status": status,
                "error": error_message,
            }
        )

        task_bar.set_postfix(task=spec.task_id, rows=len(task_rows), unique=produced, status=status)
        task_bar.update(1)

    task_bar.close()

    total_duration = time.perf_counter() - pipeline_start
    frame = pd.DataFrame(
        rows,
        columns=["task_id", "query", "ground_truth", "metric", "task_description"],
    )
    frame.to_csv(output_path, index=False)

    tasks_with_rows = [stat for stat in task_stats if stat["final_examples"] > 0]
    metadata_path = Path(output_path).with_suffix(Path(output_path).suffix + ".metadata.json")
    metadata = {
        "generated_at": datetime.now(tz=timezone.utc).isoformat(),
        "output_path": output_path,
        "requested_sample_size": sample_size,
        "total_examples": len(frame),
        "tasks": task_stats,
    }
    metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")

    if wandb_run is not None:
        wandb_run.log(
            {
                "dataset/output_path": output_path,
                "dataset/total_examples": len(frame),
                "dataset/num_tasks": len(tasks_with_rows),
                "dataset/sample_size": sample_size,
                "dataset/empty_tasks": len([s for s in task_stats if s["produced_examples"] == 0]),
                "dataset/total_duration_sec": total_duration,
            }
        )

    return frame


def _load_wandb_key() -> str:
    """Load W&B API key from ~/.ember/config.yaml (canonical secret store)."""
    ember_config_path = Path.home() / ".ember" / "config.yaml"
    if not ember_config_path.exists():
        raise FileNotFoundError(
            f"Ember config not found at {ember_config_path}. "
            "Run 'ember setup' or create the file with wandb.api_key set."
        )
    ember_config = yaml.safe_load(ember_config_path.read_text(encoding="utf-8"))
    wandb_section = ember_config.get("wandb") or {}
    api_key = str(wandb_section.get("api_key", "")).strip()
    if not api_key:
        raise ValueError(
            f"W&B API key not found in {ember_config_path}. "
            "Add 'wandb:\\n  api_key: <your-key>' to the file."
        )
    return api_key


if __name__ == "__main__":
    # Use _MODULE_DIR defined at module level for consistent path resolution
    CONFIG_PATH = _MODULE_DIR / "configs" / "config.yaml"
    config = yaml.safe_load(CONFIG_PATH.read_text(encoding="utf-8"))

    output_path = str(_MODULE_DIR / config["unified_qa_data_path"])
    sample_size = int(config["dataset_sample_size"])

    # Load W&B key from canonical secret store (~/.ember/config.yaml)
    api_key = _load_wandb_key()

    wandb.login(key=api_key, relogin=True)
    run = wandb.init(
        project="graph_router_dataset",
        job_type="dataset_generation",
        config={
            "output_path": output_path,
            "sample_size": sample_size,
            "num_tasks": len(TASK_SPECS),
        },
    )

    try:
        generate_unified_qa_dataset(
            output_path=output_path,
            sample_size=sample_size,
            show_progress=True,
            wandb_run=run,
        )
    finally:
        run.finish()
