# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Metric to calculate the accuracy for the TRAM benchmark by Wang et al. (2024)."""

import re
from typing import TypedDict

import datasets
import evaluate

VALID_ANSWER_CHOICES = frozenset({"A", "B", "C", "D"})

TRAM_ANSWER_PATTERN = r"[Tt]he final answer is \(([A-D])\)"


class AccuracyResult(TypedDict):
    accuracy: float | list[int]


_CITATION = """\
@InProceedings{auss:tram_accuracy,
title = {TRAM Accuracy},
authors={Auss Abbood},
year={2025}
}
"""

_DESCRIPTION = """\
Accuracy metric for the (multiple choice) TRAM datasets by Wang et al. (2024).
"""


_KWARGS_DESCRIPTION = """
Calculates the accuracy for the TRAM datasets by extracting the final answer from the prediction and comparing it to the reference answer.
Args:
    predictions: list of predictions to score. Each prediction
        should be a string with the model's response, which contains the final answer.
    references: list of reference for each prediction. Each
        reference a single letter representing the correct answer.
    return_average: whether to return the average accuracy or the accuracy for each prediction.
Returns:
    accuracy: the accuracy for the TRAM datasets.
"""


TRAM_ANSWER_REGEX = re.compile(TRAM_ANSWER_PATTERN)


@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class TRAMAccuracy(evaluate.Metric):
    """Calculates the accuracy for the (multiple choice) TRAM datasets by extracting the final answer from the prediction and comparing it to the reference answer."""

    def _info(self) -> evaluate.MetricInfo:
        return evaluate.MetricInfo(
            module_type="metric",
            description=_DESCRIPTION,
            citation=_CITATION,
            inputs_description=_KWARGS_DESCRIPTION,
            # This defines the format of each prediction and reference
            features=datasets.Features(
                {
                    "predictions": datasets.Value("string"),
                    "references": datasets.Value("string"),
                }
            ),
            homepage="https://huggingface.co/spaces/aauss/tram_accuracy",
            codebase_urls=[
                "https://huggingface.co/spaces/aauss/tram_accuracy/tree/main"
            ],
            reference_urls=["https://huggingface.co/datasets/Warrieryes/TRAM-Temporal"],
        )

    def _compute(
        self,
        predictions: list[str],
        references: list[str],
        return_average: bool = True,
    ) -> AccuracyResult:
        """Returns the accuracy for the (multiple choice) TRAM datasets."""
        if not predictions:
            raise ValueError("predictions cannot be empty")
        if len(predictions) != len(references):
            raise ValueError(
                f"predictions and references must have same length, "
                f"got {len(predictions)} and {len(references)}"
            )
        predictions_matches = [
            TRAM_ANSWER_REGEX.search(prediction) for prediction in predictions
        ]
        predictions_extracted = [
            match.group(1) if match is not None else None
            for match in predictions_matches
        ]
        accuracy = [
            1 if response == label else 0
            for response, label in zip(predictions_extracted, references)
        ]
        if return_average:
            return {"accuracy": sum(accuracy) / len(accuracy)}
        return {"accuracy": accuracy}