import re
from typing import Union

from inspect_ai.scorer import Score, Scorer, Target, accuracy, scorer, stderr
from inspect_ai.solver._task_state import TaskState


class ListTarget(Target):
    """A subclass of Target that returns the target as a list."""

    def __init__(self, target: list[Union[int, float]]):
        if not isinstance(target, list):
            raise ValueError(f"Expected a list, got {type(target)}")

        # Convert numbers to strings for the parent class
        str_target = [str(x) for x in target]
        super().__init__(str_target)

        # Store the original numeric values
        self._numeric_target = target

    @property
    def as_list(self) -> list[Union[int, float]]:
        """Returns the target as a list of numbers."""
        return self._numeric_target


# Logic based on source code from InfiniteBench:
# https://github.com/OpenBMB/InfiniteBench/blob/51d9b37b0f1790ead936df2243abbf7f0420e439/src/compute_scores.py#L319
@scorer(metrics=[accuracy(), stderr()])
def math_calc_scorer(model_name: str) -> Scorer:
    """
    Custom scorer for Math.Calc task which evaluates the number of correct numbers in the result.

    Args:
        model_name (str): The name of the model to score the task with.

    Returns:
        Score: The score for the task.
    """

    async def score(state: TaskState, target: Target) -> Score:
        # Convert string targets to numbers
        numeric_target = [float(x) if "." in x else int(x) for x in target.target]
        number_target = ListTarget(numeric_target)

        answer = state.output.completion
        ground_truth = number_target.as_list

        # Convert answer into a list
        # Use regex to extract the numbers from the answer string (assuming it's a list representation)
        answer_list = re.findall(r"-?\d+\.?\d*", answer)

        # The prompts result in GPT models always outputting the first number as the first value in the predicted answer.
        # This logic is based on the source code from InfiniteBench:
        # https://github.com/OpenBMB/InfiniteBench/blob/51d9b37b0f1790ead936df2243abbf7f0420e439/src/compute_scores.py#L330C5-L331C31
        if model_name.startswith("gpt"):
            answer_list = answer_list[1:]

        # Ensure ground_truth is a list of strings so that it can be compared to answer_list
        ground_truth_list = [str(item) for item in ground_truth]

        # Count correct numbers before the first error
        cnt = 0
        for i in range(len(ground_truth_list)):
            if i >= len(answer_list):
                break
            if ground_truth_list[i] == answer_list[i]:
                cnt += 1
            else:
                break  # Stop counting once a mismatch is found

        # Calculate the accuracy as the fraction of correct numbers
        score_value = cnt / len(ground_truth_list)

        return Score(
            value=score_value,
            explanation=f"Correct numbers before first error: {cnt}/{len(ground_truth_list)}",
        )

    return score
