import re
import string
from logging import getLogger
from typing import Any

from inspect_ai.scorer import (
    CORRECT,
    INCORRECT,
    Score,
    Scorer,
    Target,
    accuracy,
    scorer,
    stderr,
)
from inspect_ai.solver import TaskState

logger = getLogger(__name__)


@scorer(metrics=[accuracy(), stderr()])
def gaia_scorer() -> Scorer:
    async def gaia_scorer(state: TaskState, target: Target) -> Score:
        answer = state.output.completion

        score, explanation = question_scorer(
            model_answer=answer, ground_truth=target.text
        )

        return Score(
            value=CORRECT if score else INCORRECT,
            answer=answer,
            explanation=explanation,
        )

    return gaia_scorer


### ---- CODE BELOW THIS POINT COPIED FROM  THE OFFICIAL GAIA SCORER AT https://huggingface.co/spaces/gaia-benchmark/leaderboard/blob/main/scorer.py ---- ###


def normalize_number_str(number_str: str) -> float:
    # we replacex these common units and commas to allow
    # conversion to float
    for char in ["$", "%", ","]:
        number_str = number_str.replace(char, "")
    try:
        return float(number_str)
    except ValueError:
        logger.warning(f"String {number_str} cannot be normalized to number str.")
        return float("inf")


def split_string(
    s: str,
    char_list: list[str] = [",", ";"],
) -> list[str]:
    pattern = f"[{''.join(char_list)}]"
    return re.split(pattern, s)


def question_scorer(
    model_answer: str,
    ground_truth: str,
) -> tuple[bool, str]:
    def is_float(element: Any) -> bool:
        try:
            float(element)
            return True
        except ValueError:
            return False

    # if gt is a number
    if is_float(ground_truth):
        normalized_answer = normalize_number_str(model_answer)
        return (
            normalized_answer == float(ground_truth),
            f"Evaluated {model_answer} as a number.",
        )

    # if gt is a list
    elif any(char in ground_truth for char in [",", ";"]):
        # question with the fish: normalization removes punct
        gt_elems = split_string(ground_truth)
        ma_elems = split_string(model_answer)

        # check length is the same
        if len(gt_elems) != len(ma_elems):
            return (
                False,
                f"Evaluated {model_answer} as a comma separated list, returned False because lists have different lenghts.",
            )

        # compare each element as float or str
        comparisons = []
        for ma_elem, gt_elem in zip(ma_elems, gt_elems):
            if is_float(gt_elem):
                normalized_ma_elem = normalize_number_str(ma_elem)
                comparisons.append(normalized_ma_elem == float(gt_elem))
            else:
                # we do not remove punct since comparisons can include punct
                comparisons.append(
                    normalize_str(ma_elem, remove_punct=False)
                    == normalize_str(gt_elem, remove_punct=False)
                )
        return all(comparisons), f"Evaluated {model_answer} as a comma separated list."

    # if gt is a str
    else:
        return (
            normalize_str(model_answer) == normalize_str(ground_truth),
            f"Evaluated {model_answer} as a string.",
        )


def normalize_str(input_str: str, remove_punct: bool = True) -> str:
    """
    Normalize a string by:
    - Removing all white spaces
    - Optionally removing punctuation (if remove_punct is True)
    - Converting to lowercase
    Parameters:
    - input_str: str, the string to normalize
    - remove_punct: bool, whether to remove punctuation (default: True)
    Returns:
    - str, the normalized string
    """  # noqa
    # Remove all white spaces. Required e.g for seagull vs. sea gull
    no_spaces = re.sub(r"\s", "", input_str)

    # Remove punctuation, if specified.
    if remove_punct:
        translator = str.maketrans("", "", string.punctuation)
        return no_spaces.lower().translate(translator)
    else:
        return no_spaces.lower()
