# MIT License

# Copyright (c) 2024 The HuggingFace Team

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import logging
from typing import Callable, Literal, Sequence

import numpy as np

from lighteval.metrics.metrics_sample import (
    ExactMatches,
    F1_score,
    LoglikelihoodAcc,
    NormalizedMultiChoiceProbability,
    Probability,
)
from lighteval.metrics.normalizations import (
    LogProbNormalization,
    LogProbTokenNorm,
    get_multilingual_normalizer,
)
from lighteval.metrics.utils.extractive_match_utils import (  # noqa: F401
    ExprExtractionConfig,
    ExtractionTarget,
    LatexExtractionConfig,
    extract_target_from_pred,
    get_extraction_regexes,
)
from lighteval.metrics.utils.math_comparison import compare_gold_target
from lighteval.metrics.utils.metric_utils import SampleLevelComputation, SampleLevelMetric
from lighteval.models.model_output import ModelResponse
from lighteval.tasks.requests import Doc, SamplingMethod
from lighteval.utils.language import Language
from lighteval.utils.timeout import timeout


logger = logging.getLogger(__name__)


class LogLikelihoodAccMetric(SampleLevelMetric):
    def __init__(self, normalization: LogProbNormalization | None = None):
        """Creates an accuracy (loglikelihood) metric, which returns accuracy given normalization."""
        super().__init__(
            metric_name="acc" + (f"_{normalization.name}" if normalization else ""),
            sample_level_fn=LoglikelihoodAcc(logprob_normalization=normalization),
            category=SamplingMethod.LOGPROBS,
            corpus_level_fn=np.mean,
            higher_is_better=True,
        )


class NormalizedMultiChoiceProbMetric(SampleLevelMetric):
    def __init__(
        self,
        normalization: LogProbNormalization | None = None,
        aggregation_function: Callable[[np.ndarray], float] = np.max,
    ):
        """Creates a normalized multi-choice probability metric, which returns the probability of the gold choice / sum of probabilities of all choices (after logprobs are normalized)."""
        super().__init__(
            metric_name="normalized_mc_prob" + (f"_{normalization.name}" if normalization else ""),
            sample_level_fn=NormalizedMultiChoiceProbability(
                log_prob_normalization=normalization, aggregation_function=aggregation_function
            ),
            category=SamplingMethod.LOGPROBS,
            corpus_level_fn=np.mean,
            higher_is_better=True,
        )


class ProbabilityMetric(SampleLevelMetric):
    def __init__(
        self,
        normalization: LogProbTokenNorm | None = None,
        aggregation_function: Callable[[np.ndarray], float] = np.max,
    ):
        """Creates a probability metric, which returns the probability of the gold choice given normalization."""
        super().__init__(
            metric_name="prob" + (f"_{normalization.name}" if normalization else ""),
            sample_level_fn=Probability(normalization=normalization, aggregation_function=aggregation_function),
            category=SamplingMethod.LOGPROBS,
            corpus_level_fn=np.mean,
            higher_is_better=True,
        )


class MultilingualQuasiF1ScoreMetric(SampleLevelMetric):
    def __init__(self, language: Language, aggregation_function: Callable[[list[float]], float] = max):
        """Creates a language-aware F1 score metric, which returns the F1 score.

        Args:
            language: The language of the samples.
            aggregation_function: Aggregation samples to use when multiple golds are present.
        """
        super().__init__(
            metric_name=f"f1_{language.value}",
            sample_level_fn=F1_score(
                normalize_gold=get_multilingual_normalizer(language),
                normalize_pred=get_multilingual_normalizer(language),
                aggregation_function=aggregation_function,
            ),
            category=SamplingMethod.GENERATIVE,
            corpus_level_fn=np.mean,
            higher_is_better=True,
        )


class MultilingualQuasiExactMatchMetric(SampleLevelMetric):
    def __init__(
        self,
        language: Language,
        match_type: Literal["prefix", "suffix", "full"] = "full",
        aggregation_function: Callable[[list[float]], float] = max,
    ):
        """Creates a language-aware exact match metric, which returns the exact match score
        Args:
            language: The language of the samples.
            match_type: The type of match to use
                - "prefix": Prefixes must match
                - "suffix": Suffixes must match
                - "full": Full strings must match
            aggregation_function: Aggregation samples to use when multiple golds are present.
        """
        super().__init__(
            metric_name=f"exact_match_{language.value}_{match_type}",
            sample_level_fn=ExactMatches(
                normalize_gold=get_multilingual_normalizer(language),
                normalize_pred=get_multilingual_normalizer(language),
                aggregation_function=aggregation_function,
                type_exact_match=match_type,
            ),
            category=SamplingMethod.GENERATIVE,
            corpus_level_fn=np.mean,
            higher_is_better=True,
        )


class MultilingualExtractiveMatchMetric(SampleLevelComputation):
    def __init__(
        self,
        language: Language = Language.ENGLISH,
        gold_extraction_target: Sequence[ExtractionTarget] = (ExprExtractionConfig(),),
        pred_extraction_target: Sequence[ExtractionTarget] = (ExprExtractionConfig(), LatexExtractionConfig()),
        aggregation_function: Callable[[list[float]], float] = max,
        fallback_mode: Literal["no_fallback", "first_match"] = "first_match",
        extraction_mode: Literal["first_match", "any_match"] = "any_match",
        precision: int = 6,
        timeout_seconds: int = 5,
    ):
        """Creates a language-aware extractive match metric that extracts answers from the model's output.

        Known issues:
        - If the task is to simplify an expression, the metric might overestimate the accuracy. This is because if the model doesn't output any anchor for the extraction (e.g final answer is..),
            it's possible that the extracted prediction will be the expression to simplify. Because we do simplifications ourselves, it can thus happen that sympy will correctly simplify the expression,
            thus it will match gold, despite model not doing anything. PRs to fix this are welcome.

        - There is currently no StringExtractionConfig, so if the gold is \boxed{\text{Friday}} and model outputs Friday it will not match, because nothing will be extracted.

        Args:
            language: Language
                The language of the samples.
            gold_extraction_target: Sequence[ExtractionTarget]
                Extraction targets to use for gold answers. Defaults to extracting simple math expressions.
            pred_extraction_target: Sequence[ExtractionTarget]
                Extraction targets to use for predictions. Defaults to extracting simple math expressions.
            aggregation_function: Callable[[list[float]], float]
                Function to aggregate scores when multiple golds/predictions are present. Defaults to max.
            fallback_mode: Literal["no_fallback", "first_match"]
                How to perform extraction. Defaults to "first_match".
                - "no_fallback": Only use first successfully parsed matches
                - "first_match": Use the first successfully parsed match + first match irregardless the parsing success
            extraction_mode: Literal["first_match", "any_match"]
                - "first_match": Only tries to extract the first regex match if it fails no other matches are tried
                - "any_match": Tries to extract any regex match

            precision: int
                Number of decimal places to use when comparing numerical values. Defaults to 6.
            timeout_seconds: int
                Timeout for the extraction (each attempt) and comparison. Defaults to 5.

        """
        self.language = language
        self.gold_extraction_target = gold_extraction_target
        self.pred_extraction_target = pred_extraction_target
        self.aggregation_function = aggregation_function
        self.fallback_mode = fallback_mode
        self.extraction_mode = extraction_mode
        self.precision = precision
        self.timeout_seconds = timeout_seconds

    @timeout(2)
    def add_to_specifics_with_timeout(
        self, formatted_doc: Doc, extracted_predictions: list[list[str]], extracted_golds: list[list[str]]
    ) -> None:
        if formatted_doc.specific is None:
            formatted_doc.specific = {}

        formatted_doc.specific["extracted_predictions"] = [
            str(pred) for preds in extracted_predictions for pred in preds
        ]
        formatted_doc.specific["extracted_golds"] = [str(gold) for golds in extracted_golds for gold in golds]

    def compute(self, doc: Doc, model_response: ModelResponse) -> float:
        golds = doc.get_golds()
        predictions = model_response.final_text

        gold_extraction_regexes = get_extraction_regexes(doc, self.gold_extraction_target, self.language)
        pred_extraction_regexes = get_extraction_regexes(doc, self.pred_extraction_target, self.language)

        extracted_predictions = [
            extract_target_from_pred(
                pred, pred_extraction_regexes, self.fallback_mode, self.extraction_mode, self.timeout_seconds
            )
            for pred in predictions
        ]
        extracted_golds = [
            extract_target_from_pred(
                gold, gold_extraction_regexes, self.fallback_mode, self.extraction_mode, self.timeout_seconds
            )
            for gold in golds
        ]

        # Assert on empty gold and warn on empty pred
        if any(len(g) == 0 for g in extracted_golds):
            logger.warning(f"We did not manage to extract a gold in the correct format. Gold: {golds}")
            extracted_golds = [[gold] for gold in golds]

        if all(len(p) == 0 for p in extracted_predictions):
            logger.warning(
                f"We did not manage to extract a prediction in the correct format. Gold: {golds}, Pred: {predictions}"
            )

        # We have to use timeout because the sypmy to str conversion can be very slow
        try:
            self.add_to_specifics_with_timeout(doc, extracted_predictions, extracted_golds)
        except TimeoutError:  # noqa: E722
            logger.warning("Timeout when adding extracted predictions and golds to specific")

        return self.aggregation_function(
            [
                (
                    1.0
                    if any(
                        compare_gold_target(gold, pred, self.precision, timeout_seconds=self.timeout_seconds)
                        for gold in extracted_golds
                    )
                    else 0.0
                )
                for pred in extracted_predictions
            ]
        )
