import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets    # there are different libraries/datasets such as TorchText, TorchVision and TorchAudio
from torchvision.transforms import ToTensor


# Define device
if torch.backends.mps.is_available():
    device = torch.device("mps")
    print(f"Using MPS device. \n")
elif torch.cuda.is_available():
    device = torch.device("cuda")
    print("Using CUDA device.\n")
else:
    device = torch.device("cpu")
    print(f"MPS not available, using CPU. \n")


# download training data from open datasets
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),   # images are orginally stored as PIL images or NumPy arrays, so this converts it to Tensors
)

# dowload test data from open datasets
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)


BATCH_SIZE = 64

# create data loaders
train_dataloader = DataLoader(training_data, batch_size=BATCH_SIZE)
test_dataloader = DataLoader(test_data, batch_size=BATCH_SIZE)


for X, y in test_dataloader:
    # note: N = images in batch,
    #       C = color channels (1 = greyscale, 3 = RGB)
    #       H = height in px, W = width in px
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break


#----------------------
# define model
#----------------------
class NeuralNetwork(nn.Module):
    def __init__(self):
        """
        init() is the machinery
        forward() is the conveyor belt

        - Flatten(): unrolls the 28x28 matrix into a 784 array
        - Linear(input, output): applies linear transform (y=mx+b) to create neurons
        - ReLU(): activation function used to help model learn non-linear patterns
        """
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),  # expands 768 inputs into 512 neurons
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)  # there are 10 different clothing items 0-9
        )

    def forward(self, x):
        """Where the magic happens
        - flatten image
        - pass thru stack (Linear -> ReLu -> Linear -> ReLU -> Linear)

        Args:
            x: images
        
        Returns:
            logits: raw scores for each of the 10 classes
        """
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

# move the mathematical weights to GPU
model = NeuralNetwork().to(device)
print(model)


#-------------------------
# optimizing
#-------------------------
# measures how far off the guess was from the truth
loss_fn = nn.CrossEntropyLoss()
# nudges weights to make error smaller the next time
# model.parameters() are the parts of the models it's allowed to tweak (w and b from Linear)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)


def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()   # enables "training mode"
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)   # moves batch of images and labels to GPU

        # compute prediction error
        pred = model(X) # forward pass
        loss = loss_fn(pred, y) # calculate how wrong the guess was

        # backpropagation: calculate which weights contributed most to the error
        loss.backward() # calculates the gradient (direction & magnitude of error) for every weight in nn
        optimizer.step() # update weight
        optimizer.zero_grad() # reset, need to clear out gradients after every step so old and new gradients don't get mixed up

        # monitoring: print progress every 100 batches
        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)   # loss.item() converts loss from PyTorch tensor back to a Python number
            print(f"loss: {loss:>7f} [{current:>5d}/{size:>5d}]")

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad(): # since we aren't training, we don't need to calculate gradients
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item() # turns true/false into a 1/0
    test_loss /= num_batches # sum up losses from every batch, div by num_batches to get avg loss
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")


#-------------------
# training
#-------------------
EPOCHS = 25
for t in range(EPOCHS):
    print(f"Epoch {t+1}\n----------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")


torch.save(model.state_dict(), "01_mnist_model.pth")
print("Saved PyTorch Model State to 01_mnist_model.pth")