# ruff: noqa: N806, D105
"""Training dataset preparation and model training.

Naming conventions:

- Ensemble Dataset
  - Member-Datasets: selection based on config
- L2-Datasets: ready to use XDGGS Xarray Datasets
  - Currently implemented: embeddings (alpha-earth), ERA5-seasonal, ERA5-shoulder, ERA5-yearly, Arcticdem32m
  - All L2 Datasets can be multidimensional, but all having at least the cell_ids dimension
  - All L2 Datasets have at least one data variable
  - Dimensions of L2 Datasets are e.g. time or aggregation
"""

import hashlib
import json
from collections.abc import Generator
from dataclasses import asdict, dataclass, field
from functools import cache, cached_property
from itertools import product
from typing import Literal, TypeVar, cast

import cupy as cp
import cyclopts
import geopandas as gpd
import numpy as np
import pandas as pd
import seaborn as sns
import torch
import xarray as xr
from cuml import KMeans
from rich import pretty, traceback
from sklearn import set_config
from sklearn.model_selection import train_test_split
from stopuhr import stopwatch

import entropice.spatial.grids
import entropice.utils.paths
from entropice.ingest.era5 import unstack_era5_time
from entropice.utils.types import Grid, L2SourceDataset, TargetDataset, Task, TemporalMode

traceback.install()
pretty.install()

set_config(array_api_dispatch=True)

sns.set_theme("talk", "whitegrid")


def _collapse_to_dataframe(ds: xr.Dataset | xr.DataArray) -> pd.DataFrame:
    collapsed = ds.to_dataframe()
    # Make a dummy row to avoid empty dataframe issues
    use_dummy = collapsed.shape[0] == 0
    if use_dummy:
        collapsed.loc[tuple(range(len(collapsed.index.names)))] = np.nan
    cols = cast(list[str], list(collapsed.index.names))
    pivcols = sorted(set(cols) - {"cell_ids"})
    collapsed = collapsed.pivot_table(index="cell_ids", columns=pivcols)
    collapsed.columns = ["_".join(map(str, v)) for v in collapsed.columns]
    if use_dummy:
        collapsed = collapsed.dropna(how="all")
    expected_cols = _get_expected_collapsed_columns(ds)
    missing_cols = set(expected_cols) - set(collapsed.columns)
    if missing_cols:
        raise ValueError(
            f"Collapsed dataframe is missing expected columns: {missing_cols=} {collapsed.columns=} {expected_cols=}"
        )
    return collapsed


def _get_expected_collapsed_columns(ds: xr.Dataset | xr.DataArray) -> list[str]:
    dims = sorted(set(ds.dims) - {"cell_ids"})
    dims_product = list(product(*[ds.coords[dim].to_numpy() for dim in dims]))
    expected_cols = []
    if isinstance(ds, xr.Dataset):
        variables = list(ds.data_vars)
        for var in variables:
            for dims_values in dims_product:
                agg = "_".join(dims_values)
                expected_cols.append(f"{var}_{agg}")
    else:
        assert ds.name is not None, "DataArray must have a name to determine expected columns"
        for dims_values in dims_product:
            agg = "_".join(dims_values)
            expected_cols.append(f"{ds.name}_{agg}")

    return expected_cols


def _cell_ids_hash(cell_ids: pd.Series) -> str:
    sorted_ids = np.sort(cell_ids.to_numpy())
    return hashlib.blake2b(sorted_ids.tobytes(), digest_size=8).hexdigest()


def bin_values(
    values: pd.Series,
    task: Literal["count_regimes", "density_regimes"],
    none_val: float = 0,
) -> pd.Series:
    """Bin values into predefined intervals for different tasks.

    First, a 'none' bin is created for values equal to `none_val` usually 0.
    Then, the remaining values are binned automatically into 5 bins, all containing roughly the same number of samples.

    Args:
        values (pd.Series): Pandas Series of numerical values to bin.
        task (Literal["count_regimes", "density_regimes"]): Task type - 'count_regimes' or 'density_regimes'.
        none_val (float, optional): Value representing 'none' or 'empty' (e.g., 0 for count). Defaults to 0.

    Returns:
        pd.Series: Pandas Series of ordered categorical binned values.

    Raises:
        ValueError: If an value is NaN.

    """
    labels_dict = {
        "count_regimes": ["None", "Very Few", "Few", "Several", "Many", "Very Many"],
        "density_regimes": ["Empty", "Very Sparse", "Sparse", "Moderate", "Dense", "Very Dense"],
    }
    labels = labels_dict[task]

    if values.isna().any():
        raise ValueError("Values contain NaN")

    # Separate none values from others
    none_mask = values == none_val
    non_none_values = values[~none_mask]

    assert len(non_none_values) > 5, "Not enough non-none values to create bins."

    # Create bins without labels first to handle duplicates
    binned_non_none = pd.qcut(non_none_values, q=5, labels=False, duplicates="drop")

    # Map the bin indices to labels, using only as many labels as bins created
    n_bins = binned_non_none.max() + 1  # bins are 0-indexed
    active_labels = labels[1 : 1 + n_bins]  # Skip "None" and take n_bins labels
    binned_non_none = binned_non_none.map(dict(enumerate(active_labels))).astype("category")
    binned_non_none = binned_non_none.cat.set_categories(labels, ordered=True)

    binned = pd.Series(index=values.index, dtype="category")
    binned = binned.cat.set_categories(labels, ordered=True)
    binned.update(binned_non_none)
    binned.loc[none_mask] = labels[0]
    return binned


ArrayType = TypeVar("ArrayType", torch.Tensor, np.ndarray, cp.ndarray)


@dataclass(frozen=True, eq=False)
class SplittedArrays[ArrayType: (torch.Tensor, np.ndarray, cp.ndarray)]:
    """Small wrapper for train and test arrays."""

    train: ArrayType
    test: ArrayType

    @cached_property
    def complete(self) -> ArrayType:
        """Complete train and test arrays."""
        if isinstance(self.train, torch.Tensor) and isinstance(self.test, torch.Tensor):
            return torch.cat([self.train, self.test], dim=0)  # ty:ignore[invalid-return-type]
        elif isinstance(self.train, cp.ndarray) and isinstance(self.test, cp.ndarray):
            return cp.concatenate([self.train, self.test], axis=0)
        elif isinstance(self.train, np.ndarray) and isinstance(self.test, np.ndarray):
            return np.concatenate([self.train, self.test], axis=0)  # ty:ignore[invalid-return-type]
        else:
            raise TypeError("Incompatible types for train and test arrays.")

    def as_numpy(self) -> "SplittedArrays[np.ndarray]":
        """Return the arrays as numpy arrays."""
        train_np = (
            self.train.cpu().numpy()
            if isinstance(self.train, torch.Tensor)
            else self.train.get()
            if isinstance(self.train, cp.ndarray)
            else self.train
        )
        test_np = (
            self.test.cpu().numpy()
            if isinstance(self.test, torch.Tensor)
            else self.test.get()
            if isinstance(self.test, cp.ndarray)
            else self.test
        )
        return SplittedArrays(train=train_np, test=test_np)


@dataclass(frozen=True, eq=False)
class TrainingSet[ArrayType: (torch.Tensor, np.ndarray, cp.ndarray)]:
    """Container for the training dataset."""

    targets: gpd.GeoDataFrame
    features: pd.DataFrame
    X: SplittedArrays[ArrayType]
    y: SplittedArrays[ArrayType]
    z: pd.Series
    split: pd.Series

    @cached_property
    def target_intervals(self) -> list[tuple[float, float] | tuple[int, int] | tuple[None, None]]:
        """Calculate the intervals for each target category.

        Returns:
            list[tuple[float, float] | tuple[int, int] | tuple[None, None]]:
                List of (min, max) tuples for each category.

        """
        binned = self.targets["y"]
        raw = self.targets["z"]
        assert binned.dtype.name == "category", "Target labels are not categorical."

        # For each category get the min and max values from raw_values
        intervals: list[tuple[float, float] | tuple[int, int] | tuple[None, None]] = []
        for category in binned.cat.categories:
            category_mask = binned == category
            if category_mask.sum() == 0:
                intervals.append((None, None))
            else:
                category_raw_values = raw[category_mask]
                intervals.append((category_raw_values.min(), category_raw_values.max()))
        return intervals

    @property
    def target_labels(self) -> list[str] | None:
        """Labels of the target categories, None if not categorical."""
        binned = self.targets["y"]
        if binned.dtype.name == "category":
            return list(binned.cat.categories)
        else:
            return None

    @property
    def target_codes(self) -> list[int] | None:
        """Label-Codes of the target categories, None if not categorical."""
        binned = self.targets["y"]
        if binned.dtype.name == "category":
            return list(range(len(binned.cat.categories)))
        else:
            return None

    @property
    def feature_names(self) -> list[str]:
        """Names of the features."""
        return self.features.columns.tolist()

    @cached_property
    def clusters(self) -> pd.Series:
        """Geo-Cluster assignments for each sample, based on the geometries of the target GeoDataFrame."""
        centroids = self.targets.to_crs("EPSG:3413")["geometry"].centroid
        centroids = cp.array([centroids.x.to_numpy(), centroids.y.to_numpy()]).T
        # Use kMeans to cluster the centroids into 10 clusters
        kmeans = KMeans(n_clusters=10, random_state=42)
        clusters = kmeans.fit_predict(centroids).get()
        return pd.Series(clusters, index=self.targets.index)

    def __len__(self):
        return len(self.z)

    def to_dataframe(self, split: Literal["train", "test"] | None) -> pd.DataFrame:
        """Create a simple DataFrame for training with e.g. autogluon.

        This function creates a DataFrame with features and target labels (called "y") for training.
        The resulting dataframe contains no geometry information and drops all samples with NaN values.

        Args:
            task (Task): The task.
            split (Literal["train", "test"] | None): If specified, only return the samples for the given split.

        Returns:
            pd.DataFrame: The training DataFrame.

        """
        dataset = self.targets[["y"]].rename(columns={"y": "label"}).join(self.features)
        if split is not None:
            dataset = dataset[self.split == split]
        assert len(dataset) > 0, "No valid samples found after joining features and targets."
        return dataset


@cyclopts.Parameter("*")
@dataclass(frozen=True, kw_only=True)
class DatasetEnsemble:
    """An ensemble of datasets for training and inference."""

    grid: Grid
    level: int
    members: list[L2SourceDataset] = field(
        default_factory=lambda: [
            "AlphaEarth",
            "ArcticDEM",
            "ERA5-yearly",
            "ERA5-seasonal",
            "ERA5-shoulder",
        ]
    )
    temporal_mode: TemporalMode = "synopsis"
    # ?: We can't use L2SourceDataset as types here because cyclopts can't handle Literals as dict keys
    dimension_filters: dict[str, dict[str, list]] = field(default_factory=dict)
    variable_filters: dict[str, list[str]] = field(default_factory=dict)

    def __post_init__(self):
        # Validate filters
        for member in self.members:
            if member in self.variable_filters:
                # This enforces the assumption that the read members always are xarray Datasets and not DataArrays
                assert isinstance(self.variable_filters[member], list) and len(self.variable_filters[member]) >= 1, (
                    f"Invalid variable filter for {member=}: {self.variable_filters[member]}"
                    " Variable filter values must be a list with one or more entries."
                )
            if member in self.dimension_filters:
                for dim, values in self.dimension_filters[member].items():
                    # This enforces the assumption that we do the temporal filtering via the temporal_mode
                    assert dim not in ["year", "month", "time"], (
                        f"Invalid dimension filter for {member=}: {dim}"
                        " Filtering on 'year', 'month' or 'time' is not supported."
                    )
                    # This enforces the assumption that there are no empty dimensions in the Dataset
                    assert isinstance(values, list) and len(values) >= 1, (
                        f"Invalid dimension filter for {dim=}: {values}"
                        " Dimension filter values must be a list with one or more entries."
                    )
        if "Grid" in self.variable_filters.keys():
            filtered = set(self.variable_filters["Grid"])
            valid = {"x", "y", "cell_area", "land_area", "water_area", "land_ratio"}
            assert len(filtered - valid) == 0

    def __hash__(self):
        return int(self.id(), 16)

    def id(self):
        """Return an unique, stable identifier based on the settings of this class."""
        return hashlib.blake2b(
            json.dumps(asdict(self), sort_keys=True).encode("utf-8"),
            digest_size=16,
        ).hexdigest()

    @cached_property
    def cell_ids(self) -> pd.Series:
        """Series of all cell-ids of the grid."""
        return self.read_grid().index.to_series()

    @cached_property
    def geometries(self) -> pd.Series:
        """Series of all geometries of the grid."""
        return self.read_grid()["geometry"]

    @stopwatch("Reading grid")
    @cache
    def read_grid(self) -> gpd.GeoDataFrame:
        """Load the grid dataframe and enrich it with lat-lon information."""
        if "Grid" in self.variable_filters:
            columns_to_load = self.variable_filters["Grid"] + ["cell_id", "geometry"]
        else:
            columns_to_load = ["cell_id", "geometry", "cell_area", "land_area", "water_area", "land_ratio", "x", "y"]

        # Reading the data takes for the largest grids ~1.7s
        gridfile = entropice.utils.paths.get_grid_file(self.grid, self.level)
        grid_gdf = gpd.read_parquet(gridfile, columns=columns_to_load)

        # Convert hex cell_id to int (for the largest grids ~0.4s)
        if self.grid == "hex":
            grid_gdf["cell_id"] = grid_gdf["cell_id"].apply(lambda x: int(x, 16)).astype(np.uint64)

        grid_gdf = grid_gdf.set_index("cell_id")
        return grid_gdf

    @stopwatch.f("Getting target labels", print_kwargs=["task"])
    def get_targets(self, target: TargetDataset, task: Task) -> gpd.GeoDataFrame:
        """Create a training target labels for a specific target and task.

        The function reads the target dataset, filters it based on coverage,
        and prepares the target labels according to the specified task.

        This is quite a fast function, with the bottleneck usually being the reading of the zarr store
        (milliseconds to few seconds).

        This is intended to be used by another training Dataset creation function
        and for easy visualization of the targets.


        Args:
            target (TargetDataset): The target dataset to use.
            task (Task): The task.

        Returns:
            gpd.GeoDataFrame: GeoDataFrame with target labels, geometry and raw-target values,
                all indexed by the cell_ids.
                Columns: ["geometry", "y", "z"], where "y" is the target label and "z" the raw value.
                For categorical tasks, "y" is categorical and "z" is numerical.
                For regression tasks, both "y" and "z" are numerical and identical.

        """
        match (target, self.temporal_mode):
            case ("darts_v1" | "darts_v2", "feature" | "synopsis"):
                version: Literal["v1-l3", "v2-l3"] = target.split("_")[1] + "-l3"  # ty:ignore[invalid-assignment]
                target_store = entropice.utils.paths.get_darts_file(grid=self.grid, level=self.level, version=version)
                targets = xr.open_zarr(target_store, consolidated=False)
            case ("darts_v1" | "darts_v2", int()):
                version: Literal["v1", "v2"] = target.split("_")[1]  # ty:ignore[invalid-assignment]
                target_store = entropice.utils.paths.get_darts_file(grid=self.grid, level=self.level, version=version)
                targets = xr.open_zarr(target_store, consolidated=False)
                targets = targets.sel(year=self.temporal_mode)
            case ("darts_mllabels", str()):  # Years are not supported
                target_store = entropice.utils.paths.get_darts_file(
                    grid=self.grid, level=self.level, version="mllabels"
                )
                targets = xr.open_zarr(target_store, consolidated=False)
            case _:
                raise NotImplementedError(f"Target {target} on {self.temporal_mode} mode not supported.")
        targets = cast(xr.Dataset, targets)
        covered_cell_ids = targets["coverage"].where(targets["coverage"] > 0).dropna("cell_ids")["cell_ids"].to_series()
        targets = targets.sel(cell_ids=covered_cell_ids.to_numpy())
        match task:
            case "binary":
                z = targets["count"].to_series()
                y = (z > 0).astype("category").map({False: "No RTS", True: "RTS"})
            case "count_regimes":
                z = targets["count"].to_series()
                y = bin_values(z, task="count_regimes", none_val=0)
            case "density_regimes":
                z = targets["density"].to_series()
                y = bin_values(z, task="density_regimes", none_val=0)
            case "count":
                z = targets["count"].to_series()
                y = z
            case "density":
                z = targets["density"].to_series()
                y = z
            case _:
                raise NotImplementedError(f"Task {task} not supported.")
        cell_ids = targets["cell_ids"].to_series()
        geometries = self.geometries.loc[cell_ids]
        # ! Warning: For some stupid unknown reason, it is not possible to enforce the uint64 dtype on the index
        # This will result in joining issues later on if the dtypes do not match!
        return gpd.GeoDataFrame(
            {
                "cell_id": cell_ids,
                "geometry": geometries,
                "y": y,
                "z": z,
            }
        ).set_index("cell_id")

    @stopwatch.f("Reading member", print_kwargs=["member", "lazy"])
    def read_member(self, member: L2SourceDataset, cell_ids: pd.Series | None = None, lazy: bool = False) -> xr.Dataset:  # noqa: C901
        """Read a single member (source) of the Ensemble and applies filters based on the ensemble configuration.

        When `lazy` is False, the data is actually read into memory, which can take some time (seconds).

        Args:
            member (L2SourceDataset): The member
            cell_ids (pd.Series | None, optional): The cell IDs to read. If None, all cell IDs are read.
                Defaults to None.
            lazy (bool, optional): Whether to not load the data and instead return a lazy dataset. Defaults to False.

        Returns:
            xr.Dataset: Xarray Dataset with the member data for the given cell IDs.

        """
        match member:
            case "AlphaEarth":
                store = entropice.utils.paths.get_embeddings_store(
                    grid=self.grid, level=self.level, temporal=self.temporal_mode
                )
            case "ERA5-yearly" | "ERA5-seasonal" | "ERA5-shoulder":
                era5_agg: Literal["yearly", "seasonal", "shoulder"] = member.split("-")[1]  # ty:ignore[invalid-assignment]
                store = entropice.utils.paths.get_era5_stores(
                    era5_agg, grid=self.grid, level=self.level, temporal=self.temporal_mode
                )
            case "ArcticDEM":
                store = entropice.utils.paths.get_arcticdem_stores(grid=self.grid, level=self.level)
            case _:
                raise NotImplementedError(f"Member {member} not implemented.")
        ds = xr.open_zarr(store, consolidated=False)

        # Apply variable and dimension filters
        if member in self.variable_filters:
            ds = ds[self.variable_filters[member]]
        if member in self.dimension_filters:
            ds = ds.sel(self.dimension_filters[member])
        # Delete all coordinates which are not in the dimension
        ds = ds.drop_vars([coord for coord in ds.coords if coord not in ds.dims])

        # Only load target cell ids
        if cell_ids is None:
            cell_ids = self.cell_ids
        is_intersecting_cell_ids = np.isin(ds["cell_ids"].values, cell_ids.to_numpy(), assume_unique=True)
        if not is_intersecting_cell_ids.all():
            ds = ds[{"cell_ids": is_intersecting_cell_ids}]

        # Unstack era5 data if needed (already done for synopsis mode)
        if member.startswith("ERA5") and self.temporal_mode != "synopsis":
            era5_agg: Literal["yearly", "seasonal", "shoulder"] = member.split("-")[1]  # ty:ignore[invalid-assignment]
            ds = unstack_era5_time(ds, era5_agg)

        # Apply the temporal mode
        if isinstance(self.temporal_mode, int) and member != "ArcticDEM":
            ds = ds.sel(year=self.temporal_mode, drop=True)

        # Actually read data into memory
        if not lazy:
            ds.load()
        return ds

    def make_features(
        self,
        cell_ids: pd.Series | None = None,
        cache_mode: Literal["none", "read", "overwrite"] = "none",
    ) -> pd.DataFrame:
        """Create a feature DataFrame for the given temporal task and cell IDs.

        This function reads all members of the ensemble, prepares their features based on the temporal task,
        and combines them into a single DataFrame (no geometry).
        It is quite computation intensive (seconds to minutes), depending on the configuration and number of cell IDs.
        To speed up repeated calls, a caching mechanism is implemented.

        The indented use for this function is solely the creation of training and inference datasets.
        For visualization purposes it is recommended to use the `read_member` with `lazy=True` functions directly.

        Args:
            cell_ids (pd.Series | None, optional): The cell IDs to read. If None, all cell IDs are read.
                Defaults to None.
            cache_mode (Literal["none", "read", "overwrite"], optional): Caching mode.
                "none": No caching.
                "read": Read from cache if exists, otherwise create and save to cache.
                "overwrite": Always create and save to cache, overwriting existing cache.
                Defaults to "none".

        Returns:
            pd.DataFrame: DataFrame with features for the given cell IDs.

        """
        if cell_ids is None:
            cell_ids = self.cell_ids

        # Caching mechanism
        if cache_mode != "none":
            cells_hash = _cell_ids_hash(cell_ids)
            cache_file = entropice.utils.paths.get_features_cache(ensemble_id=self.id(), cells_hash=cells_hash)
            if cache_mode == "read" and cache_file.exists():
                with stopwatch("Loading features from cache"):
                    return pd.read_parquet(cache_file)

        # Create features
        dataset = self.read_grid().loc[cell_ids].drop(columns=["geometry"])
        member_dfs = []
        for member in self.members:
            match member.split("-"):
                case ["ERA5", era5_agg]:
                    member_dfs.append(self._prep_era5(cell_ids, era5_agg))
                case ["AlphaEarth"]:
                    member_dfs.append(self._prep_embeddings(cell_ids))
                case ["ArcticDEM"]:
                    member_dfs.append(self._prep_arcticdem(cell_ids))
                case _:
                    raise NotImplementedError(f"Member {member} not implemented.")
        with stopwatch("Joining datasets"):
            dataset = dataset.join(member_dfs)
        print(f"Prepared dataset with {len(dataset)} samples and {len(dataset.columns)} features.")

        if cache_mode in ["read", "overwrite"]:
            with stopwatch("Saving features to cache"):
                dataset.to_parquet(cache_file)
        return dataset

    @stopwatch.f("Preparing ERA5", print_kwargs=["era5_agg"])
    def _prep_era5(
        self,
        cell_ids: pd.Series,
        era5_agg: Literal["yearly", "seasonal", "shoulder"],
    ) -> pd.DataFrame:
        era5 = self.read_member("ERA5-" + era5_agg, cell_ids=cell_ids, lazy=True)
        era5_df = _collapse_to_dataframe(era5)
        era5_df.columns = [f"era5_{col}" for col in era5_df.columns]
        # Ensure all target cell_ids are present, fill missing with NaN
        era5_df = era5_df.reindex(cell_ids.to_numpy(), fill_value=np.nan)
        return era5_df

    @stopwatch("Preparing AlphaEarth Embeddings")
    def _prep_embeddings(self, cell_ids: pd.Series) -> pd.DataFrame:
        embeddings = self.read_member("AlphaEarth", cell_ids=cell_ids, lazy=True)
        # For non-synopsis modes there is only a single data variable "embeddings"
        if self.temporal_mode != "synopsis":
            embeddings = embeddings["embeddings"]
        embeddings_df = _collapse_to_dataframe(embeddings)
        embeddings_df.columns = [f"embeddings_{col}" for col in embeddings_df.columns]
        # Ensure all target cell_ids are present, fill missing with NaN
        embeddings_df = embeddings_df.reindex(cell_ids.to_numpy(), fill_value=np.nan)
        return embeddings_df

    @stopwatch("Preparing ArcticDEM")
    def _prep_arcticdem(self, cell_ids: pd.Series) -> pd.DataFrame:
        arcticdem = self.read_member("ArcticDEM", cell_ids=cell_ids, lazy=False)
        if len(arcticdem["cell_ids"]) == 0:
            # No data for these cells - create empty DataFrame with expected columns
            # Use the Dataset metadata to determine column structure
            variables = list(arcticdem.data_vars)
            aggs = arcticdem.coords["aggregations"].to_numpy()
            expected_cols = [f"arcticdem_{var}_{agg}" for var, agg in product(variables, aggs)]
            return pd.DataFrame(index=cell_ids.to_numpy(), columns=expected_cols, dtype=float)
        arcticdem_df = arcticdem.to_dataframe().pivot_table(index="cell_ids", columns="aggregations")
        arcticdem_df.columns = [f"arcticdem_{var}_{agg}" for var, agg in arcticdem_df.columns]
        # Ensure all target cell_ids are present, fill missing with NaN
        arcticdem_df = arcticdem_df.reindex(cell_ids.to_numpy(), fill_value=np.nan)
        return arcticdem_df

    def create_inference_df(
        self,
        batch_size: int | None = None,
        cache_mode: Literal["none", "overwrite", "read"] = "read",
    ) -> Generator[pd.DataFrame]:
        """Create an inference feature set generator.

        This function creates features for all cell IDs in batches.
        If `batch_size` is None or greater than the number of cell IDs, all data is returned in a single batch.

        Args:
            batch_size (int | None, optional): The batch size. If None, all data is returned in a single batch.
                Defaults to None.
            cache_mode (Literal["none", "read", "overwrite"], optional): Caching mode for feature creation.
                "none": No caching.
                "read": Read from cache if exists, otherwise create and save to cache.
                "overwrite": Always create and save to cache, overwriting existing cache.
                Defaults to "read".

        Yields:
            Generator[pd.DataFrame]: Generator yielding DataFrames with features for inference.

        """
        all_cell_ids = self.cell_ids
        if batch_size is None or batch_size >= len(all_cell_ids):
            yield self.make_features(cell_ids=all_cell_ids, cache_mode=cache_mode)
            return

        for i in range(0, len(all_cell_ids), batch_size):
            batch_cell_ids = all_cell_ids.iloc[i : i + batch_size]
            yield self.make_features(cell_ids=batch_cell_ids, cache_mode=cache_mode)

    # @stopwatch.f("Creating training Dataset", print_kwargs=["task", "target", "device", "cache_mode"])
    def create_training_set(
        self,
        task: Task,
        target: TargetDataset,
        device: Literal["cpu", "cuda", "torch"] = "cpu",
        cache_mode: Literal["none", "overwrite", "read"] = "read",
    ) -> TrainingSet:
        """Create a full training set for model training.

        This function creates a full training set with features and target labels,
        splits it into train and test sets, converts to arrays and moves to the specified device.

        Args:
            task (Task): The task.
            target (TargetDataset): The target dataset to use.
            device (Literal["cpu", "cuda", "torch"], optional): The device to move the data to. Defaults to "cpu".
            cache_mode (Literal["none", "read", "overwrite"], optional): Caching mode for feature creation.
                "none": No caching.
                "read": Read from cache if exists, otherwise create and save to cache.
                "overwrite": Always create and save to cache, overwriting existing cache.
                Defaults to "read".

        Returns:
            TrainingSet: The training set.
                Contains a lot of useful information for training and evaluation.

        """
        targets = self.get_targets(target, task)
        assert len(targets) > 0, "No target samples found."
        features = self.make_features(cell_ids=targets.index.to_series(), cache_mode=cache_mode).dropna()
        assert len(features) > 0, "No valid features found after dropping NaNs."
        assert not features.isna().all("index").any(), "Some feature columns are all NaN"

        # Create train / test split
        train_idx, test_idx = train_test_split(features.index.to_numpy(), test_size=0.2, random_state=42, shuffle=True)
        split = pd.Series(index=features.index, dtype=object)
        split.loc[train_idx] = "train"
        split.loc[test_idx] = "test"
        split = split.astype("category")

        X_train = features.loc[train_idx].to_numpy(dtype="float64")
        X_test = features.loc[test_idx].to_numpy(dtype="float64")
        if task in ["binary", "count_regimes", "density_regimes"]:
            y_train = targets.loc[train_idx, "y"].cat.codes.to_numpy(dtype="int64")
            y_test = targets.loc[test_idx, "y"].cat.codes.to_numpy(dtype="int64")
        else:
            y_train = targets.loc[train_idx, "y"].to_numpy(dtype="float64")
            y_test = targets.loc[test_idx, "y"].to_numpy(dtype="float64")

        if device == "cuda":
            X_train = cp.asarray(X_train)
            X_test = cp.asarray(X_test)
            y_train = cp.asarray(y_train)
            y_test = cp.asarray(y_test)
            print(f"Using CUDA device: {cp.cuda.runtime.getDeviceProperties(0)['name'].decode()}")
        elif device == "torch":
            torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            X_train = torch.from_numpy(X_train).to(device=torch_device)
            X_test = torch.from_numpy(X_test).to(device=torch_device)
            y_train = torch.from_numpy(y_train).to(device=torch_device)
            y_test = torch.from_numpy(y_test).to(device=torch_device)
            print(f"Using torch device: {torch.cuda.get_device_name(X_train.device) if X_train.is_cuda else 'cpu'}")
        else:
            assert device == "cpu", f"Invalid {device=} specified."

        return TrainingSet(
            targets=targets.loc[features.index].to_crs("EPSG:4326"),
            features=features,
            X=SplittedArrays(train=X_train, test=X_test),
            y=SplittedArrays(train=y_train, test=y_test),
            z=targets.loc[features.index, "z"],
            split=split,
        )
