import matplotlib.pyplot as plt
from scipy.io import loadmat
import numpy as np

# Load data for activity
in_data = loadmat('bucky.mat')
A = in_data['A']

#####################################################################################################
# conv function code source:                                                                        #
# https://towardsdatascience.com/a-guide-to-convolutional-neural-networks-from-scratch-f1e3bfc3e2de #
#####################################################################################################

def conv(x, filt, stride, pad):
    """
    - H: height (number of rows)
    - W: width (number of columns)
    - stride: stride value to apply to each image when convolving
    - pad: number of rows and columns to zero-pad around image
    """
    H, W = x.shape
    HH, WW = filt.shape
    # get output dimensions
    H_out = 1 + (H + 2 * pad - HH) // stride
    W_out = 1 + (W + 2 * pad - WW) // stride
    out = np.zeros((H_out, W_out))
    pad_widths = ((pad,), (pad,))
    xpad = np.pad(x, pad_widths, 'constant')
    Hpad, Wpad = xpad.shape
    
    # perform convolution   
    for i in range(0, Hpad-(HH-1), stride):
        for j in range(0, Wpad-(WW-1), stride):
            prod = np.sum(np.multiply(filt, xpad[i:i+HH, j:j+WW]))
            out[int(i/stride), int(j/stride)] = prod

    return out

# Display original image
plt.imshow(A,cmap='binary')
plt.title('Original')
plt.show()

# Convolve with filter and display result
filter1 = np.array([[1, 0, -1], [1, 0, -1], [1, 0, -1]])
convolved1 = conv(A, filter1, 2, 1)
plt.imshow(convolved1,cmap='binary')
plt.title('1st convolution')
plt.show()

# Part C answer (can vary)

# add another filter and convolutional layer
filter2 = # TODO: Insert diagonal edge filter here
convolved2 = # TODO: convolve the output of the first convolution layer with filter2
plt.imshow(convolved2,cmap='binary')
plt.title('2nd convolution')
plt.show()