
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.layers import MaxPooling2D
from tensorflow.keras.layers import GlobalAveragePooling2D
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Flatten
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.layers import Dropout
from tensorflow.keras.layers import BatchNormalization
from matplotlib import pyplot

"""
source for this model:

STRIVING FOR SIMPLICITY: THE ALL CONVOLUTIONAL NET
Jost Tobias Springenberg∗, Alexey Dosovitskiy∗, Thomas Brox, Martin Riedmiller
Department of Computer Science
University of FreiburgFreiburg, 79110, Germany

https://arxiv.org/pdf/1412.6806.pdf
"""


# load train and test dataset
def load_dataset():
    # load dataset
    (trainX, trainY), (testX, testY) = cifar10.load_data()
    from sklearn.model_selection import train_test_split
    valX, testX, valY, testY = train_test_split(testX, testY, test_size=0.2, random_state=42)
    # one hot encode target values
    trainY = to_categorical(trainY)
    valY = to_categorical(valY)
    testY = to_categorical(testY)
    return trainX, trainY, valX, valY, testX, testY


# scale pixels
def prep_pixels(train, val, test):
    # convert from integers to floats
    train_norm = train.astype('float32')
    val_norm = val.astype('float32')
    test_norm = test.astype('float32')
    # normalize to range 0-1
    train_norm = train_norm / 255.0
    val_norm = val_norm / 255.0
    test_norm = test_norm / 255.0
    # return normalized images
    return train_norm, val_norm, test_norm



def define_model():
    model = Sequential()
    model.add(Conv2D(96, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same', input_shape=(32, 32, 3)))
    model.add(Dropout(0.2))
    model.add(Conv2D(96, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))

    model.add(Conv2D(96, (3, 3), strides=(2, 2), activation='relu', kernel_initializer='he_uniform', padding='same'))
    model.add(Dropout(0.5))

    model.add(Conv2D(192, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
    model.add(Conv2D(192, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))

    model.add(Conv2D(96, (3, 3), strides=(2, 2), activation='relu', kernel_initializer='he_uniform', padding='same'))
    model.add(Dropout(0.5))

    model.add(MaxPooling2D((3, 3), strides=(2, 2)))
    model.add(Dropout(0.5))
    model.add(Conv2D(192, (3, 3), activation='relu', kernel_initializer='he_uniform', padding='same'))
    model.add(Conv2D(192, (1, 1), activation='relu', kernel_initializer='he_uniform', padding='same'))
    model.add(Conv2D(10, (1, 1), activation='relu', kernel_initializer='he_uniform', padding='same'))

    model.add(GlobalAveragePooling2D())

    model.add(Dense(10, activation='softmax'))
    # compile model
    opt = SGD(lr=0.001, momentum=0.9)
    model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
    return model

def summarize_diagnostics(history, filename, accuracy):
    # plot loss
    pyplot.subplot(211)
    pyplot.title('Cross Entropy Loss')
    pyplot.plot(history.history['loss'], color='blue', label='train')
    pyplot.plot(history.history['val_loss'], color='orange', label='test')
    # plot accuracy
    pyplot.subplot(212)
    pyplot.title('Classification Accuracy')
    pyplot.plot(history.history['accuracy'], color='blue', label='train')
    pyplot.plot(history.history['val_accuracy'], color='orange', label='test')
    pyplot.xlabel(f'test accuracy = {accuracy * 100:.2f} %')
    # save plot to file

    pyplot.savefig(f'resources\\{filename}_plot.png')
    pyplot.close()

def run_test_harness():
    modelname = input('please specify model name: ')
    # load dataset
    trainX, trainY, valX, valY, testX, testY = load_dataset()
    # prepare pixel data
    trainX, valX, testX = prep_pixels(trainX, valX, testX)
    # define model
    model = define_model()
    # create data generator
    datagen = ImageDataGenerator(width_shift_range=0.1, height_shift_range=0.1, horizontal_flip=True)
    # prepare iterator
    it_train = datagen.flow(trainX, trainY, batch_size=64)
    # fit model
    steps = int(trainX.shape[0] / 64)
    history = model.fit_generator(it_train, steps_per_epoch=steps, epochs=300, validation_data=(valX, valY), verbose=1)
    # evaluate model
    _, acc = model.evaluate(testX, testY, verbose=0)
    print('> %.3f' % (acc * 100.0))

    model.save(f'resources\\{modelname}.h5')
    summarize_diagnostics(history, modelname, acc)


if __name__ == '__main__':
    run_test_harness()

