# Warm Up: Pooling
from sklearn import datasets
import matplotlib.pyplot as plt
from scipy.ndimage import convolve
import numpy as np

# Import handwritten digit dataset
digits = datasets.load_digits()

# Convolve handwritten digit number '5'
image = 5
kernel = np.array([[1, 0, -1], [0, 0, 0], [1, 0, -1]])
convolved = convolve(digits.images[image], kernel)
plt.imshow(convolved,cmap='binary')
plt.show()

# Part a: Complete the two lines of code specified below

def maxPool(x, Hp, Wp, stride):
    '''
    This function performs max pooling given an input matrix x, with
    a kernel size of Hp by Wp and a given stride value.
    '''
    # get output dimensions
    H, W = x.shape
    H_out = # TODO: complete this line
    W_out = # TODO: complete this line
    out = np.zeros((H_out, W_out))
    
    # pool
    for i in range(H_out):
        for j in range(W_out):
            out[i, j] = np.max(x[i*stride:i*stride+Hp, j*stride:j*stride+Wp])

    return out

# Test the function
pooled = maxPool(convolved, 3, 3, 1)
plt.imshow(pooled,cmap='binary')
plt.show()

# Part b: Write a function for an average pooling and output the resulting image