# baseline model for the dogs vs cats dataset
import sys
from os import listdir
from random import shuffle
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.image import imread
from keras.models import Sequential
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.layers import Dense
from keras.layers import Flatten
from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array

# define cnn model
def define_model():
    model = Sequential()
    model.add(Conv2D(32, (3, 3), activation='relu', kernel_initializer='he_uniform',
                     padding='same', input_shape=(28, 28, 3)))
    model.add(MaxPooling2D((2, 2)))
    model.add(Flatten())
    model.add(Dense(128, activation='relu', kernel_initializer='he_uniform'))
    model.add(Dense(6, activation='sigmoid'))
    # compile model
    opt = SGD(lr=0.001, momentum=0.9)
    model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
    return model

# plot diagnostic learning curves
def summarize_diagnostics(history):
    # plot loss
    plt.subplot(2, 1, 1)
    plt.title('Cross Entropy Loss')
    plt.plot(history.history['loss'], color='blue', label='train')
    plt.plot(history.history['val_loss'], color='orange', label='test')
    plt.legend()
    # plot accuracy
    plt.subplot(2, 1, 2)
    plt.title('Classification Accuracy')
    plt.plot(history.history['acc'], color='blue', label='train')
    plt.plot(history.history['val_acc'], color='orange', label='test')
    plt.legend()
    # save plot to file
    filename = sys.argv[0].split('/')[-1]
    plt.savefig(filename + '_plot.png')
    plt.tight_layout()
    plt.show()
    plt.close()

# run the test harness for evaluating a model
def run_test_harness():
    # define model
    model = define_model()
    # create data generator
    datagen = ImageDataGenerator(rescale=1.0/255.0)
    # prepare iterators
    train_it = datagen.flow_from_directory('quickdraw/train/',
                                           class_mode='categorical',
                                           batch_size=64,
                                           target_size=(28, 28))
    test_it = datagen.flow_from_directory('quickdraw/test/',
                                          class_mode='categorical',
                                          batch_size=64,
                                          target_size=(28, 28))
    labels = train_it.class_indices
    # switch keys and values
    labels = {v: k for k, v in labels.items()}
    print(labels)
    # fit model
    history = model.fit_generator(train_it,
                                  steps_per_epoch=len(train_it),
                                  validation_data=test_it,
                                  validation_steps=len(test_it),
                                  epochs=50,
                                  verbose=1)
    # evaluate model
    _, acc = model.evaluate_generator(test_it,
                                      steps=len(test_it),
                                      verbose=1)

    print(f'Accuracy on test images: {acc * 100:.2f} %')
    # learning curves
    summarize_diagnostics(history)

    # load 5 random images and print the output
    images = []
    for l in labels.values():
        folder = f'quickdraw/test/{l}/'
        images += [folder + f for f in listdir(folder)]
    shuffle(images)
    for i in range(10):
        filename = images.pop(0)
        plt_img = imread(filename)
        # plot raw pixel data
        plt.imshow(plt_img, cmap="gray")
        plt.show()

        img = load_img(filename, target_size=(28, 28))
        x = img_to_array(img)
        x = np.expand_dims(x, axis=0)
        pred = model.predict(x)
        print(pred)

# entry point, run the test harness
run_test_harness()