import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import FashionMNIST
from torch.utils.data.dataloader import DataLoader
from torch.utils.data import random_split
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
%matplotlib inline

# Checking the Available GPU
if torch.cuda.is_available():
    device = torch.device('cuda:0')
    print('GPU')
else:
    device = torch.device('cpu')
    print('CPU')
print(device)

# Downloading the Data and Splitting it into (Training and Validation Sets)

transform = transforms.Compose([transforms.ToTensor()])

data = FashionMNIST(root = 'data/', train = True, download = True, transform = transform)
# test_data = MNIST(root = 'datasets/MNIST', train = False, download = True, transform = transform)

val_data_size = 10000
train_data_size = len(data) - val_data_size

train_data, val_data = random_split(data, [train_data_size, val_data_size])


print(len(train_data), len(val_data))

# Creating Batches of Training and Validation Data

batch_size = 128
train_loader = DataLoader(train_data, batch_size = batch_size, shuffle = True, num_workers = 3, pin_memory = True)
val_loader = DataLoader(val_data, batch_size = batch_size, shuffle = True, num_workers = 3, pin_memory = True)
print(len(val_loader))

# Visualizing Batch of Training Data into a Grid using make_grid() Function

for images, labels in train_loader:
    print('Printing a Grid of Batch Images\n')
    plt.figure(figsize = (16,7))
    plt.axis('off')
    plt.imshow(make_grid(images, nrow = 16).permute((1,2,0)))
    break