# coding=utf-8

"""
Train a (fairly simple) deep CNN on the CIFAR10 small images dataset.

GPU run command:
THEANO_FLAGS=mode=FAST_RUN,device=gpu,floatX=float32 python cifar10_cnn.py

Gets to about 0.5 test loss or 83% accuracy after 65 epochs.

"""

from __future__ import absolute_import
from __future__ import print_function

from keras.callbacks import ModelCheckpoint
from keras.datasets import cifar10
from keras.layers import BatchNormalization, MaxPooling2D, GaussianNoise
from keras.layers import Dense, Activation, Flatten, Conv2D
from keras.models import Sequential
from keras.preprocessing.image import ImageDataGenerator
from keras.utils import np_utils

from snntoolbox.simulation.plotting import plot_history

batch_size = 32
nb_epoch = 100

sigma = 0.01

# Data set
(X_train, y_train), (X_test, y_test) = cifar10.load_data()
Y_train = np_utils.to_categorical(y_train, 10)
Y_test = np_utils.to_categorical(y_test, 10)

model = Sequential()

model.add(GaussianNoise(sigma, input_shape=(3, 32, 32)))
model.add(Conv2D(32, 3, 3, padding='same'))
model.add(BatchNormalization(axis=1))
model.add(Activation('relu'))
model.add(GaussianNoise(sigma))
model.add(Conv2D(32, 3, 3, padding='same'))
model.add(BatchNormalization(axis=1))
model.add(Activation('relu'))
model.add(GaussianNoise(sigma))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Conv2D(64, 3, 3))
model.add(BatchNormalization(axis=1))
model.add(Activation('relu'))
model.add(GaussianNoise(sigma))
model.add(Conv2D(64, 3, 3))
model.add(BatchNormalization(axis=1))
model.add(Activation('relu'))
model.add(GaussianNoise(sigma))
model.add(MaxPooling2D(pool_size=(2, 2)))

model.add(Flatten())
model.add(Dense(512))
model.add(BatchNormalization())
model.add(Activation('relu'))
model.add(GaussianNoise(sigma))
model.add(Dense(10))
model.add(Activation('softmax'))

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

# Whether to apply global contrast normalization and ZCA whitening
gcn = False
zca = False

traingen = ImageDataGenerator(rescale=1./255, featurewise_center=gcn,
                              featurewise_std_normalization=gcn,
                              zca_whitening=zca, horizontal_flip=True,
                              rotation_range=10, width_shift_range=0.1,
                              height_shift_range=0.1)

# Compute quantities required for featurewise normalization
# (std, mean, and principal components if ZCA whitening is applied)
traingen.fit(X_train/255.)

trainflow = traingen.flow(X_train, Y_train, batch_size=batch_size)

testgen = ImageDataGenerator(rescale=1./255, featurewise_center=gcn,
                             featurewise_std_normalization=gcn,
                             zca_whitening=zca)

testgen.fit(X_test/255.)

testflow = testgen.flow(X_test, Y_test, batch_size=batch_size)

checkpointer = ModelCheckpoint(filepath='cnn_BN.{epoch:02d}-{val_acc:.2f}.h5',
                               verbose=1, save_best_only=True)

# Fit the model on the batches generated by datagen.flow()
history = model.fit_generator(trainflow, nb_epoch=nb_epoch,
                              samples_per_epoch=X_train.shape[0],
                              validation_data=testflow,
                              nb_val_samples=len(X_test),
                              callbacks=[checkpointer])
plot_history(history)

score = model.evaluate_generator(testflow, val_samples=len(X_test))
print('Test score:', score[0])
print('Test accuracy:', score[1])

model.save('{:2.2f}.h5'.format(score[1]*100))
