"""
Data loading utilities for CW2 and CW3 (numpy-based courseworks).

Provides:
    - load_numpy_data(): Load FashionMNIST from .npy files
    - train_val_split(): Split training data into train/validation sets
    - DataIterator: Mini-batch iterator with optional shuffling
    - one_hot_encode(): Convert integer labels to one-hot vectors
"""

import os
import numpy as np


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


def load_numpy_data(data_dir):
    """
    Load FashionMNIST dataset from .npy files.

    Args:
        data_dir: Path to directory containing .npy files
                  (e.g., "data/fashionmnist")

    Returns:
        dict with keys:
            'train_images': (60000, 784) float32 array, pixel values in [0, 1]
            'train_labels': (60000,) int64 array, class indices 0-9
            'test_images':  (10000, 784) float32 array
            'test_labels':  (10000,) int64 array
    """
    files = {
        "train_images": "train_images.npy",
        "train_labels": "train_labels.npy",
        "test_images": "test_images.npy",
        "test_labels": "test_labels.npy",
    }

    data = {}
    for key, filename in files.items():
        filepath = os.path.join(data_dir, filename)
        if not os.path.exists(filepath):
            raise FileNotFoundError(
                f"Data file not found: {filepath}\n"
                "Please run 'python setup_data.py' first to download the dataset."
            )
        data[key] = np.load(filepath)

    print(f"Loaded FashionMNIST: "
          f"{data['train_images'].shape[0]} train, "
          f"{data['test_images'].shape[0]} test samples")

    return data


def train_val_split(images, labels, val_ratio=0.1, seed=42):
    """
    Split data into training and validation sets (deterministic).

    Args:
        images: (N, D) array of images
        labels: (N,) array of labels
        val_ratio: Fraction of data for validation (default: 0.1)
        seed: Random seed for reproducibility

    Returns:
        train_images, train_labels, val_images, val_labels
    """
    rng = np.random.RandomState(seed)
    n = len(images)
    indices = rng.permutation(n)
    val_size = int(n * val_ratio)

    val_indices = indices[:val_size]
    train_indices = indices[val_size:]

    return (images[train_indices], labels[train_indices],
            images[val_indices], labels[val_indices])


def subset_data(images, labels, num_samples, seed=42):
    """
    Take a random subset of the data (for --quick mode).

    Args:
        images: (N, D) array
        labels: (N,) array
        num_samples: Number of samples to keep
        seed: Random seed

    Returns:
        subset_images, subset_labels
    """
    if num_samples >= len(images):
        return images, labels
    rng = np.random.RandomState(seed)
    indices = rng.choice(len(images), num_samples, replace=False)
    return images[indices], labels[indices]


class DataIterator:
    """
    Mini-batch iterator for numpy arrays.

    Usage:
        iterator = DataIterator(images, labels, batch_size=128, shuffle=True)
        for batch_images, batch_labels in iterator:
            # batch_images: (batch_size, D)
            # batch_labels: (batch_size,)
            ...
    """

    def __init__(self, images, labels, batch_size=128, shuffle=True, seed=None):
        """
        Args:
            images: (N, D) array
            labels: (N,) array
            batch_size: Number of samples per batch
            shuffle: Whether to shuffle data each epoch
            seed: Random seed (if None, uses global numpy state)
        """
        self.images = images
        self.labels = labels
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.rng = np.random.RandomState(seed) if seed is not None else np.random

    def __iter__(self):
        n = len(self.images)
        indices = np.arange(n)
        if self.shuffle:
            self.rng.shuffle(indices)

        for start in range(0, n, self.batch_size):
            batch_idx = indices[start:start + self.batch_size]
            yield self.images[batch_idx], self.labels[batch_idx]

    def __len__(self):
        """Number of batches per epoch."""
        return (len(self.images) + self.batch_size - 1) // self.batch_size


def one_hot_encode(labels, num_classes=NUM_CLASSES):
    """
    Convert integer labels to one-hot encoded vectors.

    Args:
        labels: (N,) array of integer labels
        num_classes: Number of classes (default: 10)

    Returns:
        (N, num_classes) float32 array with 1.0 at label positions
    """
    n = len(labels)
    one_hot = np.zeros((n, num_classes), dtype=np.float32)
    one_hot[np.arange(n), labels] = 1.0
    return one_hot
