import numpy as np
from numpy.lib.stride_tricks import as_strided
from PIL import Image
import matplotlib.pyplot as plt


def load_data():
    return np.asarray(Image.open("cat.png").convert("L"), dtype=np.uint8) / 255

image = load_data()

def dimension(length, padding, stride, kernel):
    return int((length + 2 * padding - kernel + 1)/stride)

def pad(image, size, strategy):
    (ih, iw) = image.shape
    (ph, pw) = size

    (rh, rw) = (ih + 2*ph, iw + 2*pw)

    if strategy == 'zero':
        base = np.zeros((rh, rw))
        base[ph:ph+ih, pw:pw+iw] = image.copy()
        return base
    elif strategy == 'mirror':
        raise NotImplementedError("Missing")

def convolute(weights):
    return lambda view : np.dot(view.flatten(), weights.flatten())

def pool(image, stride=1, kernel_size=(2,2), transform=np.max, padding_size=(0,0), padding_strategy='zero'):
    # calculate output size
    (height, width) = [dimension(image.shape[i], padding_size[i], stride, kernel_size[i]) for i in range(2)]
    
    padded = pad(image, padding_size, padding_strategy)

    next = np.zeros((height, width))

    #'zero'|'mirror'
    for i in range(height):
        for j in range(width):
            # Translate upper left corner into original image space
            (y, x) = (i*stride, j*stride)

            # Retrieve image window 
            view = padded[y:y+kernel_size[0], x:x+kernel_size[1]]

            #print(view)
            # Set pooled image value
            next[i][j] = transform(view)

    return next

weights = np.array([
    [-10,-5,0,5,10],
    [-10,-5,0,5,10],
    [-10,-5,0,5,10],
    [-10,-5,0,5,10],
    [-10,-5,0,5,10]
])

new1 = pool(image, kernel_size=(5,5), padding_size=(2,2), transform=convolute(weights))
new2 = pool(new1, stride=3, kernel_size=(30,30), transform=np.max)

fig = plt.figure()
fig.canvas.manager.set_window_title("A Cat")

plt.imshow(new2, cmap="gray")
plt.show()