'''
The task is to build an autoencoder for the MNIST database. This will learn
to form a compressed representation of the handwritten digit (encoding)
and then decode it back into the original image.

Here, I follow the tutorial found at:

    https://blog.keras.io/building-autoencoders-in-keras.html

In this rendition, we will be using convolutional neural networks to downsize the images.
Since these types of neural nets are better at storing the local structure of the
data they are compression, they should work better than a standard feedforward layer
for encoding/decoding

'''

# This first line imports a bunch of different neuron layers we can use
# Input: self-explanatory; stores input data that is fed-forward to future layers

# Dense: standard feedforward layer with output = activation_func(Weights*input + bias)

# Conv2D: convolutional layer; each neuron in this layer has a receptive field that it looks at;
#         unlike a typical feedforward layer, these neurons only have connections to the neurons 
#         in their receptive field (some subset of the previous layer). For images, this is typically
#         just set to be some square (i.e. 3 by 3 grid of pixels) in the image.
#         For each receptive field, we define a number of neurons that are attached to it called the
#         DEPTH of the CNN. All the neurons at a given depth form what is called a "filter". The idea
#         is that we collect this series of filters where each filter acts on the entire image by
#         combining the results of a bunch of local analysis of the image. Each filter extracts its
#         own set of features.
#         Shared Weights: the neurons at each depth use the same weights/bias

# maxPooling: this neuron is what actually allows for the downsampling, reducing the layers size, not Conv2D.
#             here, each neuron in this layer is connected to some subset of the previous layer
#             It chooses its activity based on the max activity in the previous layer. Typically these are
#             chosen to be a small grid for the same reasons as in Conv2D: to capture local properties of the image
#             NOTE: these grids/filters typically don't overlap?

# UpSampling: allows network to regain the dimensionality it lost during maxPooling. For each neuron in the previous
#             layer, it simply generates multiple copies of neurons matching that activity level?
from keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D
from keras.models import Model, Sequential
from keras import backend as K

use_sequential = True # use Sequential() model class rather than general Model() class; both produce same results

# note now we are not going to vectorize our data because we care about the local structure
# instead, input will be fed in as matrices.
input_img = Input(shape=(28,28,1))

# convert the image into a lower dimensional representation

# Conv2D( depth, kernel/filter size, ...)
# The depth refers to the number of neurons connected to the same receptive field in the input.
# In the case below, we have 16 neurons connected to each 3 x 3 receptive field of pixels in 
# the input. These receptive fields are generated by moving the 3 x 3 filter one pixel at a time 
# (there is another parameter called 'strides' that can specify how many pixels to move each time).
# These 16 neurons convolve with their receptive field to capture the local patterns. Important
# in cases where local structure is the most important. As these layers only need to consider their
# receptive field, they are NOT connected to any other neurons in the previous layer

# To the right, we have shown how the dimensionality of the input changes. 

# NOTE: maxpooling does NOT effect the depth of the layers; it only reduces the number of filters
# Maxpooling reduces the

# 'same' padding ensures that the output has the same size (width and height, not depth) as input

x = Conv2D(16, (3,3), activation='relu', padding='same')(input_img) # (28, 28, 1) --> (28, 28, 16) 
x = MaxPooling2D((2,2), padding='same')(x)                          # (28, 28, 16) --> (14, 14, 16)
x = Conv2D(8, (3,3), activation = 'relu', padding = 'same')(x)      # (14,14, 16) --> (14,14, 8)
x = MaxPooling2D((2,2), padding = 'same')(x)                        # (14, 14, 8) --> (7,7, 8)
x = Conv2D(8, (3,3), activation = 'relu', padding = 'same')(x)      # (7,7,8) --> (7,7,8)
encoded = MaxPooling2D((2,2), padding='same')(x)                    # (7,7,8) --> (4,4,8)

# decode the lower dimensional representation back into an image
y = Conv2D(8, (3,3), activation='relu', padding = 'same')(encoded)
y = UpSampling2D((2,2))(y)
y = Conv2D(8,(3,3), activation = 'relu', padding = 'same')(y)
y = UpSampling2D((2,2))(y)
y = Conv2D(16, (3,3), activation = 'relu')(y)
y = UpSampling2D((2,2))(y)
# want to use sigmoid as our final activation function to make the output more
# interpretable as an image = normalized pixel values are our expected output
decoded = Conv2D(1, (3,3), activation='sigmoid', padding='same')(y)

autoencoder = Model(input = input_img, output=decoded)
autoencoder.compile(optimizer = 'adadelta', loss='binary_crossentropy')


#===============================================================================`
# same as above but uses Sequential() model class
model = Sequential()
model.add(Conv2D(16, (3,3), activation='relu', padding='same', input_shape = (28,28,1)))         # (28, 28, 1) --> (28, 28, 16) 
model.add(MaxPooling2D((2,2), padding='same'))                          # (28, 28, 16) --> (14, 14, 16)
model.add(Conv2D(8, (3,3), activation = 'relu', padding = 'same'))      # (14,14, 16) --> (14,14, 8)
model.add(MaxPooling2D((2,2), padding = 'same'))                        # (14, 14, 8) --> (7,7, 8)
model.add(Conv2D(8, (3,3), activation = 'relu', padding = 'same'))      # (7,7,8) --> (7,7,8)
model.add( MaxPooling2D((2,2), padding='same'))                         # (7,7,8) --> (4,4,8)

# decode the lower dimensional representation back into an image
model.add(Conv2D(8, (3,3), activation='relu', padding = 'same'))
model.add(UpSampling2D((2,2)))
model.add(Conv2D(8,(3,3), activation = 'relu', padding = 'same'))
model.add(UpSampling2D((2,2)))
model.add(Conv2D(16, (3,3), activation = 'relu'))
model.add(UpSampling2D((2,2)))
# want to use sigmoid as our final activation function to make the output more
model.add(Conv2D(1, (3,3), activation='sigmoid', padding='same'))

model.compile(optimizer = 'adadelta', loss='mean_squared_error')
#===============================================================================`

# Load the data again!

from keras.datasets import mnist
import numpy as np

# again, MNIST gives (image, label) pairs, but we dont care about the label as we are just reconstructing 
# the image. So throw out the labels with _
(x_train, _), (x_test, _) = mnist.load_data()

# normalize pixels to be in [0,1]
x_train = x_train.astype('float32') / 255.0
x_train = np.reshape(x_train, (len(x_train), 28,28,1))
x_test = x_test.astype('float32') / 255.0
x_test = np.reshape(x_test, (len(x_test), 28,28,1))


# open up some log files for storing training progress
# opens up a local host server at http://0.0.0.0:6006
import os
import webbrowser
from keras.callbacks import TensorBoard
os.system('tensorboard --logdir=/tmp/autoencoder &') # & runs program in background so script can keep running
webbrowser.open('http://localhost:6006', new=1) # opens a webbrowser to display log information

# fit model (train network)!
if not use_sequential:
    autoencoder.fit(x_train, x_train,
                    epochs = 50, 
                    batch_size = 128, 
                    shuffle = True,
                    validation_data = (x_test, x_test),
                    callbacks = [TensorBoard(log_dir='/tmp/autoencoder')]) # store log files
    # make predictions!
    predicted_images = autoencoder.predict(x_test)
else:
    model.fit(x_train, x_train,
              epochs = 1, 
              batch_size = 128, 
              shuffle = True,
              validation_data = (x_test, x_test),
              callbacks = [TensorBoard(log_dir='/tmp/autoencoder')]) # store log files

    # make predictions!
    predicted_images = model.predict(x_test)



# plot stuff
import matplotlib.pyplot as plt


n = 10 # number of digits to display

plt.figure(figsize=(20,4))
for i in range(n):

    # display original (in top row)
    ax = plt.subplot(2, n, i+1) # which subplot to work with; 2 rows, n columns, slot i+1
    plt.imshow(x_test[i].reshape(28,28))
    plt.gray()
    ax.get_xaxis().set_visible = False
    ax.get_yaxis().set_visible = False

    # display predicted (in bottom row)
    ax = plt.subplot(2, n, i+ 1 + n) # which subplot to work with; 2 rows, n columns, slot i+1
    plt.imshow(predicted_images[i].reshape(28,28))
    plt.gray()
    ax.get_xaxis().set_visible = False
    ax.get_yaxis().set_visible = False

plt.show()
