
# coding: utf-8

# # CoDeepNEAT demo
# ## CS081 project checkpoint demo
# ### Harsha Uppli, Alan Zhao, Gabriel Meyer-Lee
#
# The following notebook demonstrates using CoDeepNEAT to solve CIFAR-10

# In[1]:


from keras.datasets import cifar10
from keras.utils.np_utils import to_categorical
from math import pi, floor
from random import random
from codeepneat import codeepneat, config, population, chromosome, genome, visualize
import pickle
import numpy as np
import keras
import sys
from keras.preprocessing.image import ImageDataGenerator
from keras.callbacks import ReduceLROnPlateau, CSVLogger, EarlyStopping, TensorBoard

batch_size = 128
epoch = 100

lr_reducer = ReduceLROnPlateau(factor=np.sqrt(0.1), cooldown=0, patience=5, min_lr=0.5e-6)
early_stopper = EarlyStopping(monitor='val_loss', min_delta=0.001, patience=20)
csv_logger = CSVLogger('cifar10.csv')
#tensor_board = TensorBoard(log_dir='./logs', histogram_freq=0, batch_size=batch_size, write_graph=True, write_grads=False, write_images=False, embeddings_freq=0, embeddings_layer_names=None, embeddings_metadata=None)


def fitness(network, data):
    network.fit(data[0], data[1],  epochs=8)
    loss, acc = network.evaluate(data[2], data[3])
    return acc

def main():
    (x_train_all, y_train_all), (x_test, y_test) = cifar10.load_data()
    y_train_all = y_train_all[:,0]
    y_test = y_test[:,0]

    x_train_all = np.reshape(x_train_all, (x_train_all.shape[0], 32, 32, 3)).astype('float32') / 255
    x_test = np.reshape(x_test, (x_test.shape[0], 32, 32, 3)).astype('float32') / 255
    y_train_all = keras.utils.np_utils.to_categorical(y_train_all)
    y_test = keras.utils.np_utils.to_categorical(y_test)
    data = [x_train_all, y_train_all, x_test, y_test]

    directory = "CIFARtest6/"
    for i in range(25):
        model_file = "CIFAR10_best_model_" + str(i)
        eval_best(directory+model_file, data)

    #model_file = "CIFAR10_best_model_17"
    #eval_best(directory+model_file, data)




def eval_best(model_file, data):
    config.load('configCIFAR10')
    model = keras.models.load_model(model_file)
    datagen = ImageDataGenerator(
        featurewise_center=False,  # set input mean to 0 over the dataset
        samplewise_center=False,  # set each sample mean to 0
        featurewise_std_normalization=False,  # divide inputs by std of the dataset
        samplewise_std_normalization=False,  # divide each input by its std
        zca_whitening=False,  # apply ZCA whitening
        rotation_range=10,  # randomly rotate images in the range (degrees, 0 to 180)
        width_shift_range=0.1,  # randomly shift images horizontally (fraction of total width)
        height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
        horizontal_flip=True,  # randomly flip images
        vertical_flip=False)  # randomly flip images

    # Compute quantities required for featurewise normalization
    # (std, mean, and principal components if ZCA whitening is applied).
    datagen.fit(data[0])
    csv_logger = CSVLogger(model_file + '_cifar10_new.csv')

    # Fit the model on the batches generated by datagen.flow().
    model.fit_generator(datagen.flow(data[0], data[1], batch_size=batch_size),
                        steps_per_epoch=data[0].shape[0] // batch_size,
                        validation_data=(data[2], data[3]),
                        epochs=epoch, verbose=1, max_queue_size=100,
                        callbacks=[lr_reducer, early_stopper, csv_logger])

    loss, fitness = model.evaluate(data[2], data[3])
    print("fitness", fitness)

if __name__=='__main__':
    main()
