# install backend library for keras
# sh> pip install tensorflow
# install keras library
# sh> pip install keras

import keras # check backend library and update keras using 'pip install --upgrade keras'
print(keras.__version__)

import numpy as np
np.random.seed(123)

from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten
from keras.layers import Conv2D, MaxPooling2D
from keras.utils import np_utils

from keras.datasets import mnist

# Loadd pre-shuffled MNIST data into train and test sets
(X_train, y_train), (X_test, y_test) = mnist.load_data()
print(X_train.shape)

import matplotlib.pyplot as plt
plt.imshow(X_train[0])

X_train = X_train.reshape(X_train.shape[0], 28, 28, 1) # theano require the shape to be (B, C, H, W)
X_test = X_test.reshape(X_test.shape[0], 28, 28, 1)
print(X_train.shape)

# Convert 1-dimensional class arrays to 10-dimensional class matrices
Y_train = np_utils.to_categorical(y_train, 10)
Y_test = np_utils.to_categorical(y_test, 10)
print(Y_train.shape)

X_train = X_train.astype('float32')
X_test = X_test.astype('float32')
X_train /= 255
X_test /= 255


if __name__ == "__main__":
    mode = 'baseline':
    if mode=='baseline':
        print("baseline nn to predict mnist db")
        model = Sequential()

        model.add(Conv2D(32, 3, 3, activation='relu', input_shape=(28, 28, 1)))
        model.add(Conv2D(32, 3, 3, activation='relu'))
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.25))

        model.add(Flatten())
        model.add(Dense(128, activation='relu'))
        model.add(Dropout(0.5))
        model.add(Dense(10, activation='softmax'))

        model.compile(loss='categorical_crossentropy',
                      optimizer='adam',
                      metrics=['accuracy'])

        # Training
        model.fit(X_train, Y_train,
                  batch_size=32, nb_epoch=1, verbose=1)

        # Prediction
        pred = model.predict(X_test, verbose=0)
        print("pred", pred.shape)

        # Evaluation
        score = model.evaluate(X_test, Y_test, verbose=0)
        print("score", score)

    elif mode=='bam':
        print("normal nn with bam to predict mnist db")
        model = Sequential()

        model.add(Conv2D(32, 3, 3, activation='relu', input_shape=(28, 28, 1)))
        model.add(Conv2D(32, 3, 3, activation='relu'))

        # bam
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.25))

        model.add(Flatten())
        model.add(Dense(128, activation='relu'))
        model.add(Dropout(0.5))
        model.add(Dense(10, activation='softmax'))

        model.compile(loss='categorical_crossentropy',
                      optimizer='adam',
                      metrics=['accuracy'])

        # Training
        model.fit(X_train, Y_train,
                  batch_size=32, nb_epoch=1, verbose=1)

        # Prediction
        pred = model.predict(X_test, verbose=0)
        print("pred", pred.shape)

        # Evaluation
        score = model.evaluate(X_test, Y_test, verbose=0)
        print("score", score)

    elif mode=='cbam':
        print("normal nn with cbam to predict mnist db")
        model = Sequential()

        model.add(Conv2D(32, 3, 3, activation='relu', input_shape=(28, 28, 1)))
        model.add(Conv2D(32, 3, 3, activation='relu'))

        # cbam
        model.add(MaxPooling2D(pool_size=(2, 2)))
        model.add(Dropout(0.25))

        model.add(Flatten())
        model.add(Dense(128, activation='relu'))
        model.add(Dropout(0.5))
        model.add(Dense(10, activation='softmax'))

        model.compile(loss='categorical_crossentropy',
                      optimizer='adam',
                      metrics=['accuracy'])

        # Training
        model.fit(X_train, Y_train,
                  batch_size=32, nb_epoch=1, verbose=1)

        # Prediction
        pred = model.predict(X_test, verbose=0)
        print("pred", pred.shape)

        # Evaluation
        score = model.evaluate(X_test, Y_test, verbose=0)
        print("score", score)
