import sys
import os
import numpy as np
import cv2
from random import shuffle
from keras.models import Sequential
from keras.layers.convolutional import Conv3D
from keras.layers.convolutional_recurrent import ConvLSTM2D
from keras.layers.normalization import BatchNormalization
from keras.layers import Dense, Flatten, TimeDistributed, Dropout, MaxPooling2D

from keras.optimizers import adadelta


# Function to generate a list of files to be processed
def generate_movies_from_list(n_samples=100, n_frames=9, channel_number=1, stage='train'):

    row = 45
    col = 45
    x = np.zeros((n_samples, n_frames, row, col, channel_number), dtype=np.float)
    y = np.zeros((n_samples, n_frames, row, col, 1),
                 dtype=np.float)

    file_list = np.genfromtxt(stage + '_dataset.txt', dtype='str')
    # file_list = file_list[::-1]

    if stage == 'test':
        label = np.zeros(n_samples, dtype=object)
    else:
        shuffle(file_list)  # If we are training, we want to Shuffle order of file list but if testing, we want to keep
        # the same order for better comparision

    current_dir = os.getcwd()

    if channel_number == 1:
        os.chdir('/home/gcx/lstm_sequences/entire_dataset/' + stage + '/predictions_' + stage + '/')
    else:
        os.chdir('/home/gcx/lstm_sequences/entire_dataset/' + stage + '/multi_channel_prediction/')

    if len(file_list)< n_samples:
        sys.exit('not enough sample in the designated set. Exiting script!')

    # Load samples for X
    for i in range(n_samples):

        current_file_list = file_list[i, :]

        if len(current_file_list) < n_frames:
            sys.exit('not enough samples per time instant . Exiting script!')
        else:
            for j in range(n_frames):
                if channel_number == 1: # If we are considering gray-scale images
                    try:
                        temp_img = cv2.cvtColor(cv2.imread(current_file_list[j]), cv2.COLOR_BGR2GRAY)
                        temp_img /= 255
                        x[i, j , :, :, 0] = temp_img
                    except:
                        sys.exit('could not load image: ')

                elif channel_number > 3:  # If we consider multi-channel images
                    try:
                        temp_img = np.load(current_file_list[j][:-4] + '.npy')
                        temp_img = temp_img / 255.0
                        x[i, j, :, :, :] = temp_img
                    except:
                        sys.exit('could not load npy: ')

                else:
                    sys.exit('incorrect number of channels in X!!')

    # Load data for Y
    os.chdir('/home/gcx/lstm_sequences/entire_dataset/' + stage + '/y_' + stage + '/')
    if len(file_list) < n_samples:
        sys.exit('not enough sample in the designated set. Exiting script!')

    for i in range(n_samples):

        current_file_list = file_list[i, :]

        if len(current_file_list) < n_frames:
            sys.exit('not enough samples per time instant . Exiting script!')
        else:
            for j in range(n_frames):
                if y.shape[4] == 1:
                    try:
                        temp_img = cv2.cvtColor(cv2.imread(current_file_list[j][:-4] + '.tiff'), cv2.COLOR_BGR2GRAY)
                        temp_img /= 255
                        y[i, j, :, :, 0] = temp_img

                        if stage == 'test' and j == n_frames-1 :
                            file_name = current_file_list[j][:-4] + '.tiff'
                            label[i] = file_name

                    except:
                        sys.exit('could not load image: ')
                else:
                    sys.exit('incorrect number of channels in y!!')

    os.chdir(current_dir)

    if stage == 'test':
        return x, y, label
    else:
        return x, y


# Function to generate a list of files to be processed
def generate_sequences_batch_from_list(start_instant, stop_instant, img_size, map_size,  n_frames=9, channel_number=3, stage='test'):

    n_samples = stop_instant - start_instant

    x = np.zeros((n_samples, n_frames, img_size, img_size, channel_number), dtype=np.float)
    y = np.zeros((n_samples, n_frames, map_size, map_size, 1),
                 dtype=np.float)

    # file_list = np.genfromtxt('run_' + stage + '_dataset.txt', dtype='str')
    file_list = np.genfromtxt(stage + '_dataset_20.txt', dtype='str')
    # file_list = file_list[::-1]

    if stage == 'test':
        label = np.zeros(n_samples, dtype=object)
    else:
        shuffle(
            file_list)  # If we are training, we want to Shuffle order of file list but if testing, we want to keep
        # the same order for better comparision

    current_dir = os.getcwd()

    if channel_number == 1:
        os.chdir('/home/gcx/lstm_sequences/entire_dataset/' + stage + '/predictions_' + stage + '/')
    elif channel_number == 3:
        os.chdir('/home/gcx/lstm_sequences/entire_dataset/' + stage + '/x_' + stage + '/')
    else:
        os.chdir('/home/gcx/lstm_sequences/entire_dataset/' + stage + '/multi_channel_prediction/')

    if len(file_list) < n_samples:
        sys.exit('not enough samples in the designated set. Exiting script!')

    # Load samples for X
    for i in range(n_samples):

        current_file_list = file_list[i + start_instant, :]

        if len(current_file_list) < n_frames:
            sys.exit('not enough samples per time instant . Exiting script!')
        else:
            for j in range(n_frames):
                if channel_number == 1:  # If we are considering gray-scale images
                    try:
                        temp_img = cv2.cvtColor(cv2.imread(current_file_list[j]), cv2.COLOR_BGR2GRAY)
                        temp_img /= 255
                        x[i, j, :, :, 0] = temp_img
                    except:
                        sys.exit('could not load image: ')

                elif channel_number == 3:
                    try:
                        temp_img = cv2.imread(current_file_list[j])
                        temp_img = cv2.resize(temp_img, (img_size, img_size))
                        temp_img = temp_img / 255.

                        x[i, j, :, :, :] = temp_img
                    except:
                        sys.exit('error loading rgb image')

                elif channel_number > 3:  # If we consider multi-channel images
                    try:
                        temp_img = np.load(current_file_list[j][:-4] + '.npy')
                        temp_img = temp_img / 255.0
                        x[i, j, :, :, :] = temp_img
                    except:
                        sys.exit('could not load npy: ')

                else:
                    sys.exit('incorrect number of channels in X!!')

    # Load data for Y
    os.chdir('/home/gcx/lstm_sequences/entire_dataset/' + stage + '/y_' + stage + '/')
    if len(file_list) < n_samples:
        sys.exit('not enough sample in the designated set. Exiting script!')

    for i in range(n_samples):

        current_file_list = file_list[i+start_instant, :]

        if len(current_file_list) < n_frames:
            sys.exit('not enough samples per time instant . Exiting script!')
        else:
            for j in range(n_frames):
                if y.shape[4] == 1:
                    try:
                        temp_img = cv2.cvtColor(cv2.imread(current_file_list[j][:-4] + '.tiff'), cv2.COLOR_BGR2GRAY)
                        temp_img /= 255
                        temp_img = cv2.resize(temp_img, (map_size, map_size), cv2.INTER_NEAREST)
                        y[i, j, :, :, 0] = temp_img

                        if stage == 'test' and j == n_frames - 1:
                            file_name = current_file_list[j][:-4] + '.tiff'
                            label[i] = file_name

                    except:
                        sys.exit('could not load gt image: ')
                else:
                    sys.exit('incorrect number of channels in y!!')

    os.chdir(current_dir)

    if stage == 'test':
        return x, y, label, len(file_list)
    else:
        return x, y


#  function to return the lstm network  for multi-channel data
def network_definition_multichannel(channel_number, reduced_size):
    seq = Sequential()
    seq.add(ConvLSTM2D(filters=int(channel_number/4.), kernel_size=(3, 3),
                       input_shape=(None, reduced_size, reduced_size, channel_number),
                       padding='same', return_sequences=True))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=int(channel_number/8.), kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=int(channel_number/16.), kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=int(channel_number/32.), kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    # seq.add(BatchNormalization())

    seq.add(Conv3D(filters=1, kernel_size=(3, 3, 3),
                   activation='sigmoid',
                   padding='same', data_format='channels_last'))

    # optimizer = adadelta(lr=0.1)
    # seq.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])
    seq.compile(loss='binary_crossentropy', optimizer='adadelta', metrics=['accuracy'])

    return seq


#  function to return the lstm network  for multi-channel data
def network_definition_rgb(channel_number, full_image_size,  reduced_size):
    # this was the network definition used for mae and bce examples, sent to Alex

    # Test 01
    seq = Sequential()
    seq.add(ConvLSTM2D(filters=8, kernel_size=(3, 3),
                       input_shape=(None, full_image_size, full_image_size, channel_number),
                       padding='same', return_sequences=True))
    seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=16, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=32, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=16, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))


    # seq.add(BatchNormalization())

    seq.add(Conv3D(filters=1, kernel_size=(3, 3, 3),
                   activation='sigmoid',
                   padding='same', data_format='channels_last'))

    # TEST 06
    # seq = Sequential()
    # seq.add(ConvLSTM2D(filters=8, kernel_size=(3, 3),
    #                    input_shape=(None, full_image_size, full_image_size, channel_number),
    #                    padding='same', return_sequences=True))
    # seq.add(TimeDistributed( MaxPooling2D((2,2), padding = 'same')))  # 720 -> 360
    # seq.add(BatchNormalization())
    #
    # seq.add(ConvLSTM2D(filters=16, kernel_size=(3, 3),
    #                    padding='same', return_sequences=True))
    # seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))  #  360 -> 180
    # seq.add(BatchNormalization())
    #
    # seq.add(ConvLSTM2D(filters=16, kernel_size=(3, 3),
    #                    padding='same', return_sequences=True))
    # seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))  # 180 -> 90
    # # seq.add(BatchNormalization())
    #
    # # seq.add(ConvLSTM2D(filters=8, kernel_size=(3, 3),
    # #                    padding='same', return_sequences=True))
    #
    # seq.add(ConvLSTM2D(filters=1, kernel_size=(3, 3),
    #                    padding='same', return_sequences=True))
    # # seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))
    # # seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))
    # # seq.add(BatchNormalization())
    #
    # # seq.add(Conv3D(filters=1, kernel_size=(3, 3, 3),
    # #                activation='sigmoid',
    # #                padding='same', data_format='channels_last'))

    # optimizer = adadelta(lr=0.1)
    # seq.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])
    seq.compile(loss='binary_crossentropy', optimizer='adadelta', metrics=['accuracy'])
    # seq.compile(loss='mean_absolute_error', optimizer='adadelta', metrics=['accuracy'])

    return seq


#  function to return the lstm network  for multi-channel data
def network_definition_rgb_review(channel_number, full_image_size,  reduced_size, n_frames, n_gt_frames):

    seq = Sequential()
    seq.add(Conv3D(filters=8, kernel_size=(10, 1, 1),
                   input_shape=(n_frames, full_image_size, full_image_size, channel_number),
                   # activation='sigmoid',
                   strides= (int(n_frames/n_gt_frames), 2, 2),
                   padding='same',
                   data_format='channels_last'))

    # seq.add(ConvLSTM2D(filters=8, kernel_size=(3, 3),
    #                    input_shape=(None, full_image_size, full_image_size, channel_number),
    #                    padding='same', return_sequences=True))
    seq.add(ConvLSTM2D(filters=8, kernel_size=(3, 3),
                       padding='same', return_sequences=True))

    seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=16, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=32, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(TimeDistributed(MaxPooling2D((2, 2), padding='same')))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=16, kernel_size=(3, 3),
                       padding='same', return_sequences=True))

    seq.add(Conv3D(filters=1, kernel_size=(2, 3, 3),
                   activation='sigmoid',
                   padding='same', data_format='channels_last'))


    seq.compile(loss='binary_crossentropy', optimizer='adadelta', metrics=['accuracy'])
    # seq.compile(loss='mean_squared_error', optimizer='adadelta', metrics=['accuracy'])

    return seq


#  function to return the lstm network  for single-channel data
def network_definition_single_channel():
    seq = Sequential()
    seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
                       input_shape=(None, 45, 45, 1),
                       padding='same', return_sequences=True))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=40, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(BatchNormalization())

    seq.add(Conv3D(filters=1, kernel_size=(3, 3, 3),
                   activation='sigmoid',
                   padding='same', data_format='channels_last'))
    seq.compile(loss='binary_crossentropy', optimizer='adadelta', metrics=['accuracy'])

    return seq


#  function to perform regression based on the map prediction of the lstm
def network_definition_regressor(channel_number, reduced_size):
    seq = Sequential()
    seq.add(ConvLSTM2D(filters=1, kernel_size=(3, 3),
                       input_shape=(None, reduced_size, reduced_size, channel_number),
                       padding='same', return_sequences=True))
    seq.add(Dropout(0.5))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=1, kernel_size=(3, 3),
                       padding='same', return_sequences=True))
    seq.add(Dropout(0.5))
    seq.add(BatchNormalization())

    seq.add(ConvLSTM2D(filters=1, kernel_size=(3, 3),
                       padding='same', return_sequences=True))

    seq.add(TimeDistributed (Flatten(input_shape=(None,reduced_size, reduced_size, 1) ) ) )

    seq.add(TimeDistributed(Dense(8) ))

    # seq.add(Dense(8))

    seq.compile(loss='mean_absolute_error', optimizer='adadelta' )

    return seq