from __future__ import print_function
from keras.preprocessing.image import ImageDataGenerator
from keras.utils import np_utils
from keras.callbacks import ReduceLROnPlateau, CSVLogger, EarlyStopping
from keras.optimizers import Adam
from keras.losses import categorical_crossentropy
from keras.models import load_model

import numpy as np

import resnet
from custom.datasets import wormdata
from custom.losses import w_categorical_crossentropy
from custom.utils import confusion_mat
from custom.augmentation import to_PIL, to_array, norm, random_brightness, random_rotation, compose

lr_reducer = ReduceLROnPlateau(factor=np.sqrt(0.1), cooldown=0, patience=5, min_lr=0.5e-6)
early_stopper = EarlyStopping(min_delta=0.001, patience=10)
csv_logger = CSVLogger('resnet18_worm.csv')

batch_size = 32
nb_classes = 5
nb_epoch = 50
data_augmentation = True

# input image dimensions
img_rows, img_cols = 224, 224
# The CIFAR10 images are RGB.
img_channels = 3

# The data, shuffled and split between train and test sets:
# (X_train, y_train), (X_test, y_test) = cifar10.load_data()
data = wormdata(filepath='./data.mat', interval=10)
X_train, y_train = data.get_input(tag='train'), data.get_target(tag='train')
X_test, y_test = data.get_input(tag='test'), data.get_target(tag='test')
X_test = X_test.transpose((0,2,3,1))
X_train = X_train.transpose((0,2,3,1))

# Convert class vectors to binary class matrices.
# Y_train = np_utils.to_categorical(y_train, nb_classes)
# Y_test = np_utils.to_categorical(y_test, nb_classes)

X_train = np.float32(X_train)
X_test = np.float32(X_test)

# subtract mean and normalize
mean_image = np.mean(X_train, axis=0)
# X_train -= mean_image
X_test -= mean_image
# X_train /= 128.
X_test /= 128.

model = resnet.ResnetBuilder.build_resnet_18((img_channels, img_rows, img_cols), nb_classes)
# model = load_model('res18.h5', custom_objects={'w_categorical_crossentropy': w_categorical_crossentropy})

def random_brightness(img):
    img = Image.fromarray(np.uint8(img), 'RGB')
    factor = np.random.uniform(0.8, 1.2)
    img = ImageEnhance.Brightness(img).enhance(factor)
    img = np.array(img)
    img = np.float32(img) - mean_image
    img /= 128.
    return img
#
custom_transform = compose(to_PIL, random_rotation, random_brightness, to_array, norm(mean_image, 128.))

model.compile(loss=w_categorical_crossentropy,
              optimizer=Adam(lr=1e-3),
              metrics=['accuracy'])

if data_augmentation:
    print('Using real-time data augmentation.')
    # This will do preprocessing and realtime data augmentation:
    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)
        channel_shift_range=0.1,  # randomly shift channels
        horizontal_flip=True,  # randomly flip images
        vertical_flip=True,  # randomly flip images
        preprocessing_function=random_brightness)  # randomly shift brightness

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

    # Fit the model on the batches generated by datagen.flow().
    model.fit_generator(datagen.flow(X_train, y_train, batch_size=batch_size),
                        steps_per_epoch=X_train.shape[0] // batch_size,
                        validation_data=(X_test, y_test),
                        epochs=nb_epoch, verbose=1, max_q_size=100,
                        callbacks=[lr_reducer, early_stopper, csv_logger])

model.save('res18.h5')
confusion_mat(model, X_test, y_test)
