# TODO: Run this cell to import relevant packages

import torch  # Main torch import for torch tensors (arrays)
import torch.nn as nn  # Neural network module for building deep learning models
import torch.nn.functional as F  # Functional module, includes activation functions
import torch.optim as optim  # Optimization module
import torchvision  # Vision / image processing package built on top of torch

from sklearn.metrics import accuracy_score  # Computing accuracy metric

# TODO: Run this cell to download the data and setup the pre-processing pipeline

# Common practice to normalize input data to neural networks (0 mean, unit variance)
transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor(),  # All inputs to PyTorch neural networks must be torch.Tensor
    torchvision.transforms.Normalize(mean=0.1307, std=0.3081)  # Subtracts mean and divides by std. Note that the raw data is between [0, 1]
])

# Download the MNIST data and lazily apply the transformation pipeline
train_data = torchvision.datasets.MNIST('./datafiles/', train=True, download=True, transform=transform)
test_data = torchvision.datasets.MNIST('./datafiles/', train=False, download=True, transform=transform)

# Setup data loaders
# Note: Iterating through the dataloader yields batches of (inputs, targets)
# where inputs is a torch.Tensor of shape (batch, 28, 28) and targets is a torch.Tensor of shape (batch,)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=1000)

class ConvMaxPoolNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        # Activation function
        self.relu = nn.ReLU()

        # TODO: Create conv and max pool layers for block 1
        self.conv1 = nn.Conv2d(
            in_channels=1, out_channels=32, kernel_size=(3, 3), padding="same"
        )
        self.max_pool1 = nn.MaxPool2d(kernel_size=(2, 2), stride=2)

        # TODO: Create conv and max pool layers for block 2
        self.conv2 = nn.Conv2d(
            in_channels=32, out_channels=64, kernel_size=(3, 3), padding="same"
        )
        self.max_pool2 = nn.MaxPool2d(kernel_size=(2, 2), stride=2)

        # TODO: Create conv and max pool layers for block 3
        self.conv3 = nn.Conv2d(
            in_channels=64, out_channels=100, kernel_size=(3, 3), padding="same"
        )
        self.max_pool3 = nn.MaxPool2d(kernel_size=(2, 2), stride=2)

        # TODO: Create output linear layer
        self.output = nn.Linear(900, 10)

        #####

    def forward(self, x):
        """
        Forward pass implementation for the network

        :param x: torch.Tensor of shape (batch, 1, 28, 28), input images

        :returns: torch.Tensor of shape (batch, 10), output logits
        """
        # TODO: Process x through block 1
        x = self.conv1(x)
        x = self.relu(x)
        x = self.max_pool1(x)

        # TODO: Process x through block 2
        x = self.conv2(x)
        x = self.relu(x)
        x = self.max_pool2(x)

        # TODO: Process x through block 3
        x = self.conv3(x)
        x = self.relu(x)
        x = self.max_pool3(x)

        # TODO: Flatten x and process through the output layer
        x = torch.flatten(x, 1)  # (B, 900)
        x = self.output(x)
        #####
        return x

# TODO: Run this cell to help test your code
# Tip: Try to build the model layer by layer and test with this cell as you go
# Final expected output shape: (4, 10)
x = torch.randn(4, 1, 28, 28)
cnn_mp = ConvMaxPoolNetwork()

y = cnn_mp(x)
y.shape

class ConvBatchNormNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        # Activation function
        self.relu = nn.ReLU()

        self.conv1 = nn.Conv2d(
            in_channels=1, out_channels=32, kernel_size=(3, 3), stride=2, padding=1
        )
        self.batchNorm1 = nn.BatchNorm2d(num_features=32)

        # TODO: Create conv and batch norm layers for block 1
        self.batchNorm2 = nn.BatchNorm2d(num_features=64)
        self.conv2 = nn.Conv2d(
            in_channels=32, out_channels=64, kernel_size=(3, 3), stride=2, padding=1
        )

        # TODO: Create conv and batch norm layers for block 1
        self.batchNorm3 = nn.BatchNorm2d(num_features=100)
        self.conv3 = nn.Conv2d(
            in_channels=64, out_channels=100, kernel_size=(3, 3), stride=2, padding=1
        )

        # TODO: Create adaptive average pooling and output linear layers
        self.pool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        self.output = nn.Linear(100, 10)

        #####

    def forward(self, x):
        """
        Forward pass implementation for the network

        :param x: torch.Tensor of shape (batch, 1, 28, 28), input images

        :returns: torch.Tensor of shape (batch, 10), output logits
        """
        # TODO: Process x through block 1
        x = self.conv1(x)
        x = self.relu(x)
        x = self.batchNorm1(x)

        # TODO: Process x through block 2
        x = self.conv2(x)
        x = self.relu(x)
        x = self.batchNorm2(x)

        # TODO: Process x through block 3
        x = self.conv3(x)
        x = self.relu(x)
        x = self.batchNorm3(x)

        # TODO: Flatten x and process through the output layer
   


        # TODO: Pool x, squeeze extra dims, and process through the output layer
        x = self.pool(x)
        x = torch.squeeze(x)  # (B, 100)
        x = self.output(x)
        #####
        return x

# TODO: Run this cell to help test your code
# Tip: Try to build the model layer by layer and test with this cell as you go
# Final expected output shape: (4, 10)
cnn_bn = ConvBatchNormNetwork()
x = torch.randn(4, 1, 28, 28)
z = cnn_bn(x)
z.shape

# TODO: Run this cell to define the train function

def train(model, train_loader, loss_fn, optimizer, epoch=-1):
    """
    Trains a model for one epoch (one pass through the entire training data).

    :param model: PyTorch model
    :param train_loader: PyTorch Dataloader for training data
    :param loss_fn: PyTorch loss function
    :param optimizer: PyTorch optimizer, initialized with model parameters
    :kwarg epoch: Integer epoch to use when printing loss and accuracy
    """
    total_loss = 0
    all_predictions = []
    all_targets = []

    model.train()  # Set model in training mode
    for i, (inputs, targets) in enumerate(train_loader):  # 1. Fetch next batch of data
        optimizer.zero_grad()  # 2. Zero parameter gradients
        outputs = model(inputs)  # 3. Compute model outputs
        loss = loss_fn(outputs, targets)  # 4. Compute loss between outputs and targets
        loss.backward()  # 5. Backpropagation for parameter gradients
        optimizer.step()  # 6. Gradient descent step

        # Track some values to compute statistics
        total_loss += loss.item()
        preds = torch.argmax(outputs, dim=-1)  # Take the class with the highest output as the prediction
        all_predictions.extend(preds.tolist())
        all_targets.extend(targets.tolist())

        # Print some statistics every 100 batches
        if i % 100 == 0:
            running_loss = total_loss / (i + 1)
            print(f"Epoch {epoch + 1}, batch {i + 1}: loss = {running_loss:.2f}")
     
    acc = accuracy_score(all_targets, all_predictions)

    # Print average loss and accuracy
    print(f"Epoch {epoch + 1} done. Average train loss = {total_loss / len(train_loader):.2f}, average train accuracy = {acc * 100:.3f}%")

# TODO: Run this cell to define the test function

def test(model, test_loader, loss_fn, epoch=-1):
    """
    Tests a model for one epoch of test data.

    Note:
        In testing and evaluation, we do not perform gradient descent optimization, so steps 2, 5, and 6 are not needed.
        For performance, we also tell torch not to track gradients by using the `with torch.no_grad()` context.

    :param model: PyTorch model
    :param test_loader: PyTorch Dataloader for test data
    :param loss_fn: PyTorch loss function
    :kwarg epoch: Integer epoch to use when printing loss and accuracy
    """
    total_loss = 0
    all_predictions = []
    all_targets = []

    model.eval()  # Set model in evaluation mode
    for i, (inputs, targets) in enumerate(test_loader):  # 1. Fetch next batch of data
        with torch.no_grad():
            outputs = model(inputs)  # 3. Compute model outputs
            loss = loss_fn(outputs, targets)  # 4. Compute loss between outputs and targets

            # Track some values to compute statistics
            total_loss += loss.item()
            preds = torch.argmax(outputs, dim=-1)  # Take the class with the highest output as the prediction
            all_predictions.extend(preds.tolist())
            all_targets.extend(targets.tolist())
      
    acc = accuracy_score(all_targets, all_predictions)

    # Print average loss and accuracy
    print(f"Epoch {epoch + 1} done. Average test loss = {total_loss / len(test_loader):.2f}, average test accuracy = {acc * 100:.3f}%")

# TODO: Run this cell to train the ConvMaxPoolNetwork for 5 epochs
cnn_mp = ConvMaxPoolNetwork()
optimizer = optim.Adam(cnn_mp.parameters(), lr=1e-2)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(5):
    train(cnn_mp, train_loader, loss_fn, optimizer, epoch)
    test(cnn_mp, test_loader, loss_fn, epoch)

# TODO: Run this cell to train the ConvBatchNormNetwork for 5 epochs
cnn_bn = ConvBatchNormNetwork()
optimizer = optim.Adam(cnn_bn.parameters(), lr=1e-2)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(5):
    train(cnn_bn, train_loader, loss_fn, optimizer, epoch)
    test(cnn_bn, test_loader, loss_fn, epoch)

# If on Google Colab, uncomment and run this
 #!pip install torchinfo

# TODO: Run this cell to import torchinfo.summary
from torchinfo import summary

# TODO: Run this cell to instantiate a MLP model
class MultiLayerPerceptron(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(784, 100)
        self.output = nn.Linear(100, 10)
    
    def forward(self, x):
        """
        Forward pass implementation for the network
        
        :param x: torch.Tensor of shape (batch, 28, 28), input images

        :returns: torch.Tensor of shape (batch, 10), output logits
        """
        x = torch.flatten(x, 1)  # shape (batch, 28*28)
        x = self.hidden(x)
        x = F.relu(x)
        x = self.output(x)
        return x

mlp = MultiLayerPerceptron()

# TODO: Run this cell to view a summary of each model and note the number of parameters
print(summary(mlp))
print()
print(summary(cnn_mp))
print()
print(summary(cnn_bn))

# TODO: Fill in the results in this cell. Then, run the cell to view a table for easy comparison.
import pandas as pd
results = [
    {
        "model": "MLP",
        "epoch_1_train_loss": 0.28,
        "epoch_5_train_loss": 0.16,
        "epoch_5_test_accuracy": 95.38,
        "parameters": None
    },
    {
        "model": "CNN MaxPool",
        "epoch_1_train_loss": None,
        "epoch_5_train_loss": None,
        "epoch_5_test_accuracy": None,
        "parameters": None
    },
    {
        "model": "CNN BatchNorm",
        "epoch_1_train_loss": None,
        "epoch_5_train_loss": None,
        "epoch_5_test_accuracy": None,
        "parameters": None
    }
]
df = pd.DataFrame(results)
df