"""
Instruction-Following Evaluation for Large Language Models

Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan,
Denny Zhou, Le Hou
https://arxiv.org/pdf/2311.07911

Based on: https://github.com/google-research/google-research/tree/master/instruction_following_eval
"""

from typing import Any, cast

import numpy as np
from inspect_ai import Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.scorer import (
    Metric,
    SampleScore,
    Score,
    Scorer,
    Target,
    Value,
    metric,
    scorer,
)
from inspect_ai.solver import TaskState, generate


@task
def ifeval() -> Task:
    """Inspect Task implementation for the Instruction-Following benchmark"""
    return Task(
        dataset=hf_dataset(
            path="google/IFEval", split="train", sample_fields=record_to_sample
        ),
        solver=[generate()],
        scorer=instruction_following(),
    )


@metric
def if_metric() -> Metric:
    def _final_accuracy_stderr(
        scores: list[SampleScore], mean_final_accuracy: float
    ) -> float:
        # This is a bit convoluted; see the README in this directory for mathematical details.
        total_num_instructions = int(
            sum(
                cast(dict[str, Any], score.score.value)["num_instructions"]
                for score in scores
            )
        )
        mean_num_instructions = total_num_instructions / len(scores)
        variance = 0.0
        cluster_count = len(scores)
        for score in scores:
            value = cast(dict[str, Any], score.score.value)
            inst_level_strict = int(value["inst_level_strict"])
            inst_level_loose = int(value["inst_level_loose"])
            prompt_level_strict = int(value["prompt_level_strict"])
            prompt_level_loose = int(value["prompt_level_loose"])
            num_instructions = int(value["num_instructions"])

            loose_only = inst_level_loose - inst_level_strict
            num_incorrect = int(num_instructions - inst_level_loose)
            prompt_adjustment = (
                0.25
                * (prompt_level_strict + prompt_level_loose)
                * mean_num_instructions
                / num_instructions
            )
            vector = [
                (0.5 + prompt_adjustment - mean_final_accuracy) * inst_level_strict,
                (0.25 + prompt_adjustment - mean_final_accuracy) * loose_only,
                (0.0 + prompt_adjustment - mean_final_accuracy) * num_incorrect,
            ]
            variance += np.outer(vector, vector).sum()

        stderr = (
            np.sqrt(variance * cluster_count / (cluster_count - 1))
            / total_num_instructions
            if cluster_count > 1
            else 0.0
        )

        return stderr

    def metric(scores: list[SampleScore]) -> Value:
        statistics: list[float] = []
        prompt_keys = ["prompt_level_strict", "prompt_level_loose"]
        instruct_keys = ["inst_level_strict", "inst_level_loose"]
        final_keys = [
            "prompt_strict_acc",
            "prompt_strict_stderr",
            "prompt_loose_acc",
            "prompt_loose_stderr",
            "inst_strict_acc",
            "inst_strict_stderr",
            "inst_loose_acc",
            "inst_loose_stderr",
            "final_acc",
            "final_stderr",
        ]

        # calculate prompt-level accuracies + stderrs
        for key in prompt_keys:
            score_lst = [
                cast(dict[str, Any], score.score.value)[key] for score in scores
            ]
            statistics.append(np.mean(score_lst).item())
            stderr = (
                np.std(score_lst, ddof=1).item() / np.sqrt(len(score_lst))
                if len(score_lst) > 1
                else 0.0
            )
            statistics.append(stderr)

        # calculate instruction-level accuracies + clustered stderrs
        for key in instruct_keys:
            flattened = []
            for score in scores:
                value = cast(dict[str, Any], score.score.value)
                num_correct = int(value[key])
                num_incorrect = int(value["num_instructions"] - value[key])
                flattened.extend([True] * num_correct + [False] * num_incorrect)

            mean = np.mean(flattened).item()
            statistics.append(mean)

            # Because the inclusion of instructions are correlated, we need to cluster
            # the standard errors by prompt. The clustered calculation follows the logic
            # in the main stderr(cluster="cluster") code found in inspect_ai.scorer.
            variance = 0.0
            cluster_count = len(scores)
            for score in scores:
                value = cast(dict[str, Any], score.score.value)
                num_correct = int(value[key])
                num_incorrect = int(value["num_instructions"] - value[key])
                vector = [num_correct * (1 - mean), num_incorrect * (0 - mean)]
                variance += np.outer(vector, vector).sum()

            stderr = (
                np.sqrt(variance * cluster_count / (cluster_count - 1)) / len(flattened)
                if cluster_count > 1
                else 0.0
            )
            statistics.append(stderr)

        # Calculate the final accuracy and its standard error
        statistics.append(
            np.mean([statistics[i] for i in range(0, len(statistics), 2)]).item()
        )
        statistics.append(_final_accuracy_stderr(scores, statistics[-1]))

        return {k: v for k, v in zip(final_keys, statistics, strict=True)}

    return metric


@scorer(metrics=[if_metric()])
def instruction_following() -> Scorer:
    from instruction_following_eval.evaluation import (  # type: ignore
        InputExample,
        ensure_nltk_resource,
        test_instruction_following,
    )

    ensure_nltk_resource()  # Required before calling test_instruction_following

    async def score(state: TaskState, target: Target) -> Score:
        # construct the input to IFEval's evaluation functions using the data class
        eval_input = InputExample(
            key=state.sample_id,
            instruction_id_list=state.metadata["instruction_id_list"],
            prompt=state.metadata["prompt"],
            kwargs=state.metadata["kwargs"],
        )

        # retrieve evaluated outputs
        out_strict = test_instruction_following(
            eval_input, state.output.completion, strict=True
        )
        out_loose = test_instruction_following(
            eval_input, state.output.completion, strict=False
        )
        ret_value = {
            "prompt_level_strict": out_strict.follow_all_instructions,
            "inst_level_strict": sum(out_strict.follow_instruction_list),
            "prompt_level_loose": out_loose.follow_all_instructions,
            "inst_level_loose": sum(out_loose.follow_instruction_list),
            "num_instructions": len(out_loose.follow_instruction_list),
        }

        # return score with resulting outputs, model answer, and the
        # expected instructions
        return Score(
            value=ret_value,
            answer=state.output.completion,
            explanation=" ".join(state.metadata["instruction_id_list"]),
        )

    return score


# map ifeval record into inspect sample
def record_to_sample(record: dict[str, Any]) -> Sample:
    new_kwargs = {}
    for index in range(len(record["instruction_id_list"])):
        # remove None values from kwargs to avoid unexpected keyword argument errors
        # in build_description method from the IFEval package.
        kwargs = {k: v for k, v in record["kwargs"][index].items() if v}
        new_kwargs[index] = kwargs

    return Sample(
        id=record["key"],
        input=record["prompt"],
        metadata={
            "prompt": record["prompt"],
            "instruction_id_list": record["instruction_id_list"],
            "kwargs": new_kwargs,
        },
    )
