from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Activation, Dense, Dropout
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten
from tensorflow.keras.utils import to_categorical, plot_model
from tensorflow.keras.datasets import mnist

# load mnist dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()

# compute the number of labels
num_labels = len(np.unique(y_train))

# convert to one-hot vector
y_train = to_categorical(y_train)
y_test = to_categorical(y_test)

# input image dimensions
image_size = x_train.shape[1]
# resize and normalize
x_train = np.reshape(x_train,[-1, image_size, image_size, 1])
x_test = np.reshape(x_test,[-1, image_size, image_size, 1])
x_train = x_train.astype('float32') / 255
x_test = x_test.astype('float32') / 255

"""
Explanation: CNN on MNIST digits classification
This example is the same as the MLP for MNIST classification. The difference is we are going to use Conv2D layers instead of Dense layers.
The model that will be costructed below is made of:

First 2 layers - Conv2D-ReLU-MaxPool
3rd layer - Conv2D-ReLU
4th layer - Dense(10)
Output Activation - softmax
Optimizer - SGD

Let us first load the packages and perform the initial pre-processing such as loading the dataset, performing normalization and conversion of labels to one-hot.
Recall that in our 3-Dense MLP example, we achieved ~95.3% accuracy at 269k parameters. Here, we can achieve ~98.5% using 105k parameters. CNN is more parameter efficient.
End of explanation
"""


# network parameters
# image is processed as is (square grayscale)
input_shape = (image_size, image_size, 1)
batch_size = 128
kernel_size = 3
filters = 64

"""
Explanation: Hyper-parameters
This hyper-parameters are similar to our MLP example. The differences are kernel_size = 3 which is a typical kernel size in most CNNs and filters = 64.
End of explanation
"""


# model is a stack of CNN-ReLU-MaxPooling
model = Sequential()
model.add(Conv2D(filters=filters,
                 kernel_size=kernel_size,
                 activation='relu',
                 padding='same',
                 input_shape=input_shape))
model.add(MaxPooling2D())
model.add(Conv2D(filters=filters,
                 kernel_size=kernel_size,
                 padding='same',
                 activation='relu'))
model.add(MaxPooling2D())
model.add(Conv2D(filters=filters,
                 kernel_size=kernel_size,
                 padding='same',
                 activation='relu'))
model.add(Flatten())
# dropout added as regularizer
# model.add(Dropout(dropout))
# output layer is 10-dim one-hot vector
model.add(Dense(num_labels))
model.add(Activation('softmax'))
model.summary()

"""
Explanation: Sequential Model Building
The model is similar to our previous example in MLP. The difference is we use Conv2D instead of Dense. Note that due to mismatch in dimensions, the output of the last Conv2D is flattened via Flatten() layer to suit the input vector dimensions of the Dense. Note that though we use Activation(softmax) as the last layer, this can also be integrated within the Dense layer in the parameter activation='softmax'. Both are the same.
End of explanation
"""


#plot_model(model, to_file='cnn-mnist.png', show_shapes=True)

# loss function for one-hot vector
# use of adam optimizer
# accuracy is good metric for classification tasks
model.compile(loss='categorical_crossentropy',
              optimizer='sgd',
              metrics=['accuracy'])
# train the network
model.fit(x_train, y_train, epochs=20, batch_size=batch_size)

loss, acc = model.evaluate(x_test, y_test, batch_size=batch_size)
print("\nTest accuracy: %.1f%%" % (100.0 * acc))


"""
Explanation: Model Training and Evaluation
After building the model, it is time to train and evaluate. This part is similar to MLP training and evaluation.
End of explanation
"""


