"""Metrics for model evaluation and hyperparameter optimization."""

from functools import partial

from sklearn.metrics import (
    accuracy_score,
    f1_score,
    jaccard_score,
    mean_absolute_error,
    mean_squared_error,
    precision_score,
    r2_score,
    recall_score,
)

from entropice.utils.types import Task


def get_metrics(task: Task) -> tuple[list[str], str]:
    """Get the list of metrics for a given task."""
    if task == "binary":
        return ["accuracy", "recall", "precision", "f1", "jaccard"], "f1"
    elif task in ["count_regimes", "density_regimes"]:
        return [
            "accuracy",  # equals "f1_micro", "precision_micro", "recall_micro", "recall_weighted"
            "f1_macro",
            "f1_weighted",
            "precision_macro",
            "precision_weighted",
            "recall_macro",
            "jaccard_micro",
            "jaccard_macro",
            "jaccard_weighted",
        ], "f1_weighted"
    else:
        return [
            "neg_mean_squared_error",
            "neg_mean_absolute_error",
            "r2",
        ], "r2"


# Compute other metrics - using predictions directly instead of re-predicting for each metric
# Use functools.partial for cleaner metric definitions with non-default parameters
metric_functions = {
    "accuracy": accuracy_score,
    "recall": recall_score,
    "precision": precision_score,
    "f1": f1_score,
    "jaccard": jaccard_score,
    "recall_macro": partial(recall_score, average="macro"),
    "recall_weighted": partial(recall_score, average="weighted"),
    "precision_macro": partial(precision_score, average="macro"),
    "precision_weighted": partial(precision_score, average="weighted"),
    "f1_macro": partial(f1_score, average="macro"),
    "f1_weighted": partial(f1_score, average="weighted"),
    "jaccard_micro": partial(jaccard_score, average="micro"),
    "jaccard_macro": partial(jaccard_score, average="macro"),
    "jaccard_weighted": partial(jaccard_score, average="weighted"),
    "neg_mean_squared_error": mean_squared_error,
    "neg_mean_absolute_error": mean_absolute_error,
    "r2": r2_score,
}
