"""Training configuration."""

from dataclasses import dataclass
from typing import Literal


@dataclass
class TrainingConfig:
    """Configuration for model training."""

    # Training parameters
    batch_size: int = 64
    max_epochs: int = 100
    learning_rate: float = 0.001
    weight_decay: float = 1e-5

    # Early stopping
    early_stopping_patience: int = 10
    early_stopping_min_delta: float = 1e-4

    # Learning rate scheduler
    lr_scheduler: Literal['plateau', 'cosine', 'none'] = 'plateau'
    lr_patience: int = 5
    lr_factor: float = 0.5

    # Device
    device: str = 'auto'  # 'auto', 'mps', 'cuda', or 'cpu'

    # Logging
    log_interval: int = 10  # Log every N batches

    # Checkpointing
    save_best: bool = True
    checkpoint_dir: str = './models/checkpoints'

    def get_device(self) -> str:
        """Get the device to use for training.

        Returns:
            Device string ('mps', 'cuda', or 'cpu')
        """
        import torch

        if self.device != 'auto':
            return self.device

        # Auto-detect best available device
        if torch.backends.mps.is_available():
            return 'mps'
        elif torch.cuda.is_available():
            return 'cuda'
        else:
            return 'cpu'


if __name__ == "__main__":
    config = TrainingConfig()
    print(f"Using device: {config.get_device()}")