import numpy as np
from keras.datasets import mnist
from keras.utils import np_utils
from keras.models import Sequential
from keras.layers import Dense,Dropout,Flatten,Conv2D,MaxPooling2D
import matplotlib.pyplot as plt
import pandas as pd

#avoid warning of AV,AVX....
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

np.random.seed(10)

def main():
    ((x_train,y_train), (x_test,y_test)) = mnist.load_data()

    x_train4D = x_train.reshape(x_train.shape[0],28,28,1).astype('float32')
    x_test4D = x_test.reshape(x_test.shape[0],28,28,1).astype('float32')

    x_train4D_normalize = x_train4D / 255
    x_test4D_normalize = x_test4D / 255

    y_TrainOntHot = np_utils.to_categorical(y_train)
    y_TestOneHot = np_utils.to_categorical(y_test)

    model = Sequential()
    model.add(Conv2D(filters=16,
                     kernel_size=(5,5),
                     padding='same',
                     input_shape=(28,28,1),
                     activation='relu'))

    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Conv2D(filters=36,
                     kernel_size=(5,5),
                     padding='same',
                     input_shape=(28,28,1),
                     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'))

    # show model's layers
    print(model.summary())
    model.compile(loss='categorical_crossentropy',optimizer='adam',metrics=['accuracy'])
    # decide optimization method
    train_history=model.fit(x=x_train4D_normalize,
                             y=y_TrainOntHot,
                             validation_split=0.2,
                             epochs=10,
                             batch_size=300,
                             verbose=2)

    show_train_history(train_history,'acc','val_acc')

    '''not complete '''
    # try:
    #     model.load_weights("SaveModels/MNIST_CNN.h5")
    #     print("load weights.")
    # except:
    #     model.save_weights("SaveModels/MNIST_CNN.h5")
    #     print("Save weights.")

    #show accuracy
    scores = model.evaluate(x_test4D_normalize, y_TestOneHot)
    print(scores[1])
    #show confuse_matrix
    confuse_matrix = pd.crosstab(y_test_label,
                prediction,
                rownames=['label'],
                colnames=['predict'])
    print(confuse_matrix)

    prediction = model.predict_classes(x_test4D_normalize)
    plot_images_labels_prediction(x_test,y_test,prediction,idx=340)


#show graph of accuracy
def show_train_history(train_history,train,validation):
    plt.plot(train_history.history[train])
    plt.plot(train_history.history[validation])
    plt.title('Train History')
    plt.ylabel(train)
    plt.xlabel('Epoch')
    plt.legend(['train', 'validation'], loc='upper left')
    plt.show()

#show MNIST recongition result
def plot_images_labels_prediction(images,labels,prediction,idx,num=10):

    fig = plt.gcf()
    fig.set_size_inches(12, 14)
    if num>25: num=25
    for i in range(0,num):
        ax=plt.subplot(5,5, 1+i)
        ax.imshow(images[idx],cmap='binary')
        title="label="+str(labels[idx])
        if len(prediction)>0:
            title+=",predict="+str(prediction[idx])

        ax.set_title(title,fontsize=10)
        ax.set_xticks([]);ax.set_yticks([])
        idx+=1
    plt.show()

if __name__ == '__main__':
    main()