"""
LAB-Bench

Measuring Capabilities of Language Models for Biology Research

Jon M. Laurent, Joseph D. Janizek, Michael Ruzo1, Michaela M. Hinks, Michael J. Hammerling, Siddharth Narayanan, Manvitha Ponnapati, Andrew D. White, Samuel G. Rodriques

FutureHouse Inc., San Francisco, CA

https://arxiv.org/abs/2407.10362
"""

from enum import Enum
from typing import Any, Callable

from inspect_ai import Epochs, Task, task
from inspect_ai.dataset import Dataset, Sample, hf_dataset
from inspect_ai.solver import multiple_choice

from inspect_evals.lab_bench.scorer import precision_choice

from .record_to_sample_helpers import (
    UNCERTAIN_ANSWER_CHOICE,
    record_to_sample_base,
    record_to_sample_figqa,
    record_to_sample_protocolqa,
    record_to_sample_suppqa,
    record_to_sample_tableqa,
)

DATASET_PATH = "futurehouse/lab-bench"


class DatasetSubsets(Enum):
    LitQA2 = "LitQA2"
    SuppQA = "SuppQA"
    FigQA = "FigQA"
    TableQA = "TableQA"
    DbQA = "DbQA"
    ProtocolQA = "ProtocolQA"
    SeqQA = "SeqQA"
    CloningScenarios = "CloningScenarios"


MULTIPLE_CHOICE_TEMPLATE = """
The following is a multiple choice question about biology.
Please answer by responding with the letter of the correct answer.

Think step by step.

Question: {question}
Options:
{choices}

You MUST include the letter of the correct answer within the following format: 'ANSWER: $LETTER' (without quotes). For example, ’ANSWER: <answer>’, where <answer> is the correct letter. Always answer in exactly this format of a single letter, even if you are unsure. We require this because we use automatic parsing.
"""


# helper function for retrieving HF dataset for all subsets
def retrieve_hf_dataset(
    name: str,
    record_to_sample_fields_func: Callable[[dict[str, Any]], Sample],
    shuffle_dataset: bool = True,
    shuffle_choices: bool = True,
) -> Dataset:
    return hf_dataset(
        path=DATASET_PATH,
        name=name,
        split="train",
        sample_fields=record_to_sample_fields_func,
        trust=True,
        shuffle=shuffle_dataset,
        shuffle_choices=shuffle_choices,
    )


@task
def lab_bench_litqa() -> Task:
    """Inspect Task implementing the LAB-Bench benchmark LitQA"""
    return task_base(
        retrieve_hf_dataset(DatasetSubsets.LitQA2.value, record_to_sample_base)
    )


@task
def lab_bench_suppqa() -> Task:
    """Inspect Task implementing the LAB-Bench benchmark SuppQA"""
    return task_base(
        retrieve_hf_dataset(DatasetSubsets.SuppQA.value, record_to_sample_suppqa)
    )


@task
def lab_bench_figqa() -> Task:
    """Inspect Task implementing the LAB-Bench benchmark FigQA"""
    return task_base(
        retrieve_hf_dataset(DatasetSubsets.FigQA.value, record_to_sample_figqa)
    )


@task
def lab_bench_tableqa() -> Task:
    """Inspect Task implementing the LAB-Bench benchmark TableQA"""
    return task_base(
        retrieve_hf_dataset(DatasetSubsets.TableQA.value, record_to_sample_tableqa)
    )


@task
def lab_bench_dbqa() -> Task:
    """Inspect Task implementing the LAB-Bench benchmark DbQA"""
    return task_base(
        retrieve_hf_dataset(DatasetSubsets.DbQA.value, record_to_sample_base)
    )


@task
def lab_bench_protocolqa() -> Task:
    """Inspect Task implementing the LAB-Bench benchmark ProtocolQA"""
    return task_base(
        retrieve_hf_dataset(
            DatasetSubsets.ProtocolQA.value, record_to_sample_protocolqa
        )
    )


@task
def lab_bench_seqqa() -> Task:
    """Inspect Task implementing the LAB-Bench benchmark SeqQA"""
    return task_base(
        retrieve_hf_dataset(DatasetSubsets.SeqQA.value, record_to_sample_base)
    )


@task
def lab_bench_cloning_scenarios() -> Task:
    """Inspect Task implementing the LAB-Bench benchmark CloningScenarios"""
    return task_base(
        retrieve_hf_dataset(
            DatasetSubsets.CloningScenarios.value, record_to_sample_base
        )
    )


def task_base(dataset: Dataset) -> Task:
    return Task(
        dataset=dataset,
        solver=[multiple_choice(template=MULTIPLE_CHOICE_TEMPLATE, cot=True)],
        scorer=precision_choice(no_answer=UNCERTAIN_ANSWER_CHOICE),
        epochs=Epochs(1, "mode"),
    )
