"""
Measuring Massive Multitask Language Understanding

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou,
Mantas Mazeika, Dawn Song, Jacob Steinhardt
https://arxiv.org/abs/2009.03300

Based on: https://github.com/openai/simple-evals/blob/main/mmlu_eval.py

# eval all subjects w/ 500 randomly selected samples
inspect eval inspect_evals/mmlu_0_shot --limit 500

# add chain of thought
inspect eval inspect_evals/mmlu_0_shot --limit 500 -T cot=true

# eval selected subjects
inspect eval inspect_evals/mmlu_0_shot -T subjects=anatomy
inspect eval inspect_evals/mmlu_0_shot -T subjects=astronomy
inspect eval inspect_evals/mmlu_0_shot -T subjects=anatomy,astronomy
"""

from typing import Any, Literal, Union

from inspect_ai import Task, task
from inspect_ai.dataset import Dataset, Sample, hf_dataset
from inspect_ai.model import GenerateConfig
from inspect_ai.scorer import choice, exact
from inspect_ai.solver import (
    Choices,
    Generate,
    Solver,
    TaskState,
    generate,
    multiple_choice,
    solver,
)

MMLU_MULTISHOT_PROMPT_TEMPLATE = r"""
The following are multiple choice questions about {subject_name}.

{subject_dev_questions}

{test_question}
""".strip()

MMLU_MULTISHOT_QUESTION_TEMPLATE = r"""
{question}
{choices}
Answer: {answer}
""".strip()


@task
def mmlu_0_shot(subjects: str | list[str] = [], cot: bool = False) -> Task:
    """
    Inspect Task implementation for MMLU, with 0-shot prompting.

    Args:
        subjects (str | list[str]): Subjects to filter to
        cot (bool): Whether to use chain of thought
    """
    return Task(
        # (shuffle so that --limit draws from multiple subjects)
        dataset=get_mmlu_dataset("test", shuffle=True, subjects=subjects),
        solver=multiple_choice(cot=cot),
        scorer=choice(),
        config=GenerateConfig(temperature=0.5),
    )


@task
def mmlu_5_shot(subjects: str | list[str] = []) -> Task:
    """
    Inspect Task implementation for MMLU, with 5-shot prompting.

    Args:
        subjects (str | list[str]): Subjects to filter to
    """
    return Task(
        # (shuffle so that --limit draws from multiple subjects)
        dataset=get_mmlu_dataset("test", shuffle=True, subjects=subjects),
        solver=[
            mmlu_5_shot_solver(),
            # To perfectly match the MMLU paper, we should be getting the logprobs
            # for the tokens 'A', 'B', 'C' and 'D' and pick the greatest one. We
            # instead get a simple completion, given logprobs are not supported
            # by many models supported by Inspect.
            generate(max_tokens=1),
        ],
        scorer=exact(),
        config=GenerateConfig(temperature=0),
    )


@solver
def mmlu_5_shot_solver() -> Solver:
    """A custom solver for MMLU 5-shot.

    It uses the dev dataset to generate a prompt with 5 example answered questions.

    It expects the TaskState to have a metadata field "subject" that is the
    subject of the question, e.g. "astronomy".

    Returns:
        Solver: A solver that uses the dev dataset to generate a prompt with
          5 example answered questions.
    """
    dev_dataset = get_mmlu_dataset("dev", shuffle=False)

    async def solve(state: TaskState, generate: Generate) -> TaskState:
        # e.g. "abstract_algebra" -> "abstract algebra"
        subject_name = state.metadata["subject"].replace("_", " ")

        subject_dev_dataset = dev_dataset.filter(
            lambda sample: sample.metadata is not None
            and sample.metadata["subject"] == state.metadata["subject"]
        )
        subject_dev_questions = "\n\n".join(
            [
                format_mmlu_question(
                    question=str(sample.input),
                    choices=Choices(
                        sample.choices if sample.choices is not None else []
                    ),
                    answer=str(sample.target),
                )
                for sample in subject_dev_dataset
            ]
        )

        test_question = format_mmlu_question(
            question=str(state.input),
            choices=state.choices,
        )

        state.user_prompt.text = MMLU_MULTISHOT_PROMPT_TEMPLATE.format(
            subject_name=subject_name,
            subject_dev_questions=subject_dev_questions,
            test_question=test_question,
        )

        return state

    return solve


def format_mmlu_question(question: str, choices: Choices, answer: str = "") -> str:
    r"""Format an MMLU-style prompt for a question, choices and answer.

    Args:
        question (str): Question to pose.
        choices (Choices): Choices to select from.
        answer (str, optional): Correct answer. Defaults to "". If this is not
          provided the prompt will still include space for an answer.

    Returns:
        Prompt for question and choices.

    Example:
    >>> format_mmlu_question("What is 5 squared?", Choices(["20", "25", "50", "55"]), "B")
    "What is 5 squared?\n(A) 20\n(B) 25\n(C) 50\n(D) 55\nAnswer: B"
    """
    return MMLU_MULTISHOT_QUESTION_TEMPLATE.format(
        question=question, choices=format_mmlu_choices(choices), answer=answer
    )


def format_mmlu_choices(choices: Choices) -> str:
    r"""Format choices in the style of the MMLU paper.

    Returns the `choices` formatted as a multiple choice question in style of the
    MMLU paper (with round parentheses around letters). This is necessary as by
    default Inspect will only put parentheses after letters.

    Args:
        choices (Choices): Choices to select from.

    Returns:
        Formatted choices.

    Example:
    >>> format_mmlu_choices(Choices(["20", "25", "50", "55"]))
    "(A) 20\n(B) 25\n(C) 50\n(D) 55"
    """
    indexes = list(range(len(choices)))

    return "\n".join(
        [f"({chr(65 + i)}) {choices[j].value}" for i, j in enumerate(indexes)]
    )


def get_mmlu_dataset(
    split: Union[Literal["test"], Literal["dev"], Literal["validation"]] = "test",
    shuffle: bool = False,
    subjects: Union[list[str], str] = [],
) -> Dataset:
    dataset = hf_dataset(
        path="cais/mmlu",
        name="all",
        split=split,
        sample_fields=record_to_sample,
        shuffle=shuffle,
        seed=42,
    )

    # filter dataset if requested
    subjects = subjects if isinstance(subjects, list) else [subjects]
    if len(subjects) > 0:
        return dataset.filter(
            name=f"{dataset.name}-{'-'.join(subjects)}",
            predicate=lambda sample: sample.metadata is not None
            and sample.metadata.get("subject") in subjects,
        )
    else:
        return dataset


def record_to_sample(record: dict[str, Any]) -> Sample:
    return Sample(
        input=record["question"],
        choices=record["choices"],
        # converts 0 -> A, 1 -> B, etc.
        target=("ABCD"[record["answer"]]),
        metadata={"subject": record["subject"]},
    )
