"""Shared types used across the entropice codebase."""

from dataclasses import dataclass
from typing import Literal

type Grid = Literal["hex", "healpix"]
type GridLevel = Literal[
    "hex3",
    "hex4",
    "hex5",
    "hex6",
    "healpix6",
    "healpix7",
    "healpix8",
    "healpix9",
    "healpix10",
]
type TargetDataset = Literal["darts_v1", "darts_mllabels"]
type L0SourceDataset = Literal["ArcticDEM", "ERA5", "AlphaEarth"]
type L2SourceDataset = Literal["ArcticDEM", "ERA5-shoulder", "ERA5-seasonal", "ERA5-yearly", "AlphaEarth"]
type Task = Literal["binary", "count_regimes", "density_regimes", "count", "density"]
# TODO: Consider implementing a "timeseries" and "event" temporal mode
type TemporalMode = Literal["feature", "synopsis", 2018, 2019, 2020, 2021, 2022, 2023]
type Model = Literal["espa", "xgboost", "rf", "knn", "autogluon"]
type Stage = Literal["train", "inference", "visualization"]
type HPSearch = Literal["random"]  # TODO Grid and Bayesian search
type Splitter = Literal[
    "kfold", "stratified", "shuffle", "stratified_shuffle", "group", "stratified_group", "shuffle_group"
]
type Scaler = Literal["standard", "minmax", "maxabs", "quantile", "none"]


@dataclass(frozen=True)
class GridConfig:
    """Grid configuration specification with metadata."""

    grid: Grid
    level: int
    id: GridLevel
    display_name: str
    res: Literal["sparse", "low", "medium"]
    sort_key: str

    @classmethod
    def from_grid_level(cls, grid_level: GridLevel | tuple[Literal["hex", "healpix"], int]) -> "GridConfig":
        """Create a GridConfig from a GridLevel string."""
        if isinstance(grid_level, str):
            if grid_level.startswith("hex"):
                grid = "hex"
                level = int(grid_level[3:])
            elif grid_level.startswith("healpix"):
                grid = "healpix"
                level = int(grid_level[7:])
            else:
                raise ValueError(f"Invalid grid level: {grid_level}")
        else:
            grid, level = grid_level
            grid_level: GridLevel = f"{grid}{level}"  # ty:ignore[invalid-assignment]
        display_name = f"{grid.capitalize()}-{level}"

        resmap: dict[str, Literal["sparse", "low", "medium"]] = {
            "hex3": "sparse",
            "hex4": "sparse",
            "hex5": "low",
            "hex6": "medium",
            "healpix6": "sparse",
            "healpix7": "sparse",
            "healpix8": "low",
            "healpix9": "low",
            "healpix10": "medium",
        }

        return cls(
            grid=grid,
            level=level,
            id=grid_level,
            display_name=display_name,
            res=resmap[grid_level],
            sort_key=f"{grid}_{level:02d}",
        )


# Note: get_args() doesn't work with Python 3.12+ type statement, so we define explicit lists
all_tasks: list[Task] = ["binary", "count_regimes", "density_regimes", "count", "density"]
all_temporal_modes: list[TemporalMode] = ["feature", "synopsis", 2018, 2019, 2020, 2021, 2022, 2023]
all_target_datasets: list[TargetDataset] = ["darts_v1", "darts_mllabels"]
all_l2_source_datasets: list[L2SourceDataset] = [
    "ArcticDEM",
    "ERA5-shoulder",
    "ERA5-seasonal",
    "ERA5-yearly",
    "AlphaEarth",
]

grid_configs: list[GridConfig] = [
    GridConfig.from_grid_level("hex3"),
    GridConfig.from_grid_level("hex4"),
    GridConfig.from_grid_level("hex5"),
    GridConfig.from_grid_level("hex6"),
    GridConfig.from_grid_level("healpix6"),
    GridConfig.from_grid_level("healpix7"),
    GridConfig.from_grid_level("healpix8"),
    GridConfig.from_grid_level("healpix9"),
    GridConfig.from_grid_level("healpix10"),
]
