"""LitQA2 task from LAB-bench.

This module adapts LAB-bench's LitQA2 task to work with inspect_ai:
https://github.com/Future-House/LAB-Bench/blob/77c14efc2192a1b3d14f695e120c668eef72cd82/LitQA2/task.py
"""

import random
import string
from typing import Literal, NamedTuple, cast

import pandas as pd
from datasets import Dataset, DatasetDict, IterableDatasetDict, load_dataset
from inspect_ai import Task, task, task_with
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.scorer import (
    Metric,
    SampleScore,
    Score,
    Scorer,
    Target,
    accuracy,
    choice,
    metric,
    scorer,
    stderr,
)
from inspect_ai.solver import TaskState, use_tools
from pydantic import BaseModel, Field

from astabench.constants import ASTA_BENCH_DATASET_REPO, ASTA_BENCH_DATASET_REVISION
from astabench.evals.utils import (
    format_multichoice_as_textin,
    mark_multichoice_answer,
    not_implemented_solver,
)
from astabench.tools import make_asta_mcp_tools, make_native_search_tools

# See https://github.com/allenai/nora-issues-research/issues/98#issuecomment-2878244637
# Set to one day later than the min insertion date of 2024-10-16
INSERTED_BEFORE = "2024-10-17"


class LabbenchChoices(NamedTuple):
    choices: list[str]
    target_idx: int
    unsure_idx: int


class LabbenchQuestion(BaseModel):
    id: str
    question: str
    ideal: str
    distractors: list[str]
    sources: list[str]
    canary: str
    version: str
    is_opensource: bool = False
    key_passage: str | None = Field(alias="key-passage")

    def full_choices(
        self, shuffle: bool = True, rng: random.Random | None = None
    ) -> LabbenchChoices:
        """Get all choices, including the ideal answer and an injected "Insufficient information" answer.

        Args:
            shuffle: Whether to shuffle the choices.

        Returns:
            A tuple of the shuffled choices, the index of the ideal answer, and the index of the "Insufficient information" answer.

        """
        choices = [self.ideal] + self.distractors

        # This is the 'unsure' answer that LAB-bench injects
        # https://github.com/Future-House/LAB-Bench/blob/77c14efc2192a1b3d14f695e120c668eef72cd82/labbench/utils.py#L25
        choices.append("Insufficient information to answer the question")
        unsure_idx = len(choices) - 1
        target_idx = 0

        if shuffle:
            if not rng:
                rng = random.Random()
            # Shuffle the choices but keep track of permutations so we can
            # adjust the indexes
            perm = list(range(len(choices)))
            rng.shuffle(perm)
            choices = [choices[i] for i in perm]
            target_idx = perm.index(0)
            unsure_idx = perm.index(unsure_idx)
        return LabbenchChoices(
            choices=choices, target_idx=target_idx, unsure_idx=unsure_idx
        )


def json_to_sample(obj: dict, rng: random.Random | None = None) -> Sample:
    """Convert LAB-bench litqa2 format to inspect_ai Sample format."""
    instance = LabbenchQuestion.model_validate(obj)

    combined_choices = instance.full_choices(shuffle=True, rng=rng)

    return Sample(
        id=str(instance.id),
        input=instance.question,
        choices=[choice for choice in combined_choices.choices],
        target=string.ascii_uppercase[combined_choices.target_idx],
        metadata={
            "is_opensource": instance.is_opensource,
            "unsure_letter": string.ascii_uppercase[combined_choices.unsure_idx],
        },
    )


@metric
def coverage() -> Metric:
    """Coverage is the fraction of sure answers (i.e. did not select
    "Insufficient information")."""

    def metric_fn(scores: list[SampleScore]) -> float:
        return sum(score.score.value["is_sure"] for score in scores) / len(scores)  # type: ignore

    return metric_fn


@metric
def precision() -> Metric:
    """Precision is the fraction of correct answers out of sure answers (i.e.
    did not select "Insufficient information")."""

    def metric_fn(scores: list[SampleScore]) -> float:
        score_vals: list[dict[str, bool]] = [
            cast(dict[str, bool], score.score.value) for score in scores
        ]

        return sum(v["is_correct"] and v["is_sure"] for v in score_vals) / max(
            1, sum(v["is_sure"] for v in score_vals)
        )

    return metric_fn


@scorer(
    metrics=[
        {"is_correct": [accuracy(), stderr()]},
        precision(),
        coverage(),
    ]
)
def score_litqa2() -> Scorer:
    """Score LitQA2 responses.  This uses InspectAI's choice() internally, but
    returns additional metrics for is_sure because litqa2 allows 'insufficient
    info' as an answer option.

    This scorer prioritizes answer choices marked in `state.choices`; if none
    are marked, it tries parsing the output text using
    `astabench.evals.utils.mark_multichoice_answer()`, which works with the
    astabench `{"answer": "<letter>"}` format or the Labbench
    `[ANSWER]<letter>[/ANSWER]` format."""

    choice_scorer = choice()

    async def score(state: TaskState, target: Target) -> Score:
        if all(c.correct is None for c in state.choices):
            # No choices marked; try to parse from output
            state = await mark_multichoice_answer(state)

        if all((not c.correct) for c in state.choices):
            # Still nothing marked
            return Score(
                value={"is_correct": False, "is_sure": False},
                answer=state.output.completion,
                explanation="No choice selection found in output",
            )

        score = await choice_scorer(state, target)
        score.value = {
            "is_correct": score.value,  # type: ignore
            "is_sure": score.answer != state.metadata["unsure_letter"],
        }
        return score

    return score


def load_litqa2(split: Literal["dev", "test", "all"] = "all") -> list[dict]:
    """Get LitQA2 v2 data from LAB-bench's HuggingFace dataset.

    Args:
        split: The split to load. Can be "dev", "test", or "all".  If "all",
               returns the full litqa2 dataset.  Otherwise, returns the
               dev/test splits from the astabench mapping, both of which are
               filtered to only include papers that have fulltext entries in
               AI2's snippet search index.
    """
    dataset = load_dataset("futurehouse/lab-bench", "LitQA2")
    assert isinstance(dataset, Dataset | DatasetDict | IterableDatasetDict)

    if split == "all":
        return list(map(dict, dataset["train"]))

    litqa_mappings = load_dataset(
        ASTA_BENCH_DATASET_REPO,
        data_files="tasks/labbench/litqa2_mapping.json",
        revision=ASTA_BENCH_DATASET_REVISION,
    )

    df = pd.merge(
        dataset["train"].to_pandas(),
        litqa_mappings["train"].to_pandas().drop("question", axis=1),
        how="left",
        on="id",
    )
    split_subset = Dataset.from_pandas(df[df["split"] == split])

    return list(map(dict, split_subset))


@task
def litqa2_inspect(
    seed: int = 0,
    with_search_tools: bool = True,
    with_native_search_tools: bool = False,
    split: Literal["dev", "test", "all"] = "all",
):
    """LitQA2 task from LAB-bench.

    This task presents multiple choice questions about scientific literature. Each question
    has a correct answer, an "Insufficient information" option, and several distractors.

    The scoring is designed to be compatible with InspectAI's `multiple_choice`
    solver; that is, it provides choices in `state.choices` and expects the
    solver to mark those choices as true or false in the `state` object.  If
    choices are not marked, it will attempt to parse output text as
    a fallback (see `astabench.evals.utils.mark_multichoice_answer()`).

    See the LAB-bench prompt for more details on how they instruct the agent to format output:
    https://github.com/Future-House/LAB-Bench/blob/77c14efc2192a1b3d14f695e120c668eef72cd82/labbench/zero_shot.py#L15-L26
    """
    # Get LAB-bench litqa2 data
    data = load_litqa2(split=split)

    # Note: we need to shuffle in json_to_sample instead of using
    # `dataset.shuffle_choices` because `shuffle_choices` only remaps the
    # target letter, but we want to track both the target letter and the
    # "unsure" letter.
    rng = random.Random(seed)
    dataset = MemoryDataset([json_to_sample(d, rng) for d in data])

    tool_setups = []
    if with_search_tools:
        tool_setups.append(
            use_tools(make_asta_mcp_tools(insertion_date=INSERTED_BEFORE))
        )
    if with_native_search_tools:
        native_search_tools = make_native_search_tools(inserted_before=INSERTED_BEFORE)
        if native_search_tools:
            tool_setups.append(use_tools(native_search_tools))

    return Task(
        dataset=dataset,
        solver=not_implemented_solver(),  # Call the function to get the Solver instance
        scorer=score_litqa2(),
        setup=tool_setups or None,
    )


@task
def litqa2(*litqa_args, **litqa_kwargs):
    """LitQA2 task from LAB-bench, formatted to use only the astabench
    textin/textout format instead of Inspect's `state.choices`.  See litqa2()
    for more task details."""

    base_task = litqa2_inspect(*litqa_args, **litqa_kwargs)

    dataset = []
    for sample in base_task.dataset:
        # Litqa base task is randomized by default, no need to shuffle
        sample = format_multichoice_as_textin(sample, shuffle=False)
        dataset.append(sample)

    return task_with(
        base_task,
        dataset=MemoryDataset(dataset),
        scorer=score_litqa2(),
    )


@task
def litqa2_validation(*litqa_args, **litqa_kwargs):
    """LitQA2 but pinned to the validation/"dev" set."""

    if litqa_kwargs.get("split", "dev") != "dev":
        raise ValueError("Must have split dev for litqa2_validation")

    litqa_kwargs["split"] = "dev"

    return litqa2(*litqa_args, **litqa_kwargs)


@task
def litqa2_test(*litqa_args, **litqa_kwargs):
    """LitQA2 but pinned to the test set."""

    if litqa_kwargs.get("split", "test") != "test":
        raise ValueError("Must have split test for litqa2_test")

    litqa_kwargs["split"] = "test"

    return litqa2(*litqa_args, **litqa_kwargs)
