"""MuSR: Testing the Limits of Chain-of-thought with Multistep Soft Reasoning

Zayne Sprague, Xi Ye, Kaj Bostrom, Swarat Chaudhuri, and Greg Durrett.
https://arxiv.org/abs/2310.16049

# Example: Eval MuSR with team_allocation and CoT+
inspect eval musr.py -T domain=team_allocation -T prompt_technique=cot+

"""

import ast
from typing import Any, Dict, Literal

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, system_message

from inspect_evals.musr.prompts import (
    COT_PLUS_PROMPT,
    COT_PROMPT,
    MURDER_MYSTERIES_EXAMPLE,
    MURDER_MYSTERIES_HINT,
    OBJECT_PLACEMENTS_EXAMPLE,
    OBJECT_PLACEMENTS_HINT,
    REGULAR_PROMPT,
    SYSTEM_PROMPT,
    TEAM_ALLOCATION_EXAMPLE,
    TEAM_ALLOCATION_HINT,
)

DomainType = Literal["murder_mysteries", "object_placements", "team_allocation"]
PromptTechniqueType = Literal["regular", "cot", "cot+"]

DEFAULT_DOMAIN: DomainType = "murder_mysteries"
DEFAULT_PROMPT_TECHNIQUE: PromptTechniqueType = "regular"
DEFAULT_EXAMPLE_COUNT = 0


@task
def musr(
    domain: DomainType = DEFAULT_DOMAIN,
    prompt_technique: PromptTechniqueType = DEFAULT_PROMPT_TECHNIQUE,
    example_count: int = DEFAULT_EXAMPLE_COUNT,
) -> Task:
    """Inspect task implementing the MuSR benchmark.

    Args:
        domain (Literal["murder_mysteries", "object_placements", "team_allocation"]): Which domain in the dataset to evaluate.
            Defaults to "murder_mysteries".
        prompt_technique (Literal["regular", "cot", "cot+"]): The prompt technique to use. "regular" includes only the narrative
            and the question. "cot" uses chain-of-thought prompting. "cot+" includes a hint. Defaults to "regular".
        example_count (int): Number of solved examples to include at the beginning of each prompt. Defaults to 0. Currently only supports 1 example.
    """
    prompt = get_domain_prompt(domain, prompt_technique, example_count)

    dataset = hf_dataset(
        path="TAUR-Lab/MuSR",
        split=domain,
        sample_fields=record_to_sample,
        shuffle=True,
        auto_id=True,
    )

    return Task(
        dataset=dataset,
        solver=[
            system_message(SYSTEM_PROMPT),
            multiple_choice(template=prompt),
        ],
        scorer=choice(),
    )


def get_domain_prompt(
    domain: DomainType = DEFAULT_DOMAIN,
    prompt_technique: PromptTechniqueType = DEFAULT_PROMPT_TECHNIQUE,
    example_count: int = DEFAULT_EXAMPLE_COUNT,
) -> str:
    domain_info = {
        "murder_mysteries": {
            "hint": MURDER_MYSTERIES_HINT,
            "example": MURDER_MYSTERIES_EXAMPLE,
        },
        "object_placements": {
            "hint": OBJECT_PLACEMENTS_HINT,
            "example": OBJECT_PLACEMENTS_EXAMPLE,
        },
        "team_allocation": {
            "hint": TEAM_ALLOCATION_HINT,
            "example": TEAM_ALLOCATION_EXAMPLE,
        },
    }

    if domain not in domain_info:
        raise ValueError(
            "Unknown domain. Valid domains are murder_mysteries (default), object_placements, and team_allocation"
        )

    if prompt_technique == "regular":
        prompt = REGULAR_PROMPT
    elif prompt_technique == "cot":
        prompt = COT_PROMPT
    elif prompt_technique == "cot+":
        prompt = COT_PLUS_PROMPT.format(hint=domain_info[domain]["hint"])
    else:
        raise ValueError(
            "Unknown prompt technique. Valid prompt techniques are regular (default), cot, and cot+."
        )

    if example_count > 1:
        raise ValueError(">1 examples currently not supported")
    if example_count == 1:
        return f"Here is an example of solving the task:\n\n{domain_info[domain]['example']}\n\nThis is the end of the example. The real task is below.\n\n---\n\n{prompt}"
    else:
        return prompt


def record_to_sample(record: Dict[str, Any]) -> Sample:
    return Sample(
        input=f"{record['narrative']}\n\n{record['question']}",
        choices=ast.literal_eval(record["choices"]),
        target=chr(ord("A") + int(record["answer_index"])),
    )
