# To run this notebook on a local computer, go back to the Jupyter home page and click the upload
# button. A dialog box will appear. Locate the file data_lab_2b.pt and upload.
#
# To run on Google Colab, upload the data_lab_2B.pt file to a folder in your Google Drive and 
# uncomment the next four lines of code. The code assumes the data_lab_2B.pt file is in a folder 
# labeled "ese2000/Lab2B." Change the name to the proper folder if needed.

## import os
# from google.colab import drive

## Mount the drive. It prompts for an authorization.
# drive.mount('/content/drive')

## Specify the directory where the data is located. Change "ese2000/Lab2B" to your own folder name.
# folder_path = '/content/drive/My Drive/ese2000/Lab2B'

## Change the directory to this folder to perform operations within it
# os.chdir(folder_path)

# Import libraries
import torch
import torch.nn as nn
import torch.functional as F
import torch.optim as optim
from torch.nn.functional import relu
from torch.utils.data import TensorDataset, DataLoader
import numpy as np
import matplotlib.pyplot as plt
from IPython.display import Audio, display
import random
np.random.seed(1)
torch.manual_seed(1)
random.seed(1)

print('\n Done \n')

# Define the convolutional function as in LAB2A
def convolution_1D(signal,filter):
    taps = filter.shape[-1]
    z = torch.zeros(signal.shape)
    for k in range(taps):
        z = z + signal * filter[k]
        signal = torch.roll(signal, shifts=1, dims=-1)
        signal[:,:,0] = 0
    return z

# Define the convolutional layer that inherits from nn.Module
# and implements (1) using convolution_1D
class Convolutional_Layer(nn.Module):
    def __init__(self,taps):
        super().__init__()
        self.filter = torch.nn.Parameter(torch.ones(taps))

    def forward(self,x):
        return convolution_1D(x,self.filter)

# Define the CNN, which also inherits from nn.Module
# and has L layers, where each layer instantiates the convolutional_layer class
class CNN(nn.Module):
    def __init__(self, L,taps):
        '''
        inputs:
            L (int): number of layers
            taps (list[int]): list with number of taps for each layer
        '''
        super().__init__()
        self.L = L
        assert len(taps) == L, f"There are {L} layers, but taps is a vector of length {len(taps)}"
        self.conv_layers = []
        for i in range(L):
            self.conv_layers.append(Convolutional_Layer(taps[i]))


    def forward(self, x):
        for i in range(self.L):
            # compute the convolution
            x = self.conv_layers[i](x)
            if i != self.L-1:
                x = torch.tanh(x)
        return x

print('\n Done \n')

# Define the function that implementss a MIMO convolution
def MIMO_convolution_1D(signal,filter):
    taps = filter.shape[0]
    G = filter.shape[1]
    N = signal.shape[-1]
    batch = signal.shape[0]
    z = torch.zeros((batch, G, N)).to(device)
    for k in range(taps):
        z += torch.matmul(filter[k], signal)
        signal = torch.roll(signal, shifts=1, dims=-1)
        signal[:,:,0] = 0
    return z

# Define the convolutional layer class that inherits from nn.module
class MIMO_convolutional_layer(nn.Module):
    def __init__(self,F,G,K):
        '''
        inputs:
            G(int): Number of output features
            F(int): Number of input features
            K(int): Number of filter taps
        '''
        super().__init__()
        self.filter = torch.nn.Parameter(torch.empty(K,G,F))
        nn.init.xavier_normal_(self.filter)
    def forward(self,x):
        return MIMO_convolution_1D(x,self.filter)

print('\n Done \n')

# Define the CNN

class MIMO_CNN(nn.Module):
    def __init__(self, L,taps,features):
        '''
        inputs:
            L (int): number of layers
            taps (list[int]): list with number of taps for each layer [K1, ..., KL]
            features: list with number of features of each layer [F0, ..., FL]
        '''
        super().__init__()
        self.L = L
        assert len(taps) == L, f"There are {L} layers, but taps is a vector of length {len(taps)}"
        assert len(features) == L+1, f"There are {L} layers, but features is a vector of length {len(features)} (There has to be L+1 features)"
        self.layers = nn.ModuleList()
        for i in range(L):
            self.layers.append(MIMO_convolutional_layer(F=features[i],G=features[i+1],K=taps[i]))
            
    def forward(self, x):
        for i in range(self.L):
            x = self.layers[i](x)
            if i != self.L-1:
                x = torch.tanh(x)
        return x

print('\n Done \n')

dataset_name = 'data_lab_2B.pt'
sample_rate = 32000

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device("cpu")
dataB = torch.load(dataset_name, map_location=device, weights_only = False)
print(f'Device: {device}')

# Take one sample and separate x and y
data_0 = dataB[0]
x_0 = data_0[0].cpu()
y_0 = data_0[1].cpu()

print(len(dataB))
print(x_0.shape)

# Noisy signal plot and play
plt.plot(x_0[0])
plt.show()
display(Audio(x_0[0],rate=sample_rate))

# Original signal plot and play
plt.plot(y_0[0])
plt.show()
display(Audio(y_0,rate=sample_rate))

# Split into train and test

test_size = 100
train_size = len(dataB) - test_size

test_set_B, train_set_B =  torch.utils.data.random_split(dataB, [test_size, train_size])

print('\n Done \n')

def evaluate(dataloader, estimator):
    loss = 0
    with torch.no_grad():
        for x_batch, y_batch in dataloader:
            yHat = estimator.forward(x_batch)
            loss += torch.mean((yHat-y_batch)**2)
    numBatchs = len(dataloader)
    loss /= numBatchs
    return loss.item()

print('\n Done \n')

# SAME TRAINING LOOP as in previous lab BUT WITH MIMO_CNN AS ESTIMATOR

L = 2
taps = [40, 40]
features = [1,5,1]
estimator = MIMO_CNN(L=L,taps=taps,features=features)
lr = 0.5
batch_size = 32
n_batches = np.ceil(train_size/batch_size).astype(int)

optimizer = optim.SGD(estimator.parameters(), lr=lr)
loss_L1 = nn.L1Loss()
estimator.to(device)

# Instantiate Data Loaders
train_loader_CL = DataLoader(
    train_set_B,
    batch_size=batch_size,
    shuffle=True,
)

test_loader_CL = DataLoader(
    test_set_B,
    batch_size=batch_size,
    shuffle=False,
)

loss_evolution = []
loss_evolution_test = []

# Training loop
n_epochs = 20

print('\nStart Training Loop\n')

# Iterate n_epochs times over the whole dataset.
for ep in range(n_epochs):
    i=0
    #iterate over all batches in the dataset
    for x_batch, y_batch in train_loader_CL:
        i+=1

        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)

        # Set gradients to zero
        estimator.zero_grad()

        # Compute predictions
        yHat = estimator.forward(x_batch)

        # Compute loss
        loss = loss_L1(yHat,y_batch)
        loss.backward()

        # Update parameters
        optimizer.step()

        print(f'epoch {ep+1} / {n_epochs}, batch {i} / {n_batches},  ', end='\r')
    
    epoch_loss_filter = evaluate(train_loader_CL, estimator)
    loss_evolution.append(epoch_loss_filter)

    lossTest_filter = evaluate(test_loader_CL, estimator)
    loss_evolution_test.append(lossTest_filter)
    print(f'epoch {ep+1} / {n_epochs}, batch {i} / {n_batches}, loss {lossTest_filter}  ', end='\r')

print('\n\nDone\n')

# Plot results
plt.plot(loss_evolution, ".", label="train_filter")
plt.title("Mean squared error evolution")
plt.xlabel("Epoch")
plt.ylabel("L1 loss")
plt.plot(loss_evolution_test, ".", label="test_filter")
plt.legend()
plt.savefig("cnn_mse.pdf")
plt.show()

print('\nDone\n')

# Listen to the samples

# (change to try other audios)
nmbr_example = 0

x_batch, y_batch = next(iter(test_loader_CL))
sample_x, sample_y = x_batch[nmbr_example], y_batch[nmbr_example]
sample_x = sample_x.reshape(sample_x.shape[0], 1, sample_x.shape[1])
predicted_y = estimator(sample_x).detach()[0]

sample_x = sample_x.cpu()
sample_y = sample_y.cpu()
predicted_y = predicted_y.cpu()

print(f'Original Audio')
display(Audio(sample_y, rate=sample_rate))

print(f'Noisy Audio')
display(Audio(sample_x.reshape(sample_x.shape[0],sample_x.shape[2]), rate=sample_rate))

print(f'Audio recovered with filter')
display(Audio(predicted_y, rate=sample_rate))