from keras.models import Sequential
from keras.layers import  Dropout, Conv2D,MaxPooling2D, Flatten, Dense
from keras_preprocessing.image import ImageDataGenerator
from sklearn.metrics import classification_report, accuracy_score,confusion_matrix
import matplotlib.pyplot as plt
import keras_metrics as km
import numpy as np
import os
import split_folders
import shutil



seedNumbers = [2,3,5,7,11]
batchSize = 32
test_dir_name = '/Users/haikristianlethanh/Desktop/test/val'
train_dir_name = '/Users/haikristianlethanh/Desktop/test/train'
path_to_dataset = '/Users/haikristianlethanh/Desktop/CNN-keras-deep-learning/data_untouch'
splited_dir_path = '/Users/haikristianlethanh/Desktop/test'
presnost = []

downloalded = '/Users/haikristianlethanh/Desktop/CNN-keras-deep-learning/Test_downloaded/COMP/'
reference_FRI = '/Users/haikristianlethanh/Desktop/FIRST_Data/113.98142_R165.jpg'
transformed_data_target ='/Users/haikristianlethanh/Desktop/CNN-keras-deep-learning/GREY/COMP/'


def trainingLoss(history):
    loss = history.history['loss']
    epochs = range(1, len(loss) + 1)
    plt.plot(epochs, loss, color='red', label='Training loss')
    plt.title('Training  loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    plt.show()

def metrics(test_set, classifier):
    test_set.reset()
    numfiles = sum([len(files) for r, d, files in os.walk(test_dir_name)])
    print(numfiles)
    Y_pred = classifier.predict_generator(test_set,steps=(numfiles // 32) + 1)
    print(Y_pred.shape)
    classes = test_set.classes[test_set.index_array]
    print(classes)
    test_set.classes
    y_pred = np.argmax(Y_pred, axis=-1)
    accuracy = accuracy_score(y_pred, test_set.classes)
    print("Accuracy score: ", accuracy)
    y_pred
    cm = confusion_matrix(test_set.classes, y_pred)
    print(cm)
    print(classification_report(test_set.classes, y_pred,target_names=test_set.class_indices.keys()))
    return accuracy

def initializeModel():
    classifier =  Sequential()
    classifier.add(Conv2D(32, (11, 11), input_shape=(150,150,3), activation='relu'))
    classifier.add(MaxPooling2D(pool_size=(2,2)))
    classifier.add(Conv2D(64, (3, 3),input_shape=(150,150,3), activation='relu'))
    classifier.add(MaxPooling2D(pool_size=(2,2)))
    classifier.add(Conv2D(194, (3, 3),input_shape=(150,150,3), activation='relu'))
    classifier.add(MaxPooling2D(pool_size=(2,2)))
    classifier.add(Dropout(0.2))
    classifier.add(Flatten())
    classifier.add(Dense(units = 194, activation='relu'))
    classifier.add(Dense(units = 4, activation='softmax'))
    recall = km.binary_recall(label=0)
    precision = km.binary_precision(label=1)
    c_precision = km.categorical_precision()
    classifier.compile(optimizer= 'Adam', loss = 'categorical_crossentropy', metrics = [c_precision, 'accuracy', recall, precision])
    classifier.summary()
    classifier.metrics_names
    return classifier


def main():
    for i in seedNumbers:
        split_folders.ratio(path_to_dataset, output=splited_dir_path, seed=i, ratio=(.8, .2)) # default values
        classifier = initializeModel()
        numfiles = sum([len(files) for r, d, files in os.walk(path_to_dataset)])
        train_datagen_1 = ImageDataGenerator(brightness_range=[1,1.5])
        train_datagen_2 = ImageDataGenerator(horizontal_flip=True,vertical_flip=True)
        training_set_1 = train_datagen_1.flow_from_directory(train_dir_name,target_size=(150,150),batch_size=batchSize,class_mode='categorical')
        training_set_2 = train_datagen_2.flow_from_directory(train_dir_name,target_size=(150,150),batch_size=batchSize,class_mode='categorical')
        test_datagen = ImageDataGenerator()
        test_set = test_datagen.flow_from_directory(test_dir_name,
                                                    target_size= (150,150),
                                                    batch_size=batchSize,
                                                    class_mode='categorical',
                                                    shuffle=False)
        history = classifier.fit_generator(training_set_1, steps_per_epoch=numfiles/batchSize, nb_epoch=20)
        history = classifier.fit_generator(training_set_2, steps_per_epoch=numfiles/batchSize, nb_epoch=20)
        trainingLoss(history)
        presnost.append(metrics(test_set, classifier))
        shutil.rmtree(splited_dir_path)
        print('Removed')