"""Utility functions related to the text-to-text task group."""

import logging
import typing as t

import evaluate
import numpy as np
from evaluate import EvaluationModule

from ..constants import METRIC_ATTRIBUTES_TAKING_UP_MEMORY
from ..data_models import BenchmarkConfig, DatasetConfig, GenerativeModelOutput
from ..exceptions import InvalidBenchmark
from ..utils import (
    HiddenPrints,
    clear_memory,
    raise_if_model_output_contains_nan_values,
)

if t.TYPE_CHECKING:
    from ..types import Labels, Predictions


logger = logging.getLogger("scandeval")


def compute_metrics(
    model_outputs_and_labels: tuple["Predictions", "Labels"],
    dataset_config: "DatasetConfig",
    benchmark_config: "BenchmarkConfig",
) -> dict[str, float]:
    """Compute the metrics needed for evaluation.

    Args:
        model_outputs_and_labels:
            The first sequence contains the model outputs and the second sequence
            contains the true labels.
        dataset_config:
            The configuration of the dataset.
        benchmark_config:
            The configuration of the benchmark.

    Returns:
        A dictionary with the names of the metrics as keys and the metric values as
        values.
    """
    model_outputs, labels = model_outputs_and_labels
    raise_if_model_output_contains_nan_values(model_output=model_outputs)

    metrics = {
        metric_cfg.name: (
            evaluate.load(
                path=metric_cfg.huggingface_id, cache_dir=benchmark_config.cache_dir
            )
            if metric_cfg.huggingface_id != ""
            else None
        )
        for metric_cfg in dataset_config.task.metrics
    }

    model_output_dtype = np.asarray(model_outputs).dtype
    output_is_prob = model_output_dtype in [np.float16, np.float32, np.float64]
    if output_is_prob:
        predictions = np.asarray(model_outputs).argmax(axis=-1)
    else:
        predictions = model_outputs

    results: dict[str, float] = dict()
    for cfg in dataset_config.task.metrics:
        metric = metrics[cfg.name]
        assert isinstance(metric, EvaluationModule)

        # Some metrics can be computed on hardware accelerators. In this case we
        # start by setting the device to the same device as the model
        if cfg.compute_kwargs.get("device", None) == "auto":
            cfg.compute_kwargs["device"] = benchmark_config.device.type

        while True:
            try:
                with HiddenPrints():
                    score_dict: dict[str, float] | None = metric.compute(
                        predictions=predictions, references=labels, **cfg.compute_kwargs
                    )

                # Clear the cache of the BERTScorer to avoid memory leaks
                for attribute in METRIC_ATTRIBUTES_TAKING_UP_MEMORY:
                    if hasattr(metric, attribute):
                        delattr(metric, attribute)

                clear_memory()
                break
            except Exception as e:
                # Clear the cache of the BERTScorer to avoid memory leaks
                if hasattr(metric, "cached_bertscorer"):
                    del metric.cached_bertscorer
                    clear_memory()

                oom_error = [
                    "CUDA out of memory",
                    "CUDA error",
                    "MPS backend out of memory",
                ]
                if not any(error in str(e) for error in oom_error):
                    raise InvalidBenchmark(str(e))

                if cfg.compute_kwargs.get("batch_size", 1) > 1:
                    batch_size = cfg.compute_kwargs["batch_size"]
                    cfg.compute_kwargs["batch_size"] = batch_size // 2
                    logger.debug(
                        "Out of memory error occurred during the computation of "
                        f"the metric {cfg.pretty_name}. Reducing the batch size to "
                        f"{cfg.compute_kwargs['batch_size']}."
                    )
                elif cfg.compute_kwargs.get("device", "cpu") != "cpu":
                    cfg.compute_kwargs["batch_size"] = 32
                    cfg.compute_kwargs["device"] = "cpu"
                    logger.debug(
                        "Out of memory error occurred during the computation of "
                        f"the metric {cfg.pretty_name}. Moving the computation to "
                        "the CPU."
                    )
                else:
                    raise InvalidBenchmark(str(e))

        # The metric returns None if we are running on multi-GPU and the current
        # process is not the main process
        if score_dict is not None:
            scores = score_dict[cfg.results_key]
            if isinstance(scores, list):
                scores = sum(scores) / len(scores)
            results[cfg.name] = scores

    return results


def extract_labels_from_generation(
    input_batch: dict[str, list], model_output: "GenerativeModelOutput"
) -> list[t.Any]:
    """Extract the predicted labels from the generated output.

    Args:
        input_batch:
            The input batch, where the keys are the feature names and the values
            are lists with the feature values.
        model_output:
            The raw generated output of the model.

    Returns:
        The predicted labels.
    """
    return model_output.sequences
