"""
BIG-Bench-Hard: Challenging BIG-Bench Tasks and Whether Chain-of-Thought Can Solve Them

Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay,
Hyung Won Chung, Aakanksha Chowdhery, Quoc V. Le, Ed H. Chi, Denny Zhou, Jason Wei

https://arxiv.org/pdf/2210.09261

# eval w/ 25 randomly selected samples from the date_understanding dataset with a 3-shot Chain-of-Thought prompt
inspect eval inspect_evals/bbh -T "dataset_name=date_understanding" -T "prompt_type=chain_of_thought" --limit 25
"""

from dataclasses import dataclass
from enum import Enum, auto
from typing import Any, Callable, List, Tuple

import datasets  # type: ignore
from inspect_ai import Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.model import GenerateConfig
from inspect_ai.scorer import choice, match
from inspect_ai.solver import generate, multiple_choice


@task
def bbh(
    dataset_name: str = "date_understanding", prompt_type: str = "answer_only"
) -> Task:
    """Inspect task implementing the BBH benchmark.

    Args:
        dataset_name: Name of dataset to use. See
                      BBHDatasetRegistry.MULTIPLE_CHOICE_DATASETS,
                      BBHDatasetRegistry.BINARY_CHOICE_DATASETS,
                      BBHDatasetRegistry.EXACT_MATCH_DATASETS or
                      BBHDatasetRegistry.DYCK_DATASET for valid options.
                     Defaults to "date_understanding".
        prompt_type: Type of prompt to use. One of ["zero_shot", "answer_only",
                    "chain_of_thought"]. Defaults to "answer_only".
    """
    try:
        prompt_type_enum = PromptType[prompt_type.upper()]
    except KeyError as e:
        raise ValueError(
            f"Invalid prompt type: {prompt_type}. "
            f"Must be one of: {[pt.value for pt in PromptType]}"
        ) from e

    dataset_config = BBHDatasetRegistry.get_dataset_config(dataset_name)

    prompt_manager = BBHPromptManager()
    prompt = prompt_manager.get_prompt(dataset_name, prompt_type_enum)

    def record_to_sample(record: dict[str, Any]) -> Sample:
        """Convert a dataset record to a Sample object with prompt."""
        sample = dataset_config.sample_converter(record)
        if prompt:
            if dataset_config.type == DatasetType.MULTIPLE_CHOICE:
                sample.input = f"{prompt}\n\nQUESTION: {sample.input}\nOPTIONS:"
            elif dataset_config.type in [
                DatasetType.BINARY_CHOICE,
                DatasetType.EXACT_MATCH,
                DatasetType.DYCK,
            ]:
                sample.input = f"{prompt}\n\nQUESTION: {sample.input}"
            else:
                raise NotImplementedError(
                    f"Prompt not implemented for dataset type: {dataset_config.type}"
                )
        return sample

    dataset = hf_dataset(
        path="Joschka/big_bench_hard",
        name=dataset_name,
        split=dataset_name,
        sample_fields=record_to_sample,
        trust=True,
        auto_id=True,
        shuffle=True,
    )

    if dataset_config.type == DatasetType.MULTIPLE_CHOICE:
        return Task(
            dataset=dataset,
            solver=multiple_choice(),
            scorer=choice(),
            config=GenerateConfig(temperature=0),
        )

    elif dataset_config.type in [
        DatasetType.BINARY_CHOICE,
        DatasetType.EXACT_MATCH,
        DatasetType.DYCK,
    ]:
        return Task(
            dataset=dataset,
            solver=generate(),
            scorer=match(location="end"),
            config=GenerateConfig(temperature=0),
        )

    else:
        raise NotImplementedError(
            f"Solver not implemented for dataset type: {dataset_config.type}"
        )


class DatasetType(Enum):
    """Enumeration for dataset types in the BIG-Bench-Hard benchmark."""

    MULTIPLE_CHOICE = auto()
    BINARY_CHOICE = auto()
    EXACT_MATCH = auto()
    DYCK = auto()


class PromptType(Enum):
    """Enumeration for types of prompts used in dataset processing."""

    ZERO_SHOT = "zero_shot"
    ANSWER_ONLY = "answer_only"
    CHAIN_OF_THOUGHT = "chain_of_thought"


@dataclass
class DatasetConfig:
    """Configuration for datasets, including name, type, and sample conversion method."""

    name: str
    type: DatasetType
    sample_converter: Callable[[dict[str, Any]], Sample]


class BBHDatasetRegistry:
    """Registry for BBH datasets with their configurations."""

    MULTIPLE_CHOICE_DATASETS: List[str] = [
        "date_understanding",
        "disambiguation_qa",
        "geometric_shapes",
        "hyperbaton",
        "logical_deduction_five_objects",
        "logical_deduction_seven_objects",
        "logical_deduction_three_objects",
        "movie_recommendation",
        "penguins_in_a_table",
        "reasoning_about_colored_objects",
        "ruin_names",
        "salient_translation_error_detection",
        "snarks",
        "temporal_sequences",
        "tracking_shuffled_objects_five_objects",
        "tracking_shuffled_objects_seven_objects",
        "tracking_shuffled_objects_three_objects",
    ]

    BINARY_CHOICE_DATASETS: List[Tuple[str, str, str]] = [
        ("boolean_expressions", "True", "False"),
        ("causal_judgement", "Yes", "No"),
        ("formal_fallacies", "valid", "invalid"),
        ("navigate", "Yes", "No"),
        ("sports_understanding", "yes", "no"),
        ("web_of_lies", "Yes", "No"),
    ]

    EXACT_MATCH_DATASETS: List[Tuple[str, str]] = [
        ("multistep_arithmetic_two", "integer"),
        ("object_counting", "natural"),
        ("word_sorting", "words"),
    ]

    DYCK_DATASET: List[str] = ["dyck_languages"]

    @classmethod
    def get_dataset_config(cls, dataset_name: str) -> DatasetConfig:
        """Get the configuration for a dataset."""
        if dataset_name in cls.MULTIPLE_CHOICE_DATASETS:
            return DatasetConfig(
                name=dataset_name,
                type=DatasetType.MULTIPLE_CHOICE,
                sample_converter=cls._multiple_choice_converter,
            )

        elif dataset_name in [name for name, _, _ in cls.BINARY_CHOICE_DATASETS]:
            return DatasetConfig(
                name=dataset_name,
                type=DatasetType.BINARY_CHOICE,
                sample_converter=cls._binary_choice_converter,
            )

        elif dataset_name in [name for name, _ in cls.EXACT_MATCH_DATASETS]:
            return DatasetConfig(
                name=dataset_name,
                type=DatasetType.EXACT_MATCH,
                sample_converter=cls._exact_match_converter,
            )
        elif dataset_name in [name for name in cls.DYCK_DATASET]:
            return DatasetConfig(
                name=dataset_name,
                type=DatasetType.DYCK,
                sample_converter=cls._dyck_converter,
            )

        raise ValueError(f"Unknown dataset: {dataset_name}")

    @staticmethod
    def _multiple_choice_converter(record: dict[str, Any]) -> Sample:
        """Convert a multiple choice dataset record to a Sample object."""
        choices = [str(text) for text in record["choices"]["text"]]
        return Sample(
            input=record["question"],
            choices=choices,
            target=record["target"],
        )

    @staticmethod
    def _binary_choice_converter(record: dict[str, Any]) -> Sample:
        """Convert a binary choice dataset record to a Sample object."""
        return Sample(input=record["question"], target=record["target"])

    @staticmethod
    def _exact_match_converter(record: dict[str, Any]) -> Sample:
        """Convert an exact match dataset record to a Sample object."""
        return Sample(input=record["question"], target=record["target"])

    @staticmethod
    def _dyck_converter(record: dict[str, Any]) -> Sample:
        """Convert a Dyck language record to a Sample object."""
        return Sample(input=record["question"], target=f"ANSWER: {record['target']}")


class BBHPromptManager:
    """Manages prompts for BBH tasks."""

    def __init__(self) -> None:
        self.prompts_dataset = datasets.load_dataset(
            path="Joschka/big_bench_hard",
            name="few_shot_prompts",
            split="few_shot_prompts",
            trust_remote_code=True,
        )

    def get_prompt(self, dataset_name: str, prompt_type: PromptType) -> str:
        """Get the appropriate prompt for a dataset and prompt type."""
        if prompt_type == PromptType.ZERO_SHOT:
            return ""

        prompt_row: dict[str, str] | None = next(
            (r for r in self.prompts_dataset if r["dataset_name"] == dataset_name), None
        )
        if prompt_row is None:
            raise ValueError(f"No prompts found for dataset: {dataset_name}")

        return prompt_row[prompt_type.value + "_prompt"]
