import tensorflow as tf
from keras.layers import *
from keras.models import Model
from keras.utils import plot_model

import numpy as np

import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import skimage
from skimage import transform


data_path  = 'pictures/'
img_width  = 1400
img_height = 1400
img_depth  = 3


# see: https://keras.io/preprocessing/image/
def dataAugmentation(dataToAugment):
    arrayToFill = []

    # faster computation with values between 0 and 1 ?
    dataToAugment = np.divide(dataToAugment, 255.)

    # TODO: switch from RGB channels to CbCrY

    # adding the normal images   (8)
    for i in range(len(dataToAugment)):
        arrayToFill.append(dataToAugment[i])
    # vertical axis flip         (-> 16)
    for i in range(len(arrayToFill)):
        arrayToFill.append(np.fliplr(arrayToFill[i]))
    # horizontal axis flip       (-> 32)
    for i in range(len(arrayToFill)):
        arrayToFill.append(np.flipud(arrayToFill[i]))

    # downsizing by scale of 4   (-> 64 images of 350x350x3)
    for i in range(len(arrayToFill)):
        arrayToFill.append(skimage.transform.resize(
            arrayToFill[i],
            (img_width/4, img_height/4),
            mode='reflect',
            anti_aliasing=True))

    # # Sanity check: display the images
    # plt.figure(figsize=(10, 10))
    # for i in range(64):
    #     plt.subplot(8, 8, i + 1)
    #     plt.imshow(arrayToFill[i], cmap=plt.cm.binary)
    # plt.show()

    return arrayToFill


def setUpImages():

    # Setting up paths
    path1  = data_path + 'Moi.jpg'
    path2  = data_path + 'ASaucerfulOfSecrets.jpg'
    path3  = data_path + 'AtomHeartMother.jpg'
    path4  = data_path + 'Animals.jpg'
    path5  = data_path + 'DivisionBell.jpg'          # validator
    path6  = data_path + 'lighter.jpg'
    path7  = data_path + 'Meddle.jpg'                # validator
    path8  = data_path + 'ObscuredByClouds.jpg'      # validator
    path9  = data_path + 'TheDarkSideOfTheMoon.jpg'
    path10 = data_path + 'TheWall.jpg'
    path11 = data_path + 'WishYouWereHere.jpg'

    # Extracting images (1400x1400)
    train = [mpimg.imread(path1),
             mpimg.imread(path2),
             mpimg.imread(path3),
             mpimg.imread(path4),
             mpimg.imread(path6),
             mpimg.imread(path9),
             mpimg.imread(path10),
             mpimg.imread(path11)]
    finalTest = [mpimg.imread(path5),
                 mpimg.imread(path8),
                 mpimg.imread(path7)]

    # Augmenting data
    trainData = dataAugmentation(train)
    testData  = dataAugmentation(finalTest)

    setUpData(trainData, testData)


def setUpData(trainData, testData):

    print(type(trainData))  # <class 'list'>
    print(len(trainData))  # 64
    print(type(trainData[0]))  # <class 'numpy.ndarray'>
    print(trainData[0].shape)  # (1400, 1400, 3)
    print(trainData[len(trainData)//2].shape)  # (350, 350, 3)

    # (num_img, width, height, channels=3)

    # TODO: substract mean of all images to all images

    # Separating the training data
    validateData = trainData[:len(trainData)//2]    # First half is the unaltered data
    trainingData = trainData[len(trainData)//2:]    # Second half is the deteriorated data

    validateTestData = testData[:len(testData)//2]  # First half is the unaltered data
    trainingTestData = testData[len(testData)//2:]  # Second half is the deteriorated data

    # # Sanity checks
    # print(type(validateData))  # <class 'list'>
    # print(len(validateData))  # 32
    # print(type(validateData[0].shape))  # <class 'tuple'>
    # print(validateData[0].shape)  # (1400, 1400, 3)
    #
    # print(type(trainingData))  # <class 'list'>
    # print(len(trainingData))  # 32
    # print(type(trainingData[0].shape))  # <class 'tuple'>
    # print(trainingData[0].shape)  # (350, 350, 3)

    # # Sanity check: display four images (2x HR/LR) TODO: retest (changed architecture from np to list)
    # plt.figure(figsize=(10, 10))
    # for i in range(2):
    #     plt.subplot(2, 2, i + 1)
    #     plt.imshow(validateData[i], cmap=plt.cm.binary)
    # for i in range(2):
    #     plt.subplot(2, 2, i + 1 + 2)
    #     plt.imshow(trainingData[i], cmap=plt.cm.binary)
    # plt.show()

    setUpModel(validateData, trainingData, validateTestData, trainingTestData)


def setUpModel(validateData, trainingData, validateTestData, trainingTestData):

    # # exemple de merge de deux networks: merge = concatenate([network1, network2])
    # # exemple de deux inputs pour un seul model: model = Model(inputs=[visible1, visible2], outputs=output)

    filters = 256
    kernel_size = 3
    strides = 1
    factor = 4  # the factor of upscaling

    inputLayer = Input(shape=(img_height/factor, img_width/factor, img_depth))
    conv1 = Conv2D(filters, kernel_size, strides=strides, padding='same')(inputLayer)

    res = Conv2D(filters, kernel_size, strides=strides, padding='same')(conv1)
    act = ReLU()(res)
    res = Conv2D(filters, kernel_size, strides=strides, padding='same')(act)
    res_rec = Add()([conv1, res])

    for i in range(15):  # 16-1
        res1 = Conv2D(filters, kernel_size, strides=strides, padding='same')(res_rec)
        act = ReLU()(res1)
        res2 = Conv2D(filters, kernel_size, strides=strides, padding='same')(act)
        res_rec = Add()([res_rec, res2])

    conv2 = Conv2D(filters, kernel_size, strides=strides, padding='same')(res_rec)
    a = Add()([conv1, conv2])
    up = UpSampling2D(size=4)(a)
    outputLayer = Conv2D(filters=3,
                         kernel_size=1,
                         strides=1,
                         padding='same')(up)

    model = Model(inputs=inputLayer, outputs=outputLayer)

    # Sanity checks
    print(model.summary())
    plot_model(model, to_file='CNN_graph.png')

    train(model, validateData, trainingData, validateTestData, trainingTestData)


def train(model, validateData, trainingData, validateTestData, trainingTestData):
    model.compile(optimizer=tf.train.AdamOptimizer(),
                  loss='sparse_categorical_crossentropy',  # TODO: Customize loss function?
                  metrics=['accuracy'])
    # TODO: possibly need to be LISTS instead of np.array ?
    model.fit(trainingData,
              validateData,
              epochs=5,
              verbose=2,
              batch_size=4)  # 32 images-> 8 batches of 4       TODO is it multi-fold testing?

    # Now use the TEST dataset to calculate performance
    test_loss, test_acc = model.evaluate(trainingTestData, validateTestData)
    print('Test accuracy:', test_acc)


    # TODO: eventually look into https://blog.keras.io/building-powerful-image-classification-models-using-very-little-data.html

    ###########################
    #      PREDICTIONS        #
    ###########################

    # Trying to make predictions on a single image
    predictions = model.predict(trainingTestData)

    # "model.predict" works in batches, so extracting a single prediction:
    img = trainingTestData[0]                   # Grab an image from the test dataset
    img = (np.expand_dims(img, 0))              # Add the image to a batch where it's the only member.
    predictions_single = model.predict(img)     # returns a list of lists, one for each image in the batch of data
    print(predictions_single)

    ###########################
    #        DRAWINGS         #
    ###########################

    # def plot_image(i, predictions_array, true_label, img):
    #     predictions_array, true_label, img = predictions_array[i], true_label[i], img[i]
    #     plt.grid(False)
    #     plt.xticks([])
    #     plt.yticks([])
    #
    #     plt.imshow(img, cmap=plt.cm.binary)
    #
    #     predicted_label = np.argmax(predictions_array)
    #     if predicted_label == true_label:
    #         color = 'blue'
    #     else:
    #         color = 'red'
    #
    #     plt.xlabel("{} {:2.0f}% ({})".format(class_names[predicted_label],
    #                                          100 * np.max(predictions_array),
    #                                          class_names[true_label]),
    #                color=color)
    #
    # def plot_value_array(i, predictions_array, true_label):
    #     predictions_array, true_label = predictions_array[i], true_label[i]
    #     plt.grid(False)
    #     plt.xticks([])
    #     plt.yticks([])
    #     thisplot = plt.bar(range(10), predictions_array, color="#777777")
    #     plt.ylim([0, 1])
    #     predicted_label = np.argmax(predictions_array)
    #     plt.xticks(range(10))  # adding the class-index below prediction graph
    #
    #     thisplot[predicted_label].set_color('red')
    #     thisplot[true_label].set_color('blue')
    #
    # # def draw_prediction(index):
    # #     plt.figure(figsize=(6, 3))
    # #     plt.subplot(1, 2, 1)
    # #     plot_image(index, predictions, test_labels, test_images)
    # #     plt.subplot(1, 2, 2)
    # #     plot_value_array(index, predictions, test_labels)
    # #     plt.show()
    # #
    # # To draw a single prediction
    # # draw_prediction(0)
    # # draw_prediction(12)
    #
    # # Plot the first X test images, their predicted label, and the true label
    # # Color correct predictions in blue, incorrect predictions in red
    # num_rows = 5
    # num_cols = 3
    # num_images = num_rows * num_cols
    # plt.figure(figsize=(2 * 2 * num_cols, 2 * num_rows))

    # Adding a title to the plot
    plt.suptitle("Check it out!")

    # for i in range(num_images):
    #     plt.subplot(num_rows, 2 * num_cols, 2 * i + 1)
    #     plot_image(i, predictions, test_labels, test_images)
    #     plt.subplot(num_rows, 2 * num_cols, 2 * i + 2)
    #     plot_value_array(i, predictions, test_labels)
    # plt.show()


if __name__ == '__main__':
    setUpImages()