"""
Setup script for downloading and preparing the FashionMNIST dataset.

Run this script once before starting any coursework:
    python setup_data.py

This will download FashionMNIST and save it as numpy arrays in data/fashionmnist/.
"""

import os
import numpy as np

def download_and_prepare():
    """Download FashionMNIST via torchvision and save as .npy files."""
    # Import torchvision only here — CW2/CW3 do not need PyTorch
    try:
        import torchvision
    except ImportError:
        print("Error: torchvision is required to download the dataset.")
        print("Install it with: pip install torchvision")
        return

    data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data", "fashionmnist")
    os.makedirs(data_dir, exist_ok=True)

    # Check if data already exists
    required_files = ["train_images.npy", "train_labels.npy", "test_images.npy", "test_labels.npy"]
    if all(os.path.exists(os.path.join(data_dir, f)) for f in required_files):
        print("Dataset already exists in data/fashionmnist/. Skipping download.")
        print("Delete the data/fashionmnist/ directory to re-download.")
        return

    print("Downloading FashionMNIST dataset...")

    # Download using torchvision (saves to a temporary location)
    raw_dir = os.path.join(data_dir, "_raw")
    train_dataset = torchvision.datasets.FashionMNIST(root=raw_dir, train=True, download=True)
    test_dataset = torchvision.datasets.FashionMNIST(root=raw_dir, train=False, download=True)

    # Convert to numpy arrays
    # Images: flatten to (N, 784), normalize to [0, 1] as float32
    train_images = train_dataset.data.numpy().reshape(-1, 784).astype(np.float32) / 255.0
    train_labels = train_dataset.targets.numpy().astype(np.int64)
    test_images = test_dataset.data.numpy().reshape(-1, 784).astype(np.float32) / 255.0
    test_labels = test_dataset.targets.numpy().astype(np.int64)

    # Save as .npy files
    np.save(os.path.join(data_dir, "train_images.npy"), train_images)
    np.save(os.path.join(data_dir, "train_labels.npy"), train_labels)
    np.save(os.path.join(data_dir, "test_images.npy"), test_images)
    np.save(os.path.join(data_dir, "test_labels.npy"), test_labels)

    # Print dataset statistics
    class_names = [
        "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
        "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
    ]

    print("\n=== FashionMNIST Dataset Statistics ===")
    print(f"Training set: {train_images.shape[0]} samples")
    print(f"Test set:     {test_images.shape[0]} samples")
    print(f"Image shape:  {train_images.shape[1]} features (28x28 flattened)")
    print(f"Pixel range:  [{train_images.min():.1f}, {train_images.max():.1f}]")
    print(f"Number of classes: {len(class_names)}")
    print(f"\nClass distribution (training set):")
    for i, name in enumerate(class_names):
        count = np.sum(train_labels == i)
        print(f"  {i}: {name:<15s} — {count} samples")

    print(f"\nData saved to: {data_dir}")
    print("You can now start working on the courseworks!")

    # Clean up raw torchvision download
    import shutil
    if os.path.exists(raw_dir):
        shutil.rmtree(raw_dir)
        print("Cleaned up temporary download files.")


if __name__ == "__main__":
    download_and_prepare()
