import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np
import time

# Image normalization
# Converts the image into a tesor and normalizes to acheive a mean of 0 and variance of 1
transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
        ])

# Download MNIST train and test set and normalize images
train_set = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_set = datasets.MNIST('./data', train=False, download=True, transform=transform)

# Define train and test dataloaders, used to get batches sequentially
train_loader = torch.utils.data.DataLoader(train_set, batch_size=128)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=1024)

# Sample Visualization
img_index = 20
image = train_set[img_index][0].squeeze()
plt.imshow(image, cmap = 'binary_r')
plt.title(f"Label: {train_set[img_index][1]}")
print(f"Image size {image.shape} \n")

def conv2D(image, filters):
    """
    Performs 2D convolution on the given image with the given filters.
    Args:
        image (torch.Tensor): input image of shape (height, width).
        filters (torch.Tensor): filters of shape (num_filters, filter_size, filter_size).
                            where filter_size is 2*K+1

    Returns:
        torch.Tensor: convolved image of shape (num_kernels, height, width).
    """
    
    # We are assuming square filters. The true size is filter_size^2
    num_filters, filter_size, _ = filters.shape
    y, x = image.shape
    
    K = (filter_size-1)//2

    # Add zero padding around the border of the image.
    # This is needed to prevent the filter from "hanging off the edge" of the image
    # when we are performing convolutions along the edge pixels.
    # Adding this padding allows us to write simpler code without having to explictly handle this edge case.
    y_pad = y + filter_size - 1
    x_pad = x + filter_size - 1

    # Initialize padded image
    result = torch.zeros((num_filters, y_pad, x_pad))
    result[:, K:y_pad-K, K:x_pad-K] = image

    # Initialize output
    output = torch.zeros((num_filters, y, x))
    
    # Iterate over the kernels (filters are often referred to as "kernels")
    for idx in range(num_filters):
        filt = filters[idx]
        # Iterate over the kernel entries (these are the filter taps AKA the parameters we are learning)
        for k in range(-K, K+1):
            for l in range(-K, K+1):
                # Shift the image vertically and then horizontally
                rolled_image = torch.roll(torch.roll(result[idx], -k, dims=0), -l, dims=1)
                # Multiply the kernel with the rolled image and sum the result
                output[idx] += rolled_image[K: K+y, K: K+x] * filt[k+K, l+K]

    return output

class FilterBank(nn.Module):
    """
    A filter bank
    Args:
        num_filters (Int): The number of filters in the bank (these filters are in parallel!)
        K (Int): Defines the size of each filter in the bank. The true size is (2K+1)^2. 
                    Recall that we sweep from -K to K. So if K is 1, each filter in our bank is a 3x3 grid.
                    The filter is assumed to be square.
    """
    
    def __init__(self,  num_filters, K):
        super(FilterBank, self).__init__()

        # L is the number of pixels in each dimension (For the example above L = (1) * 2 + 1 = 3 and each filter is LxL = 3x3)
        L = K*2+1

        # The number of filters in the bank determines the number of output channels AKA output features
        self.num_filters = num_filters
        self.filters = torch.zeros((num_filters, L, L))

    def forward(self, x):

        # b is the batch size, h is the image height (y) and w is the image width (x)
        b, h, w = x.shape

        # For each image in the batch, we run our filter bank over the image and each filter in the bank produces 1 image (output feature).
        # These images at the output are all in parallel and have the same dimensions as the input image.
        out = torch.zeros((b, self.num_filters, h, w))

        # Compute the convolution for each image in the batch and fill up the output tensor.
        for i, img in enumerate(x):
            out[i] = conv2D(img, self.filters)

        return out

# Take any sample image
img = train_set[0][0]

# Plot an original image (no convolution) to compare against
fig, ax = plt.subplots(1,4, figsize=(14,3.5))
ax[0].imshow(img.squeeze(), cmap = 'binary_r')
ax[0].set_title("Original image")

K_vals = [2,4,8]

for i, K in enumerate(K_vals):
   # Instantiate Filter Bank with one filter (single output channel)
    filt = FilterBank(num_filters = 1, K = K)
    
    # Instantiate each filter in the bank with a set of uniform values
    # (Later on, these filter values will be learned)
    filt.filters = torch.ones((1, 2*K+1, 2*K+1))*(1/(2*K+1)**2)
    
    # Perform 2D convolution
    out = filt(img)
    
    # Plot the convolved image for each K
    ax[i+1].imshow(out.squeeze(), cmap = 'binary_r')
    ax[i+1].set_title(f"Covolved image with K={K}")

# Grab batch of 100 images from the dataset
imgs = next(iter(train_loader))[0][:100]

# Perform convolution with function from torch
K = 8
conv = nn.Conv2d(in_channels = 1, out_channels = 1, kernel_size = 2*K+1, padding = 'same', bias = False)
# Time convolution
start = time.time()
out = conv(imgs)
end = time.time()

print(f"PyTorch's Convolution took: {end-start:.3f} seconds")

# Perform convolution with your function
filterbank = FilterBank(num_filters = 1, K = 8)
# Initialize filters with the torch initialization
filterbank.filters = conv.weight.squeeze(0)

start = time.time()
out_manual = filterbank(imgs.squeeze())
end = time.time()
print(f"Manual Convolution took: {end-start:.3f} seconds")

# Check that outputs match
print(f"Manual convolution and torch convolution output the same tensor: {torch.allclose(out_manual, out, atol = 1e-5)}")

class FilterBank(nn.Module):
    """
    A filter bank. This class is very similar to the one in the previous task, however now we are using torch's built in convolution functions.
    Args:
        num_filters (Int): The number of filters in the bank (these filters are in parallel!)
        K (Int): Defines the size of each filter in the bank. The true size is (2K+1)^2. 
                    Recall that we sweep from -K to K. So if K is 1, each filter in our bank is a 3x3 grid.
                    The filter is assumed to be square.
    """
    def __init__(self, num_filters, K):
        super(FilterBank, self).__init__()

        # In the MNIST dataset each class corresponds to a digit
        num_classes = 10
        
        # Filter Bank with num_filter filters and num_taps taps
        self.conv1 = nn.Conv2d(in_channels = 1, out_channels = num_filters, kernel_size = 2*K+1, stride = 1, padding = 'same')
        
        # Linear layer. This layer takes our num_filters number of output channels and transforms it into an output with dimension num_classes
        self.linear = nn.Linear(num_filters, num_classes)

    def forward(self, x):
        # Pass through filter bank
        x = self.conv1(x)
        
        # Compute energy of each output feature
        # Notice that this condenses filter bank output into a vector of shape (b, num_filters)
        x = torch.sum( x**2, dim=(2,3) )
        
        # Linear map between feature energies and class scores
        x = self.linear(x)
        return x

# Instatiate model
model = FilterBank(num_filters = 30, K = 1)

# Instantiate SGD optimizer
optimizer = optim.SGD(model.parameters(), lr=0.0001)

# Instantiate cross entropy loss
cross_entropy = nn.CrossEntropyLoss()

# Set model in training mode
model.train()

loss_evol = []

# Number of through the entire training set
num_epochs = 10

for epoch in range(num_epochs):
    loss_epoch = 0.
    acc_epoch = 0.

    for data, target in train_loader:
        # Clear gradients of previous batch
        optimizer.zero_grad()
        
        # Compute batch predictions
        output = model(data)
        
        # Compute loss
        loss = cross_entropy(output, target)
        
        # Compute gradients of loss w respect to current params
        loss.backward()
        
        # Update params using SGD
        optimizer.step()

        # Accumulate batch losses
        loss_epoch += loss.item()

        # Accumulate correct predictions, highest score is predicted class.
        pred = output.argmax(dim=-1)
        acc_epoch += torch.sum(pred == target)

    loss_evol.append(loss_epoch/len(train_loader))
    # Print metrics 
    print(f"Epoch: {epoch} \t Loss: {loss_epoch/len(train_loader):.3f} \t Acc: {acc_epoch*100/len(train_loader.dataset):.2f}")

plt.plot(loss_evol)
plt.title("Cross entropy evolution")
plt.xlabel("Epoch")