from typing import Any

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.scorer import match
from inspect_ai.solver import (
    Solver,
    generate,
    prompt_template,
)

USER_PROMPT_TEMPLATE = """
Solve the following math problem step by step.
The last line of your response should be of the form "ANSWER: $ANSWER" (without quotes) where $ANSWER is the answer to the problem.

{prompt}

Remember to put your answer on its own line at the end in the form "ANSWER: $ANSWER" (without quotes) where $ANSWER is the answer to the problem, and you do not need to use a \\boxed command.
""".strip()


@task
def aime2024() -> Task:
    """Inspect Task implementation for the AIME 2024 benchmark."""
    dataset = hf_dataset(
        path="Maxwell-Jia/AIME_2024",
        split="train",
        trust=True,
        sample_fields=record_to_sample,
    )

    return Task(
        dataset=dataset,
        solver=aime2024_solver(),
        scorer=[
            match(),
        ],
    )


def aime2024_solver() -> list[Solver]:
    """Build solver for AIME 2024 task."""
    solver = [prompt_template(USER_PROMPT_TEMPLATE), generate()]
    return solver


def record_to_sample(record: dict[str, Any]) -> Sample:
    sample = Sample(
        id=record["ID"],
        input=record["Problem"],
        target=str(record["Answer"]),
        metadata={
            "solution": record["Solution"],
        },
    )
    return sample
