import datetime
import os
import random as rn
from unittest import TestCase

import keras
import keras.backend as K
import numpy as np
import tensorflow as tf
from keras.callbacks import ModelCheckpoint, TensorBoard, LearningRateScheduler
from keras.datasets import cifar10
from keras.layers import Conv2D, Activation, GlobalAveragePooling2D, Dropout, PReLU
from keras.models import Sequential
from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator
from keras.utils import np_utils

from linked_neurons import LKReLU, LKPReLU, LKSwish, LKSELU, swish


class TestAllCnnSchedule(TestCase):
    def setUp(self):
        self.metrics = ['accuracy']
        self.allcnn_histogram_freq = 1
        self.epochs = 3000
        self.lr = 0.001
        # Set seed
        self.seed = None
        self.seed = 100
        os.environ['PYTHONHASHSEED'] = str(self.seed)
        np.random.seed(self.seed)
        rn.seed(self.seed)
        session_conf = tf.ConfigProto(intra_op_parallelism_threads=1, inter_op_parallelism_threads=1)
        tf.set_random_seed(self.seed)
        self.sess = tf.Session(graph=tf.get_default_graph(), config=session_conf)
        K.set_session(self.sess)

    def test_relu(self):
        batch_size = 32
        classes = 10

        (X_train, y_train), (X_test, y_test) = cifar10.load_data()
        print('X_train shape:', X_train.shape)
        print(X_train.shape[0], 'train samples')
        print(X_test.shape[0], 'test samples')

        print(X_train.shape[1:])

        Y_train = np_utils.to_categorical(y_train, classes)
        Y_test = np_utils.to_categorical(y_test, classes)

        model = Sequential()

        model.add(Conv2D(96, (3, 3), padding='same', input_shape=(32, 32, 3)))
        model.add(Activation('relu'))
        model.add(Conv2D(96, (3, 3), padding='same'))
        model.add(Activation('relu'))
        model.add(Conv2D(96, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation('relu'))
        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation('relu'))
        model.add(Conv2D(192, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation('relu'))
        model.add(Conv2D(192, (1, 1), padding='valid'))
        model.add(Activation('relu'))
        model.add(Conv2D(10, (1, 1), padding='valid'))

        model.add(GlobalAveragePooling2D())
        model.add(Activation('softmax'))
        sgd = SGD(lr=self.lr, decay=1e-6, momentum=0.9, nesterov=True)
        model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=self.metrics)

        print(model.summary())

        X_train = X_train.astype('float32')
        X_test = X_test.astype('float32')
        X_train /= 255
        X_test /= 255

        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=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total basicnet_width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)

        datagen.fit(X_train)
        filepath = 'checkpoints/allcnn_relu-cifar10-schedule-{epoch:02d}-{val_loss:.2f}.hdf5'
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                                     save_weights_only=False, mode='max')

        log_dir = 'summaries/allcnn_relu-cifar10-schedule-{}-{}-{}'.format(self.lr, self.seed,
                                                                            datetime.datetime.now())
        tensorboard = TensorBoard(log_dir=log_dir)
        lr_scheduler = LearningRateScheduler(schedule=self.schedule)
        # tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1, write_grads=True)
        callbacks_list = [checkpoint, tensorboard, lr_scheduler]
        # Fit the model on the batches generated by datagen.flow().
        history_callback = model.fit_generator(datagen.flow(X_train, Y_train,
                                                            batch_size=batch_size),
                                               epochs=self.epochs,
                                               validation_data=(X_test, Y_test),
                                               callbacks=callbacks_list,
                                               steps_per_epoch=int(X_train.shape[0] / batch_size),
                                               verbose=0)

    def test_lkrelu(self):
        batch_size = 128
        classes = 10

        (x_train, y_train), (x_test, y_test) = cifar10.load_data()
        # x_train = block_reduce(x_train, (1, 5, 5, 1), func=np.mean)
        # x_test = block_reduce(x_test, (1, 5, 5, 1), func=np.mean)

        x_train = x_train.astype('float32')
        x_test = x_test.astype('float32')
        x_train /= 255
        x_test /= 255
        y_train = np_utils.to_categorical(y_train, classes)
        y_test = np_utils.to_categorical(y_test, classes)

        model = Sequential()

        model.add(Conv2D(96, (3, 3), padding='same', input_shape=x_train.shape[1:]))
        model.add(LKReLU())
        model.add(Conv2D(96, (3, 3), padding='same'))
        model.add(LKReLU())
        model.add(Conv2D(96, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKReLU())
        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKReLU())
        model.add(Conv2D(192, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKReLU())
        model.add(Conv2D(192, (1, 1), padding='valid'))
        model.add(LKReLU())
        model.add(Conv2D(10, (1, 1), padding='valid'))

        model.add(GlobalAveragePooling2D())
        model.add(Activation('softmax'))
        opt = SGD(lr=self.lr, decay=1e-6, momentum=0.9, nesterov=True)

        model.compile(loss=keras.losses.categorical_crossentropy,
                      optimizer=opt,
                      metrics=self.metrics)

        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=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total basicnet_width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)

        datagen.fit(x_train)

        now = datetime.datetime.now()
        filepath = 'checkpoints/allcnn_lkrelu-cifar10-schedule-{epoch:02d}-%s.hdf5' % now
        checkpoint = ModelCheckpoint(filepath, monitor='acc', verbose=0, save_best_only=False,
                                     save_weights_only=False, mode='max', period=100)
        lr_scheduler = LearningRateScheduler(schedule=self.schedule)

        log_dir = 'summaries/allcnn_lkrelu-cifar10-schedule-{}-{}-{}'.format(self.lr, self.seed, now)
        tensorboard = TensorBoard(log_dir=log_dir)

        model.fit_generator(datagen.flow(x_train, y_train,
                                         batch_size=batch_size),
                            epochs=self.epochs,
                            steps_per_epoch=int(x_train.shape[0] / batch_size),
                            verbose=1,
                            validation_data=(x_test, y_test),
                            callbacks=[checkpoint, tensorboard, lr_scheduler])
        score = model.evaluate(x_test, y_test, verbose=0)
        print('Test loss:', score[0])
        print('Test accuracy:', score[1])

    def test_prelu(self):
        batch_size = 32
        classes = 10

        (X_train, y_train), (X_test, y_test) = cifar10.load_data()
        print('X_train shape:', X_train.shape)
        print(X_train.shape[0], 'train samples')
        print(X_test.shape[0], 'test samples')

        print(X_train.shape[1:])

        Y_train = np_utils.to_categorical(y_train, classes)
        Y_test = np_utils.to_categorical(y_test, classes)

        model = Sequential()

        model.add(Conv2D(96, (3, 3), padding='same', input_shape=(32, 32, 3)))
        model.add(PReLU())
        model.add(Conv2D(96, (3, 3), padding='same'))
        model.add(PReLU())
        model.add(Conv2D(96, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(PReLU())
        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(PReLU())
        model.add(Conv2D(192, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(PReLU())
        model.add(Conv2D(192, (1, 1), padding='valid'))
        model.add(PReLU())
        model.add(Conv2D(10, (1, 1), padding='valid'))

        model.add(GlobalAveragePooling2D())
        model.add(Activation('softmax'))
        sgd = SGD(lr=self.lr, decay=1e-6, momentum=0.9, nesterov=True)
        model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=self.metrics)

        print(model.summary())

        X_train = X_train.astype('float32')
        X_test = X_test.astype('float32')
        X_train /= 255
        X_test /= 255

        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=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total basicnet_width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)

        datagen.fit(X_train)
        filepath = 'checkpoints/allcnn_prelu-cifar10-schedule-{epoch:02d}-{val_loss:.2f}.hdf5'
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                                     save_weights_only=False, mode='max')

        log_dir = 'summaries/allcnn_prelu-cifar10-schedule-{}-{}-{}'.format(self.lr, self.seed, datetime.datetime.now())
        tensorboard = TensorBoard(log_dir=log_dir)
        lr_scheduler = LearningRateScheduler(schedule=self.schedule)
        # tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1, write_grads=True)
        callbacks_list = [checkpoint, tensorboard, lr_scheduler]
        # Fit the model on the batches generated by datagen.flow().
        history_callback = model.fit_generator(datagen.flow(X_train, Y_train,
                                                            batch_size=batch_size),
                                               epochs=self.epochs,
                                               validation_data=(X_test, Y_test),
                                               callbacks=callbacks_list,
                                               steps_per_epoch=int(X_train.shape[0] / batch_size),
                                               verbose=0)

    def test_selu(self):
        batch_size = 32
        classes = 10

        (X_train, y_train), (X_test, y_test) = cifar10.load_data()
        print('X_train shape:', X_train.shape)
        print(X_train.shape[0], 'train samples')
        print(X_test.shape[0], 'test samples')

        print(X_train.shape[1:])

        Y_train = np_utils.to_categorical(y_train, classes)
        Y_test = np_utils.to_categorical(y_test, classes)

        model = Sequential()

        model.add(Conv2D(96, (3, 3), padding='same', input_shape=(32, 32, 3)))
        model.add(Activation('selu'))
        model.add(Conv2D(96, (3, 3), padding='same'))
        model.add(Activation('selu'))
        model.add(Conv2D(96, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation('selu'))
        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation('selu'))
        model.add(Conv2D(192, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation('selu'))
        model.add(Conv2D(192, (1, 1), padding='valid'))
        model.add(Activation('selu'))
        model.add(Conv2D(10, (1, 1), padding='valid'))

        model.add(GlobalAveragePooling2D())
        model.add(Activation('softmax'))
        sgd = SGD(lr=self.lr, decay=1e-6, momentum=0.9, nesterov=True)
        model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=self.metrics)

        print(model.summary())

        X_train = X_train.astype('float32')
        X_test = X_test.astype('float32')
        X_train /= 255
        X_test /= 255

        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=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total basicnet_width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)

        datagen.fit(X_train)
        filepath = 'checkpoints/allcnn_selu-cifar10-schedule-{epoch:02d}-{val_loss:.2f}.hdf5'
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                                     save_weights_only=False, mode='max')

        log_dir = 'summaries/allcnn_selu-cifar10-schedule-{}-{}-{}'.format(self.lr, self.seed, datetime.datetime.now())
        tensorboard = TensorBoard(log_dir=log_dir)
        lr_scheduler = LearningRateScheduler(schedule=self.schedule)
        # tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1, write_grads=True)
        callbacks_list = [checkpoint, tensorboard, lr_scheduler]
        # Fit the model on the batches generated by datagen.flow().
        history_callback = model.fit_generator(datagen.flow(X_train, Y_train,
                                                            batch_size=batch_size),
                                               epochs=self.epochs,
                                               validation_data=(X_test, Y_test),
                                               callbacks=callbacks_list,
                                               steps_per_epoch=int(X_train.shape[0] / batch_size),
                                               verbose=0)

    def test_swish(self):
        batch_size = 32
        classes = 10

        (X_train, y_train), (X_test, y_test) = cifar10.load_data()
        print('X_train shape:', X_train.shape)
        print(X_train.shape[0], 'train samples')
        print(X_test.shape[0], 'test samples')

        print(X_train.shape[1:])

        Y_train = np_utils.to_categorical(y_train, classes)
        Y_test = np_utils.to_categorical(y_test, classes)

        model = Sequential()

        model.add(Conv2D(96, (3, 3), padding='same', input_shape=(32, 32, 3)))
        model.add(Activation(swish))
        model.add(Conv2D(96, (3, 3), padding='same'))
        model.add(Activation(swish))
        model.add(Conv2D(96, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation(swish))
        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation(swish))
        model.add(Conv2D(192, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(Activation(swish))
        model.add(Conv2D(192, (1, 1), padding='valid'))
        model.add(Activation(swish))
        model.add(Conv2D(10, (1, 1), padding='valid'))

        model.add(GlobalAveragePooling2D())
        model.add(Activation('softmax'))
        sgd = SGD(lr=self.lr, decay=1e-6, momentum=0.9, nesterov=True)
        model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=self.metrics)

        print(model.summary())

        X_train = X_train.astype('float32')
        X_test = X_test.astype('float32')
        X_train /= 255
        X_test /= 255

        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=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total basicnet_width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)

        datagen.fit(X_train)
        filepath = 'checkpoints/allcnn_swish-cifar10-schedule-{epoch:02d}-{val_loss:.2f}.hdf5'
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                                     save_weights_only=False, mode='max')

        log_dir = 'summaries/allcnn_swish-cifar10-schedule-{}-{}-{}'.format(self.lr, self.seed, datetime.datetime.now())
        tensorboard = TensorBoard(log_dir=log_dir)
        lr_scheduler = LearningRateScheduler(schedule=self.schedule)
        # tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1, write_grads=True)
        callbacks_list = [checkpoint, tensorboard, lr_scheduler]
        # Fit the model on the batches generated by datagen.flow().
        history_callback = model.fit_generator(datagen.flow(X_train, Y_train,
                                                            batch_size=batch_size),
                                               epochs=self.epochs,
                                               validation_data=(X_test, Y_test),
                                               callbacks=callbacks_list,
                                               steps_per_epoch=int(X_train.shape[0] / batch_size),
                                               verbose=0)

    def test_lkprelu(self):
        print('LR: {}'.format(self.lr))
        batch_size = 32
        classes = 10

        (X_train, y_train), (X_test, y_test) = cifar10.load_data()
        print('X_train shape:', X_train.shape)
        print(X_train.shape[0], 'train samples')
        print(X_test.shape[0], 'test samples')

        print(X_train.shape[1:])

        Y_train = np_utils.to_categorical(y_train, classes)
        Y_test = np_utils.to_categorical(y_test, classes)

        model = Sequential()

        model.add(Conv2D(96, (3, 3), padding='same', input_shape=(32, 32, 3)))
        model.add(LKPReLU())
        model.add(Conv2D(96, (3, 3), padding='same'))
        model.add(LKPReLU())
        model.add(Conv2D(96, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKPReLU())
        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKPReLU())
        model.add(Conv2D(192, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKPReLU())
        model.add(Conv2D(192, (1, 1), padding='valid'))
        model.add(LKPReLU())
        model.add(Conv2D(10, (1, 1), padding='valid'))

        model.add(GlobalAveragePooling2D())
        model.add(Activation('softmax'))
        sgd = SGD(lr=self.lr, decay=1e-6, momentum=0.9, nesterov=True)
        model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=self.metrics)

        print(model.summary())

        X_train = X_train.astype('float32')
        X_test = X_test.astype('float32')
        X_train /= 255
        X_test /= 255

        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=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total basicnet_width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)

        datagen.fit(X_train)
        filepath = 'checkpoints/allcnn_lkprelu-cifar10-schedule-{epoch:02d}-{val_loss:.2f}.hdf5'
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                                     save_weights_only=False, mode='max')

        log_dir = 'summaries/allcnn_lkprelu-cifar10-schedule-{}-{}-{}'.format(self.lr, self.seed,
                                                                           datetime.datetime.now())
        tensorboard = TensorBoard(log_dir=log_dir)
        lr_scheduler = LearningRateScheduler(schedule=self.schedule)
        # tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1, write_grads=True)
        callbacks_list = [checkpoint, tensorboard, lr_scheduler]
        # Fit the model on the batches generated by datagen.flow().
        history_callback = model.fit_generator(datagen.flow(X_train, Y_train,
                                                            batch_size=batch_size),
                                               epochs=self.epochs,
                                               validation_data=(X_test, Y_test),
                                               callbacks=callbacks_list,
                                               steps_per_epoch=int(X_train.shape[0] / batch_size),
                                               verbose=0)

    def test_lkselu(self):
        print('LR: {}'.format(self.lr))
        batch_size = 32
        classes = 10

        (X_train, y_train), (X_test, y_test) = cifar10.load_data()
        print('X_train shape:', X_train.shape)
        print(X_train.shape[0], 'train samples')
        print(X_test.shape[0], 'test samples')

        print(X_train.shape[1:])

        Y_train = np_utils.to_categorical(y_train, classes)
        Y_test = np_utils.to_categorical(y_test, classes)

        model = Sequential()

        model.add(Conv2D(96, (3, 3), padding='same', input_shape=(32, 32, 3)))
        model.add(LKSELU())
        model.add(Conv2D(96, (3, 3), padding='same'))
        model.add(LKSELU())
        model.add(Conv2D(96, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKSELU())
        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKSELU())
        model.add(Conv2D(192, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKSELU())
        model.add(Conv2D(192, (1, 1), padding='valid'))
        model.add(LKSELU())
        model.add(Conv2D(10, (1, 1), padding='valid'))

        model.add(GlobalAveragePooling2D())
        model.add(Activation('softmax'))
        sgd = SGD(lr=self.lr, decay=1e-6, momentum=0.9, nesterov=True)
        model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=self.metrics)

        print(model.summary())

        X_train = X_train.astype('float32')
        X_test = X_test.astype('float32')
        X_train /= 255
        X_test /= 255

        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=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total basicnet_width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)

        datagen.fit(X_train)
        filepath = 'checkpoints/allcnn_lkselu-cifar10-schedule-{epoch:02d}-{val_loss:.2f}.hdf5'
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                                     save_weights_only=False, mode='max')

        log_dir = 'summaries/allcnn_lkselu-cifar10-schedule-{}-{}-{}'.format(self.lr, self.seed,
                                                                          datetime.datetime.now())
        tensorboard = TensorBoard(log_dir=log_dir)
        lr_scheduler = LearningRateScheduler(schedule=self.schedule)
        # tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1, write_grads=True)
        callbacks_list = [checkpoint, tensorboard, lr_scheduler]
        # Fit the model on the batches generated by datagen.flow().
        history_callback = model.fit_generator(datagen.flow(X_train, Y_train,
                                                            batch_size=batch_size),
                                               epochs=self.epochs,
                                               validation_data=(X_test, Y_test),
                                               callbacks=callbacks_list,
                                               steps_per_epoch=int(X_train.shape[0] / batch_size),
                                               verbose=0)

    def test_lkswish(self):
        batch_size = 32
        classes = 10

        (X_train, y_train), (X_test, y_test) = cifar10.load_data()
        print('X_train shape:', X_train.shape)
        print(X_train.shape[0], 'train samples')
        print(X_test.shape[0], 'test samples')

        print(X_train.shape[1:])

        Y_train = np_utils.to_categorical(y_train, classes)
        Y_test = np_utils.to_categorical(y_test, classes)

        model = Sequential()

        model.add(Conv2D(96, (3, 3), padding='same', input_shape=(32, 32, 3)))
        model.add(LKSwish())
        model.add(Conv2D(96, (3, 3), padding='same'))
        model.add(LKSwish())
        model.add(Conv2D(96, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKSwish())
        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKSwish())
        model.add(Conv2D(192, (3, 3), padding='same', strides=(2, 2)))
        model.add(Dropout(0.5))

        model.add(Conv2D(192, (3, 3), padding='same'))
        model.add(LKSwish())
        model.add(Conv2D(192, (1, 1), padding='valid'))
        model.add(LKSwish())
        model.add(Conv2D(10, (1, 1), padding='valid'))

        model.add(GlobalAveragePooling2D())
        model.add(Activation('softmax'))
        sgd = SGD(lr=self.lr, decay=1e-6, momentum=0.9, nesterov=True)
        model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=self.metrics)

        print(model.summary())

        X_train = X_train.astype('float32')
        X_test = X_test.astype('float32')
        X_train /= 255
        X_test /= 255

        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=0,  # randomly rotate images in the range (degrees, 0 to 180)
            width_shift_range=0.1,  # randomly shift images horizontally (fraction of total basicnet_width)
            height_shift_range=0.1,  # randomly shift images vertically (fraction of total height)
            horizontal_flip=True,  # randomly flip images
            vertical_flip=False)

        datagen.fit(X_train)
        filepath = 'checkpoints/allcnn_lkswish-cifar10-schedule-{epoch:02d}-{val_loss:.2f}.hdf5'
        checkpoint = ModelCheckpoint(filepath, monitor='val_acc', verbose=1, save_best_only=True,
                                     save_weights_only=False, mode='max')

        log_dir = 'summaries/allcnn_lkswish-cifar10-schedule-{}-{}-{}'.format(self.lr, self.seed,
                                                                           datetime.datetime.now())
        tensorboard = TensorBoard(log_dir=log_dir)
        lr_scheduler = LearningRateScheduler(schedule=self.schedule)
        # tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1, write_grads=True)
        callbacks_list = [checkpoint, tensorboard, lr_scheduler]
        # Fit the model on the batches generated by datagen.flow().
        history_callback = model.fit_generator(datagen.flow(X_train, Y_train,
                                                            batch_size=batch_size),
                                               epochs=self.epochs,
                                               validation_data=(X_test, Y_test),
                                               callbacks=callbacks_list,
                                               steps_per_epoch=int(X_train.shape[0] / batch_size),
                                               verbose=0)

    def schedule(self, epoch):
        if epoch < 500:
            return self.lr
        elif epoch < 2000:
            return self.lr * 0.75
        elif epoch < 2350:
            return self.lr * 0.5
        elif epoch < 2750:
            return self.lr * 0.1
        else:
            return self.lr * 0.05
