"""
WMDP: Weapons of Mass Destruction Proxy

Nathaniel Li, Alexander Pan, Anjali Gopal, Summer Yue,
Daniel Berrios, Alice Gatti, Justin D Li, Ann-Kathrin
Dombrowski, Shashwat Goel, Long Phan, et al.
https://arxiv.org/pdf/2403.03218

# eval w/ 500 randomly selected samples
inspect eval inspect_evals/wmdp_bio --limit 500
inspect eval inspect_evals/wmdp_chem --limit 500
inspect eval inspect_evals/wmdp_cyber --limit 500
"""

from typing import Any

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.scorer import choice
from inspect_ai.solver import multiple_choice


@task
def wmdp_bio() -> Task:
    return create_task("wmdp-bio")


@task
def wmdp_chem() -> Task:
    return create_task("wmdp-chem")


@task
def wmdp_cyber() -> Task:
    return create_task("wmdp-cyber")


def create_task(dataset_name: str) -> Task:
    dataset = hf_dataset(
        path="cais/wmdp",
        name=dataset_name,
        sample_fields=record_to_sample,
        split="test",
    )

    return Task(
        dataset=dataset,
        solver=multiple_choice(),
        scorer=choice(),
    )


def record_to_sample(record: dict[str, Any]) -> Sample:
    choices = {0: "A", 1: "B", 2: "C", 3: "D"}
    return Sample(
        input=record["question"],
        choices=record["choices"],
        target=choices[record["answer"]],
    )
