"""Formatters for dashboard display."""

from dataclasses import dataclass
from datetime import datetime
from typing import Literal

from entropice.utils.types import Grid, Model, TargetDataset, Task


@dataclass
class ModelDisplayInfo:
    keyword: Model
    short: str
    long: str

    def __str__(self) -> str:
        return f"{self.short} ({self.long})"


model_display_infos: dict[Model, ModelDisplayInfo] = {
    "espa": ModelDisplayInfo(
        keyword="espa",
        short="eSPA",
        long="entropy-optimal Scalable Probabilistic Approximations algorithm",
    ),
    "xgboost": ModelDisplayInfo(keyword="xgboost", short="XGBoost", long="Extreme Gradient Boosting"),
    "rf": ModelDisplayInfo(keyword="rf", short="Random Forest", long="Random Forest Classifier"),
    "knn": ModelDisplayInfo(keyword="knn", short="k-NN", long="k-Nearest Neighbors Classifier"),
}


@dataclass
class TaskDisplayInfo:
    keyword: Task
    display_name: str
    explanation: str


task_display_infos: dict[Task, TaskDisplayInfo] = {
    "binary": TaskDisplayInfo(
        keyword="binary",
        display_name="Binary Classification",
        explanation="Classify each grid cell as containing or not containing the target feature.",
    ),
    "count": TaskDisplayInfo(
        keyword="count",
        display_name="Count Prediction",
        explanation="Predict the number of target features present within each grid cell.",
    ),
    "density": TaskDisplayInfo(
        keyword="density",
        display_name="Density Estimation",
        explanation="Estimate the density of target features within each grid cell.",
    ),
}


@dataclass
class TrainingResultDisplayInfo:
    task: Task
    target: TargetDataset
    model: Model | Literal["autogluon"]
    grid: Grid
    level: int
    timestamp: datetime

    def get_display_name(self, format_type: Literal["task_first", "model_first"] = "task_first") -> str:
        task = self.task.capitalize()
        target = self.target.replace("_", " ").title()
        model = self.model.upper()
        grid = self.grid.capitalize()
        level = self.level
        timestamp = self.timestamp.strftime("%Y-%m-%d %H:%M")

        if format_type == "model_first":
            return f"{model} - {task}@{target} - {grid}-{level} ({timestamp})"
        else:  # task_first
            return f"{task}@{target} - {model} - {grid}-{level} ({timestamp})"


def format_metric_name(metric: str) -> str:
    """Format metric name for display.

    Args:
        metric: Raw metric name (e.g., 'f1_micro', 'precision_macro').

    Returns:
        Formatted metric name (e.g., 'F1 Micro', 'Precision Macro').

    """
    # Split by underscore and capitalize each part
    parts = metric.split("_")
    # Special handling for F1
    formatted_parts = []
    for part in parts:
        if part.lower() == "f1":
            formatted_parts.append("F1")
        else:
            formatted_parts.append(part.capitalize())
    return " ".join(formatted_parts)
