"""
CW4: Dataset and DataLoader Setup (PROVIDED — do not modify).

Uses PyTorch's torchvision to load FashionMNIST with proper transforms.
Unlike CW2/CW3 which use numpy arrays, CW4 uses PyTorch's native data pipeline.
"""

import torch
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms


# FashionMNIST class names
CLASS_NAMES = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]
NUM_CLASSES = 10

# FashionMNIST statistics (pre-computed)
FASHION_MEAN = 0.2860
FASHION_STD = 0.3530


def get_transforms(augment=False):
    """
    Get image transforms for training and evaluation.

    Args:
        augment: If True, apply data augmentation for training

    Returns:
        train_transform, test_transform
    """
    if augment:
        train_transform = transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(28, padding=4),
            transforms.ToTensor(),
            transforms.Normalize((FASHION_MEAN,), (FASHION_STD,)),
        ])
    else:
        train_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((FASHION_MEAN,), (FASHION_STD,)),
        ])

    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((FASHION_MEAN,), (FASHION_STD,)),
    ])

    return train_transform, test_transform


def get_dataloaders(config):
    """
    Create train, validation, and test DataLoaders.

    Args:
        config: Config dataclass

    Returns:
        train_loader, val_loader, test_loader
    """
    train_transform, test_transform = get_transforms(config.use_augmentation)

    # Download and load datasets
    full_train_dataset = datasets.FashionMNIST(
        root=config.data_dir, train=True, download=True,
        transform=train_transform
    )
    test_dataset = datasets.FashionMNIST(
        root=config.data_dir, train=False, download=True,
        transform=test_transform
    )

    # Split training set into train and validation
    val_size = int(len(full_train_dataset) * config.val_ratio)
    train_size = len(full_train_dataset) - val_size
    train_dataset, val_dataset = random_split(
        full_train_dataset, [train_size, val_size],
        generator=torch.Generator().manual_seed(config.seed)
    )

    # Create DataLoaders
    train_loader = DataLoader(
        train_dataset, batch_size=config.batch_size,
        shuffle=True, num_workers=config.num_workers
    )
    val_loader = DataLoader(
        val_dataset, batch_size=config.batch_size,
        shuffle=False, num_workers=config.num_workers
    )
    test_loader = DataLoader(
        test_dataset, batch_size=config.batch_size,
        shuffle=False, num_workers=config.num_workers
    )

    print(f"Dataset: FashionMNIST")
    print(f"  Train: {train_size}, Val: {val_size}, Test: {len(test_dataset)}")
    print(f"  Augmentation: {'Yes' if config.use_augmentation else 'No'}")
    print(f"  Batch size: {config.batch_size}")

    return train_loader, val_loader, test_loader
