#!/usr/bin/struture_function1 python
# File: train_tp.py
# Author: Yasmeen George

import tensorflow as tf
#from tensorflow import keras
import argparse
from tensorpack.tfutils.summary import *
from oct_dataflow_tp import *
import tensorflow.contrib.slim as slim
from keras import backend as K
from contextlib import contextmanager
def model_summary():
    model_vars = tf.trainable_variables()
    slim.model_analyzer.analyze_vars(model_vars, print_info=True)
from tensorpack.utils.viz import stack_patches
from vft_utils import perf_measures

def get_features(image,scope):
    with tf.compat.v1.variable_scope(scope):
        l = tf.layers.conv3d(image, 32, 3, padding='SAME',name='conv0')# input_shape=input_shape)
        l=tf.nn.relu(l,name='relu0')
        l = tf.layers.conv3d( l, 16, 3, padding='SAME',name = 'conv1')
        l = tf.nn.relu(l,name='relu1')
        i = 2
        name =""
        for nbchannel in nfilters:
            l =  tf.layers.conv3d(l, nbchannel, 3, padding='SAME',name='conv'+str(i))
            l = tf.layers.batch_normalization(l,axis=-1, momentum=0.8) # input_shape=(input_shape[0], input_shape[1], input_shape[2], nbchannel)
            l = tf.nn.relu(l,name='relu'+str(i))
            name = l.name
            l = tf.layers.max_pooling3d(l,2,2,name = 'maxpool3d'+str(i))
            i +=1

    return l,name


def get_keras_model(l):
    l = tf.layers.conv3d(l, 32, 3, padding='valid',name='conv0')  # input_shape=input_shape)
    l = tf.nn.relu(l,name='relu0')
    i=1
    name = ""
    for nbchannel in nfilters_merged:
        l =  tf.layers.conv3d(l, nbchannel, 3,padding='valid',name='conv'+str(i))
        l = tf.layers.batch_normalization(l,axis=-1, momentum=0.8)# input_shape=(input_shape[0],input_shape[1],input_shape[2],nbchannel),
        l = tf.nn.relu(l,name='relu'+str(i))
        name = l.name
        i+=1

    if CNN_OUT_GarWayHeathmap: # PREDICT GARWAY HEATHMAP AVERAGE REGIONS
        l = tf.reduce_mean(l, axis=[1,2,3])
        l=tf.layers.dense(l, 64, tf.nn.relu)
        l = tf.layers.dropout(l,rate=0.5)
        l = tf.layers.dense(l, out_num)

    else: # predict VFT THRESHOLD VALUES
        l = tf.layers.conv3d(l,1,2,padding='valid')
        l = tf.reduce_mean(l,axis=(3,4))
        #l = tf.layers.conv2d(l,10,2,padding='valid')
        #l = tf.nn.softmax(l,name='pred')
        #l = tf.math.sigmoid(l,name='pred')



    return l,tf.math.sigmoid(l,name='pred'),name


@contextmanager
def guided_relu():
        """
        Returns:
            A context where the gradient of :meth:`tf.nn.relu` is replaced by
            guided back-propagation, as described in the paper:
            `Striving for Simplicity: The All Convolutional Net
            <https://arxiv.org/abs/1412.6806>`_
        """
        from tensorflow.python.ops import gen_nn_ops  # noqa

        @tf.RegisterGradient("GuidedReLU")
        def GuidedReluGrad(op, grad):
            return tf.where(0. < grad,
                            gen_nn_ops.relu_grad(grad, op.outputs[0]),
                            tf.zeros(grad.get_shape()))

        g = tf.get_default_graph()
        with g.gradient_override_map({'Relu': 'GuidedReLU'}):
            yield

def saliency_map(output, input, name="saliency_map"):
        """
        Produce a saliency map as described in the paper:
        `Deep Inside Convolutional Networks: Visualising Image Classification Models and Saliency Maps
        <https://arxiv.org/abs/1312.6034>`_.
        The saliency map is the gradient of the max element in output w.r.t input.

        Returns:
            tf.Tensor: the saliency map. Has the same shape as input.
        """
        max_outp = tf.reduce_max(output, 1)
        saliency_op = tf.gradients(max_outp, input)[:][0]
        return tf.identity(saliency_op, name=name)


class Model(ModelDesc):
    def inputs(self):
        return [tf.TensorSpec((None,)+(SHAPE, SHAPE, dpth),tf.uint8, 'input1'),
            tf.TensorSpec((None,out_num), tf.float32, 'label'),
            tf.TensorSpec((None,) + vft_shape, tf.float32, 'vft_threshold'),
            tf.TensorSpec((None,), tf.string, 'uid')]
    
    def build_graph(self, image ,label,vft_threshold,uid):

        image = tf.expand_dims(tf.cast(image,tf.float32),axis=-1) / 128.0 - 1

        f1,n1=get_features(image, 'pathway1')
        pred,sig_pred,n = get_keras_model(f1)

        model_summary()
        print(f1)
        print(pred)


        '''
        
        with guided_relu():
            saliency_map(pred, tf.get_default_graph().get_tensor_by_name(n1), name="saliency_p1")
            saliency_map(pred, tf.get_default_graph().get_tensor_by_name(n), name="saliency_p5")
        '''
        def dice_coef_loss(y_true, y_pred):
            def dice_coef(y_true, y_pred, smooth=1):
                """
                Dice = (2*|X & Y|)/ (|X|+ |Y|)
                     =  2*sum(|A*B|)/(sum(A^2)+sum(B^2))
                ref: https://arxiv.org/pdf/1606.04797v1.pdf
                """
                intersection = K.sum(K.abs(y_true * y_pred), axis=-1)
                return (2. * intersection + smooth) / (
                            K.sum(K.square(y_true), -1) + K.sum(K.square(y_pred), -1) + smooth)

            return 1 - dice_coef(y_true, y_pred)

        if CNN_OUT_GarWayHeathmap:
            y_true, y_pred = label, sig_pred
        else:
            y_true, y_pred = vft_threshold, sig_pred  # vft_threshold[..., 1:], pred[..., 1:]
            print(y_true, y_pred)

            # dice_loss = dice_coef_loss(y_true, y_pred)
            # ce_loss = tf.keras.losses.categorical_crossentropy(y_true, y_pred)
            # dice_cost = tf.reduce_mean(dice_loss, name='dice_loss')
            # ce_cost = tf.reduce_mean(ce_loss, name='cross_entropy_loss')

        sce_loss = tf.nn.sigmoid_cross_entropy_with_logits(labels=y_true, logits=pred)
        mse_loss = tf.keras.losses.mean_squared_error(y_true, y_pred)
        mae_loss = tf.keras.losses.MAE(y_true, y_pred)


        mse_cost = tf.reduce_mean(mse_loss, name='mean_squared_error')
        mae_cost = tf.reduce_mean(mae_loss, name='mean_absolute_error')
        sce_cost = tf.reduce_mean(sce_loss, name='sigmoid_cross_entropy')

        print(sce_loss, mse_loss, mae_loss)
        print("READUCED_MEAN")
        print(sce_cost, mse_cost, mae_cost)

        # weight decay on all W

        wd_cost = tf.multiply(1e-4, regularize_cost('.*/W', tf.nn.l2_loss), name='wd_cost')
        add_moving_summary(sce_cost, mse_cost, mae_cost, wd_cost)

        add_param_summary(('.*/W', ['histogram']))  # monitor W
        self.cost = tf.add_n([sce_cost, wd_cost], name='cost')
        print(self.cost)
        return self.cost
    def optimizer(self):
        lr = tf.compat.v1.get_variable('learning_rate', initializer=0.01, trainable=False)
        tf.summary.scalar('learning_rate', lr)
        return tf.compat.v1.train.MomentumOptimizer(lr, 0.9, use_nesterov=True)

def test(ds, model_path='',oct_type = 'onh',csv_save_path = 'test_results.csv',vft_type= 'THRESHOLD'):

    in_names = ['input1', 'label', 'vft_threshold', 'uid']

    pred = PredictConfig(
        session_init=SmartInit(model_path),
        model=Model(),
        input_names=in_names,
        output_names=['uid', 'vft_threshold', 'pred', 'logistic_loss', 'Mean_1', 'Mean_2'])

    df_result = perf_measures(ds, pred=pred, oct_type=oct_type,vft_type= vft_type)
    df_result.to_csv(csv_save_path)
    return df_result
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--out_dir', default=os.getcwd() , help='output dir name')# metavar='out_dir'
    parser.add_argument('--data_dir', default=None, help='data dir name') # ,metavar='data_dir'
    parser.add_argument('--gpu', help='comma separated list of GPU(s) to use.')  # nargs='*' in multi mode
    parser.add_argument('--load', help='load model')
    parser.add_argument('--drop_1', default=70, help='Epoch to drop learning rate to 0.01.')#150  # nargs='*' in multi mode
    parser.add_argument('--drop_2', default=120, help='Epoch to drop learning rate to 0.001')#225
    parser.add_argument('--depth', default=20, help='The depth of densenet')# 40
    parser.add_argument('--max_epoch', default=150, help='max epoch') #300
    parser.add_argument('--task', help='task to perform: "train" or "test" or "all" ',
                        choices=['all', 'test', 'train'], default='train')
    parser.add_argument('--oct_type', help='OCT type to use "onh" or "mac" ',
                        choices=['onh', 'mac'], default='onh')
    parser.add_argument('--fold', help='fold number ',
                        choices=['1', '2', '3', '4', '5', '6', '7', '8', '9', '10'], default='1')
    parser.add_argument('--pred', help='Prediction map',
                        choices=['THRESHOLD', 'PATTERN'], default='THRESHOLD')
    parser.add_argument('--load_model', help='load model directory '
                        , default=None)
    parser.add_argument('--model_name', help='model name e.g. model-150000 '
                        , default=None)
    args = parser.parse_args()

    # prepare dataset
    if args.data_dir is None:
        #octdatadir,vftdatadir='../sample_data/oct/', '/Users/gyasmeen/Desktop/Results/nyu_vft_xml/'
        base_dir = '/dccstor/aurmmaret1/Datasets/NYU/'
        octdatadir,vftdatadir = base_dir+'MAC_ONH_1pairPerVisit/MAC_ONH_1pairPerVisit/',base_dir +'nyu_vft_xml/'
    else:
        octdatadir,vftdatadir = args.data_dir +'/', args.data_dir+'/'


    # oct-vft_si_exp10
    # log_dir = '/mnt/results/structure-function-results/training-ybXjrevMg/train_log/'

    # oct-vft_si_exp11_linear
    # log_dir = '/mnt/results/structure-function-results/training-1HwkNRdMg/train_log/'

    # oct-vft_si_exp13_global
    #log_dir = '/mnt/results/structure-function-results/training-h2xVagdMg/train_log/'

    if args.load_model is None:
        log_dir = args.out_dir + "/train_log"
    else:
        log_dir = args.load_model
        #log_dir = '/mnt/results/structure-function-results/training-mocNhn5Gg/train_log/'

    # onh: log_dir = '/mnt/results/structure-function-results/training-SycBtn5Gg/train_log/'

    if args.task != 'train':
        # oct-onh-vft_si_exp11_linear (set: 3000)
        #model_path = '/mnt/results/structure-function-results/training-1HwkNRdMg/train_log/'
        #model_name = 'model-116250'

        # oct-mac-vft_si_exp15_vft-linear (set: 3000)
        if args.load_model is None:
            print('You must enter model path directory')
            exit()
        if args.model_name is None:
            print('You must enter model name')
            exit()
        model_path = args.load_model
        model_name = args.model_name


        dataset_test, te_batch_num = get_data(octdatadir, vftdatadir, SHAPE=SHAPE, BATCH=BATCH, task= args.task,Multi_Input=False,OCT_TYPE=args.oct_type,vft_type = args.pred)
        df_result = test(dataset_test,model_path=model_path+model_name,oct_type=args.oct_type,csv_save_path=model_path+'perf_measures_oct-'+args.oct_type+'-f'+str(args.fold)+'_input-single.csv', vft_type = args.pred)
        print('Test is finished for {} samples', len(df_result))
    elif args.task =='train':


        if args.out_dir is None:
            logger.auto_set_dir()
        else:
            logger_dir = os.path.join(log_dir)
            logger.set_logger_dir(logger_dir,action='k')

        dataset_train,batch_num = get_data(octdatadir,vftdatadir, SHAPE=SHAPE,BATCH=BATCH ,task=args.task,Multi_Input=False,OCT_TYPE=args.oct_type,fold = args.fold,vft_type = args.pred)

        steps_per_epoch = batch_num
        dataset_val,v_batch_num = get_data(octdatadir,vftdatadir, SHAPE=SHAPE,BATCH=BATCH ,task='val',Multi_Input=False,OCT_TYPE=args.oct_type,fold = args.fold,vft_type = args.pred)
        config = tf.ConfigProto()
        config.gpu_options.per_process_gpu_memory_fraction = 0.9
        session = tf.Session(config=config)

        '''
        
        extra_display = ["cost"]
        extra_callbacks = [
            ProgressBar(extra_display),
        ]
        monitors = [
            ScalarPrinter(enable_step=True),
        ]
        '''
        #cfg = TrainConfig(
        cfg = AutoResumeTrainConfig(
            model=Model(),
            dataflow=dataset_train,
            callbacks=[
                PeriodicTrigger(ModelSaver(max_to_keep=10,keep_checkpoint_every_n_hours=1), every_k_epochs=5),
                InferenceRunner(
                    dataset_val,
                    ScalarStats(['sigmoid_cross_entropy', 'mean_squared_error', 'mean_absolute_error'])),
                    #ScalarStats(['dice_loss','cross_entropy_loss','mean_squared_error','mean_absolute_error'])),
                # record GPU utilization during training
                GPUUtilizationTracker(),
                ScheduledHyperParamSetter('learning_rate',
                                          [(args.drop_1, 0.001), (args.drop_2, 0.0001)]),

            ],
            steps_per_epoch=steps_per_epoch,
            max_epoch=args.max_epoch,sess=session
        )


        print('1#1'*100)
        print(get_num_gpu())
        if get_num_gpu() <= 1:
            # single GPU:
            launch_train_with_config(cfg, SimpleTrainer())
        else:
            # multi GPU:
            launch_train_with_config(cfg, SyncMultiGPUTrainerParameterServer(get_num_gpu()))
            # "Replicated" multi-gpu trainer is not supported for Keras model
            # since Keras does not respect variable scopes.

tensorflow.layers.dropout tensorflow.cast tensorflow.nn.sigmoid_cross_entropy_with_logits tensorflow.keras.losses.MAE tensorflow.get_default_graph tensorflow.add_n tensorflow.summary.scalar tensorflow.layers.batch_normalization tensorflow.gradients tensorflow.layers.dense tensorflow.ConfigProto tensorflow.Session tensorflow.trainable_variables tensorflow.compat.v1.variable_scope tensorflow.contrib.slim.model_analyzer.analyze_vars tensorflow.compat.v1.get_variable tensorflow.identity tensorflow.math.sigmoid tensorflow.RegisterGradient tensorflow.layers.conv3d tensorflow.nn.relu tensorflow.reduce_max tensorflow.reduce_mean tensorflow.compat.v1.train.MomentumOptimizer tensorflow.keras.losses.mean_squared_error tensorflow.python.ops.gen_nn_ops.relu_grad tensorflow.TensorSpec
[(14, 'tensorflow.trainable_variables', 'tf.trainable_variables', ([], {}), True, 'import tensorflow as tf\n'), (15, 'tensorflow.contrib.slim.model_analyzer.analyze_vars', 'slim.model_analyzer.analyze_vars', (['model_vars'], {'print_info': '(True)'}), True, 'import tensorflow.contrib.slim as slim\n'), (39, 'tensorflow.layers.conv3d', 'tf.layers.conv3d', (['l', '(32)', '(3)'], {'padding': '"""valid"""', 'name': '"""conv0"""'}), True, 'import tensorflow as tf\n'), (40, 'tensorflow.nn.relu', 'tf.nn.relu', (['l'], {'name': '"""relu0"""'}), True, 'import tensorflow as tf\n'), (79, 'tensorflow.RegisterGradient', 'tf.RegisterGradient', (['"""GuidedReLU"""'], {}), True, 'import tensorflow as tf\n'), (85, 'tensorflow.get_default_graph', 'tf.get_default_graph', ([], {}), True, 'import tensorflow as tf\n'), (99, 'tensorflow.reduce_max', 'tf.reduce_max', (['output', '(1)'], {}), True, 'import tensorflow as tf\n'), (101, 'tensorflow.identity', 'tf.identity', (['saliency_op'], {'name': 'name'}), True, 'import tensorflow as tf\n'), (190, 'vft_utils.perf_measures', 'perf_measures', (['ds'], {'pred': 'pred', 'oct_type': 'oct_type', 'vft_type': 'vft_type'}), False, 'from vft_utils import perf_measures\n'), (194, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), False, 'import argparse\n'), (20, 'tensorflow.compat.v1.variable_scope', 'tf.compat.v1.variable_scope', (['scope'], {}), True, 'import tensorflow as tf\n'), (21, 'tensorflow.layers.conv3d', 'tf.layers.conv3d', (['image', '(32)', '(3)'], {'padding': '"""SAME"""', 'name': '"""conv0"""'}), True, 'import tensorflow as tf\n'), (22, 'tensorflow.nn.relu', 'tf.nn.relu', (['l'], {'name': '"""relu0"""'}), True, 'import tensorflow as tf\n'), (23, 'tensorflow.layers.conv3d', 'tf.layers.conv3d', (['l', '(16)', '(3)'], {'padding': '"""SAME"""', 'name': '"""conv1"""'}), True, 'import tensorflow as tf\n'), (24, 'tensorflow.nn.relu', 'tf.nn.relu', (['l'], {'name': '"""relu1"""'}), True, 'import tensorflow as tf\n'), (45, 'tensorflow.layers.batch_normalization', 'tf.layers.batch_normalization', (['l'], {'axis': '(-1)', 'momentum': '(0.8)'}), True, 'import tensorflow as tf\n'), (51, 'tensorflow.reduce_mean', 'tf.reduce_mean', (['l'], {'axis': '[1, 2, 3]'}), True, 'import tensorflow as tf\n'), (52, 'tensorflow.layers.dense', 'tf.layers.dense', (['l', '(64)', 'tf.nn.relu'], {}), True, 'import tensorflow as tf\n'), (53, 'tensorflow.layers.dropout', 'tf.layers.dropout', (['l'], {'rate': '(0.5)'}), True, 'import tensorflow as tf\n'), (54, 'tensorflow.layers.dense', 'tf.layers.dense', (['l', 'out_num'], {}), True, 'import tensorflow as tf\n'), (57, 'tensorflow.layers.conv3d', 'tf.layers.conv3d', (['l', '(1)', '(2)'], {'padding': '"""valid"""'}), True, 'import tensorflow as tf\n'), (58, 'tensorflow.reduce_mean', 'tf.reduce_mean', (['l'], {'axis': '(3, 4)'}), True, 'import tensorflow as tf\n'), (65, 'tensorflow.math.sigmoid', 'tf.math.sigmoid', (['l'], {'name': '"""pred"""'}), True, 'import tensorflow as tf\n'), (153, 'tensorflow.nn.sigmoid_cross_entropy_with_logits', 'tf.nn.sigmoid_cross_entropy_with_logits', ([], {'labels': 'y_true', 'logits': 'pred'}), True, 'import tensorflow as tf\n'), (154, 'tensorflow.keras.losses.mean_squared_error', 'tf.keras.losses.mean_squared_error', (['y_true', 'y_pred'], {}), True, 'import tensorflow as tf\n'), (155, 'tensorflow.keras.losses.MAE', 'tf.keras.losses.MAE', (['y_true', 'y_pred'], {}), True, 'import tensorflow as tf\n'), (158, 'tensorflow.reduce_mean', 'tf.reduce_mean', (['mse_loss'], {'name': '"""mean_squared_error"""'}), True, 'import tensorflow as tf\n'), (159, 'tensorflow.reduce_mean', 'tf.reduce_mean', (['mae_loss'], {'name': '"""mean_absolute_error"""'}), True, 'import tensorflow as tf\n'), (160, 'tensorflow.reduce_mean', 'tf.reduce_mean', (['sce_loss'], {'name': '"""sigmoid_cross_entropy"""'}), True, 'import tensorflow as tf\n'), (172, 'tensorflow.add_n', 'tf.add_n', (['[sce_cost, wd_cost]'], {'name': '"""cost"""'}), True, 'import tensorflow as tf\n'), (176, 'tensorflow.compat.v1.get_variable', 'tf.compat.v1.get_variable', (['"""learning_rate"""'], {'initializer': '(0.01)', 'trainable': '(False)'}), True, 'import tensorflow as tf\n'), (177, 'tensorflow.summary.scalar', 'tf.summary.scalar', (['"""learning_rate"""', 'lr'], {}), True, 'import tensorflow as tf\n'), (178, 'tensorflow.compat.v1.train.MomentumOptimizer', 'tf.compat.v1.train.MomentumOptimizer', (['lr', '(0.9)'], {'use_nesterov': '(True)'}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.layers.batch_normalization', 'tf.layers.batch_normalization', (['l'], {'axis': '(-1)', 'momentum': '(0.8)'}), True, 'import tensorflow as tf\n'), (82, 'tensorflow.python.ops.gen_nn_ops.relu_grad', 'gen_nn_ops.relu_grad', (['grad', 'op.outputs[0]'], {}), False, 'from tensorflow.python.ops import gen_nn_ops\n'), (100, 'tensorflow.gradients', 'tf.gradients', (['max_outp', 'input'], {}), True, 'import tensorflow as tf\n'), (106, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['((None,) + (SHAPE, SHAPE, dpth))', 'tf.uint8', '"""input1"""'], {}), True, 'import tensorflow as tf\n'), (107, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['(None, out_num)', 'tf.float32', '"""label"""'], {}), True, 'import tensorflow as tf\n'), (108, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['((None,) + vft_shape)', 'tf.float32', '"""vft_threshold"""'], {}), True, 'import tensorflow as tf\n'), (109, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['(None,)', 'tf.string', '"""uid"""'], {}), True, 'import tensorflow as tf\n'), (275, 'tensorflow.ConfigProto', 'tf.ConfigProto', ([], {}), True, 'import tensorflow as tf\n'), (277, 'tensorflow.Session', 'tf.Session', ([], {'config': 'config'}), True, 'import tensorflow as tf\n'), (113, 'tensorflow.cast', 'tf.cast', (['image', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (136, 'keras.backend.abs', 'K.abs', (['(y_true * y_pred)'], {}), True, 'from keras import backend as K\n'), (138, 'keras.backend.square', 'K.square', (['y_true'], {}), True, 'from keras import backend as K\n'), (138, 'keras.backend.square', 'K.square', (['y_pred'], {}), True, 'from keras import backend as K\n')]