"""Dataset utilities."""

from __future__ import annotations

from dataclasses import dataclass
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split


@dataclass
class SplitData:
    X_train: np.ndarray
    y_train: np.ndarray
    X_val: np.ndarray
    y_val: np.ndarray
    X_test: np.ndarray
    y_test: np.ndarray


def load_digits_splits(seed: int = 7, test_size: float = 0.20, val_size: float = 0.20) -> SplitData:
    """Load standard sklearn Digits and create stratified train/val/test splits."""
    digits = load_digits()
    X = digits.images.astype(np.float32) / 16.0
    y = digits.target.astype(np.int64)
    X_trainval, X_test, y_trainval, y_test = train_test_split(
        X, y, test_size=test_size, random_state=seed, stratify=y
    )
    val_fraction_of_trainval = val_size / (1.0 - test_size)
    X_train, X_val, y_train, y_val = train_test_split(
        X_trainval, y_trainval, test_size=val_fraction_of_trainval, random_state=seed, stratify=y_trainval
    )
    return SplitData(X_train, y_train, X_val, y_val, X_test, y_test)