###################################################
#
#   Script to:
#   - Load the images and extract the patches
#   - Define the neural network
#   - define the training
#
##################################################


import numpy as np
import configparser


###################################################
import tensorflow as tf
from keras import backend as K
sess = tf.Session()
#sess.run(tf.global_variables_initializer())
#sess.as_default()
K.set_session(sess)
###################################################

from keras.models import Model
from keras.layers import Input, Concatenate, Conv2D, MaxPooling2D, UpSampling2D, Reshape, core, Dropout, Add, Lambda
from keras.optimizers import SGD, Adam


from keras.callbacks import ModelCheckpoint, LearningRateScheduler
from keras.utils.vis_utils import plot_model as plot
#from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator

######
#from keras.layers import Lambda

import sys
sys.path.insert(0, '../lib/')
from help_functions import *

#function to obtain data for training/testing (validation)
from extract_patches import get_data_training
from networks import *
from pixel_dcn import pixel_dcl
from inception_res_v4 import *


def pdcn_unet(patch_height, patch_width, n_ch):
    inputs = Input(shape=(patch_height, patch_width, n_ch))
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same',data_format='channels_last')(inputs)
    conv1 = Dropout(0.2)(conv1)
    conv1 = Conv2D(32, (3, 3), activation='relu', padding='same',data_format='channels_last')(conv1)
    pool1 = MaxPooling2D((2, 2), data_format='channels_last')(conv1)
    
    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same',data_format='channels_last')(pool1)
    conv2 = Dropout(0.2)(conv2)
    conv2 = Conv2D(64, (3, 3), activation='relu', padding='same',data_format='channels_last')(conv2)
    pool2 = MaxPooling2D((2, 2), data_format='channels_last')(conv2)
    
    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same',data_format='channels_last')(pool2)
    conv3 = Dropout(0.2)(conv3)
    conv3 = Conv2D(128, (3, 3), activation='relu', padding='same',data_format='channels_last')(conv3)

    up1 = UpSampling2D(size=(2, 2), data_format='channels_last')(conv3)
    up1 = Concatenate(axis=3)([conv2, up1])
    
    conv4 = Conv2D(64, (3, 3), activation='relu', padding='same',data_format='channels_last')(up1)
    conv4 = Dropout(0.2)(conv4)
    conv4 = Conv2D(64, (3, 3), activation='relu', padding='same',data_format='channels_last')(conv4)
    
    up2 = Lambda(lambda x: pixel_dcl(inputs=x, out_num=64, kernel_size=(2, 2), scope='conv_pdcl', d_format='NHWC'))(conv4)
    #up2 = UpSampling2D(size=(2, 2), data_format='channels_last')(conv4)
    up2 = Concatenate(axis=3)([conv1, up2])
    
    conv5 = Conv2D(32, (3, 3), activation='relu', padding='same',data_format='channels_last')(up2)
    conv5 = Dropout(0.2)(conv5)
    conv5 = Conv2D(32, (3, 3), activation='relu', padding='same',data_format='channels_last')(conv5)
    
    conv6 = Conv2D(2, (1, 1), activation='relu',padding='same',data_format='channels_last')(conv5)
    conv6 = core.Reshape((patch_height*patch_width, 2))(conv6)
    
    conv7 = core.Activation('softmax')(conv6)

    model = Model(inputs=inputs, outputs=conv7)
    print(model.summary())

    opt = Adam(lr=0.01, beta_1=0.9, beta_2=0.999, epsilon=0.1, decay=1e-5)  # amsgrad=True doesn't exist in Keras 2.1.1
    sess.run(tf.global_variables_initializer())
    model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])

    return model



#========= Load settings from Config file
config = configparser.RawConfigParser()
config.read('../configuration.txt')
#patch to the datasets
path_data = config.get('data paths', 'path_local')
#Experiment name
name_experiment = config.get('experiment name', 'name')
#training settings
N_epochs = int(config.get('training settings', 'N_epochs'))
batch_size = int(config.get('training settings', 'batch_size'))
zca = bool(config.get('training settings', 'zca'))



#============ Load the data and divided in patches
patches_imgs_train, patches_masks_train = get_data_training(
    DRIVE_train_imgs_original = path_data + config.get('data paths', 'train_imgs_original'),
    DRIVE_train_groudTruth = path_data + config.get('data paths', 'train_groundTruth'),  #masks
    patch_height = int(config.get('data attributes', 'patch_height')),
    patch_width = int(config.get('data attributes', 'patch_width')),
    N_subimgs = int(config.get('training settings', 'N_subimgs')),
    inside_FOV = config.getboolean('training settings', 'inside_FOV') #select the patches only inside the FOV  (default == True)
)


#========= Save a sample of what you're feeding to the neural network ==========
N_sample = min(patches_imgs_train.shape[0],40)
visualize(group_images(patches_imgs_train[0:N_sample,:,:,:],5),'../'+name_experiment+'/'+"sample_input_imgs")#.show()
visualize(group_images(patches_masks_train[0:N_sample,:,:,:],5),'../'+name_experiment+'/'+"sample_input_masks")#.show()


#=========== Construct and save the model arcitecture =====
n_ch = patches_imgs_train.shape[1]
patch_height = patches_imgs_train.shape[2]
patch_width = patches_imgs_train.shape[3]
################################################################MODIFY####################################################################
#model = get_unet(n_ch, patch_height, patch_width)  #the U-net model
model = pdcn_unet(patch_height, patch_width, n_ch)  #the U-net model
################################################################MODIFY####################################################################
print("Check: final output of the network:")
print(model.output_shape)
#plot(model, to_file='../'+name_experiment+'/'+name_experiment + '_model.png')   #check how the model looks like
json_string = model.to_json()
open('../'+name_experiment+'/'+name_experiment +'_architecture.json', 'w').write(json_string)



#============  Training ==================================
checkpointer = ModelCheckpoint(filepath='../'+name_experiment+'/'+name_experiment +'_best_weights.h5', verbose=1, monitor='val_loss', mode='auto', save_best_only=True, save_weights_only=False) #save at each epoch if the validation decreased

patches_masks_train = masks_Unet(patches_masks_train)  #reduce memory consumption

# Apply ZCA to batches
#datagen = ImageDataGenerator(zca_whitening=zca, data_format='channels_last', validation_split=0.1)
#datagen.fit(patches_imgs_train)
#train_generator = datagen.flow(patches_imgs_train, patches_masks_train, batch_size=batch_size, subset='training')
#validation_generator = datagen.flow(patches_imgs_train, patches_masks_train, batch_size=batch_size, subset='validation')

#model.fit_generator(train_generator, epochs=N_epochs, verbose=2, shuffle=True, validation_data=validation_generator, callbacks=[checkpointer])

# Roll the axis from channels_first to channels_last
patches_imgs_train = np.rollaxis(patches_imgs_train, 1, 4)
print("train PATCHES images/masks final shape:")
print(patches_imgs_train.shape)

model.fit(patches_imgs_train, patches_masks_train, nb_epoch=N_epochs, batch_size=batch_size, verbose=2, shuffle=True, validation_split=0.1, callbacks=[checkpointer])

#========== Save and test the last model ===================
pdcn = bool(config.get('training settings', 'pdcn'))
if pdcn == False:
    model.save_weights('../'+name_experiment+'/'+name_experiment +'_last_weights.h5', overwrite=True)
else:
    saver = tf.train.Saver()
    sess = K.get_session()
    saver.save(sess, '../' + name_experiment + '/keras_model')