import tensorflow as tf

import DataGenerator

import time

learning_rate = 1e-5

batch_size = 2

prefetch = 4

no_of_epochs = 10000

smoothing = 0.00001

# placeholder for training mode

is_training = tf.placeholder(tf.bool)

# input data generator

trainTransforms = [

    DataGenerator.RandomFlip(),

    DataGenerator.HistogramMatching(data_dir='train-data', train_size=40, prob=0.5),

    DataGenerator.RandomSmoothing(prob=0.3),

    DataGenerator.RandomNoise(prob=0.5),

    DataGenerator.Normalization()

    ]

valTransforms = [

    DataGenerator.Normalization()

    ]

TrainDataset = DataGenerator.DataGenerator(

    data_dir='train-data',

    transforms=trainTransforms,

    train=True

    )

ValDataset = DataGenerator.DataGenerator(

    data_dir='val-data',

    transforms=valTransforms,

    train=False

    )

trainDataset = TrainDataset.get_dataset()

trainDataset = trainDataset.shuffle(buffer_size=5)

trainDataset = trainDataset.batch(batch_size)

trainDataset = trainDataset.prefetch(prefetch)

valDataset = ValDataset.get_dataset()

valDataset = valDataset.shuffle(buffer_size=5)

valDataset = valDataset.batch(batch_size)

valDataset = valDataset.prefetch(prefetch)

iterator = tf.data.Iterator.from_structure(trainDataset.output_types, trainDataset.output_shapes)

training_init_op = iterator.make_initializer(trainDataset)

validation_init_op = iterator.make_initializer(valDataset)

next_item = iterator.get_next()

# convolution layer

def conv3d(x, no_of_input_channels, no_of_filters, filter_size, strides, padding, name):

    with tf.variable_scope(name) as scope:

        initializer = tf.variance_scaling_initializer()

        filter_size.extend([no_of_input_channels, no_of_filters])

        weights = tf.Variable(initializer(filter_size), name='weights')

        biases = tf.Variable(initializer([no_of_filters]), name='biases')

        conv = tf.nn.conv3d(x, weights, strides=strides, padding=padding, name=name)

        conv += biases

        return conv

# transposed convolution layer

def upsamp(x, no_of_kernels, name):

    with tf.variable_scope(name) as scope:

        upsamp = tf.layers.conv3d_transpose(x, no_of_kernels, [2,2,2], 2, padding='VALID', use_bias=True, reuse=tf.AUTO_REUSE)

        return upsamp

# PReLu layer

def prelu(x, scope=None):

    with tf.variable_scope(name_or_scope=scope, default_name="prelu", reuse=tf.AUTO_REUSE):

        alpha = tf.get_variable("prelu", shape=x.get_shape()[-1], dtype=x.dtype, initializer=tf.constant_initializer(0.1))

        prelu_out = tf.maximum(0.0, x) + alpha * tf.minimum(0.0, x)

        return prelu_out

# model graph

def graph_encoder(x):

    fine_grained_features = {}

    conv1 = conv3d(x,1,16,[3,3,3],[1,1,1,1,1],'SAME','Conv1_1')

    conv1 = conv3d(conv1,16,16,[3,3,3],[1,1,1,1,1],'SAME','Conv1_2')

    conv1 = tf.layers.batch_normalization(conv1, training=is_training)

    conv1 = prelu(conv1,'prelu1')

    res1 = tf.add(x,conv1)

    fine_grained_features['res1'] = res1

    down1 = conv3d(res1,16,32,[2,2,2],[1,2,2,2,1],'VALID','DownSampling1')

    down1 = prelu(down1,'down_prelu1')

    conv2 = conv3d(down1,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv2_1')

    conv2 = conv3d(conv2,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv2_2')

    conv2= tf.layers.batch_normalization(conv2, training=is_training)

    conv2 = prelu(conv2,'prelu2')

    conv3 = conv3d(conv2,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv3_1')

    conv3 = conv3d(conv3,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv3_2')

    conv3 = tf.layers.batch_normalization(conv3, training=is_training)

    conv3 = prelu(conv3,'prelu3')

    res2 = tf.add(down1,conv3)

    fine_grained_features['res2'] = res2

    down2 = conv3d(res2,32,64,[2,2,2],[1,2,2,2,1],'VALID','DownSampling2')

    down2 = prelu(down2,'down_prelu2')

    conv4 = conv3d(down2,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv4_1')

    conv4 = conv3d(conv4,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv4_2')

    conv4 = tf.layers.batch_normalization(conv4, training=is_training)

    conv4 = prelu(conv4,'prelu4')

    conv5 = conv3d(conv4,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv5_1')

    conv5 = conv3d(conv5,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv5_2')

    conv5 = tf.layers.batch_normalization(conv5, training=is_training)

    conv5 = prelu(conv5,'prelu5')

    conv6 = conv3d(conv5,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv6_1')

    conv6 = conv3d(conv6,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv6_2')

    conv6 = tf.layers.batch_normalization(conv6, training=is_training)

    conv6 = prelu(conv6,'prelu6')

    res3 = tf.add(down2,conv6)

    fine_grained_features['res3'] = res3

    down3 = conv3d(res3,64,128,[2,2,2],[1,2,2,2,1],'VALID','DownSampling3')

    down3 = prelu(down3,'down_prelu3')

    conv7 = conv3d(down3,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv7_1')

    conv7 = conv3d(conv7,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv7_2')

    conv7 = tf.layers.batch_normalization(conv7, training=is_training)

    conv7 = prelu(conv7,'prelu7')

    conv8 = conv3d(conv7,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv8_1')

    conv8 = conv3d(conv8,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv8_2')

    conv8 = tf.layers.batch_normalization(conv8, training=is_training)

    conv8 = prelu(conv8,'prelu8')

    conv9 = conv3d(conv8,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv9_1')

    conv9 = conv3d(conv9,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv9_2')

    conv9 = tf.layers.batch_normalization(conv9, training=is_training)

    conv9 = prelu(conv9,'prelu9')

    res4 = tf.add(down3,conv9)

    fine_grained_features['res4'] = res4

    down4 = conv3d(res4,128,256,[2,2,2],[1,2,2,2,1],'VALID','DownSampling4')

    down4 = prelu(down4,'down_prelu4')

    conv10 = conv3d(down4,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv10_1')

    conv10 = conv3d(conv10,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv10_2')

    conv10 = tf.layers.batch_normalization(conv10, training=is_training)

    conv10 = prelu(conv10,'prelu10')

    conv11 = conv3d(conv10,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv11_1')

    conv11 = conv3d(conv11,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv11_2')

    conv11 = tf.layers.batch_normalization(conv11, training=is_training)

    conv11 = prelu(conv11,'prelu11')

    conv12 = conv3d(conv11,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv12_1')

    conv12 = conv3d(conv12,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv12_2')

    conv12 = tf.layers.batch_normalization(conv12, training=is_training)

    conv12 = prelu(conv12,'prelu12')

    res5 = tf.add(down4,conv12)

    fine_grained_features['res5'] = res5

    return fine_grained_features

def graph_decoder(features):

    inp = features['res5']

    upsamp1 = upsamp(inp,128,'Upsampling1')

    upsamp1 = prelu(upsamp1,'prelu_upsamp1')

    concat1 = tf.concat([upsamp1,features['res4']],axis=4)

    conv13 = conv3d(concat1,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv13_1')

    conv13 = conv3d(conv13,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv13_2')

    conv13 = tf.layers.batch_normalization(conv13, training=is_training)

    conv13 = prelu(conv13,'prelu13')

    conv14 = conv3d(conv13,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv14_1')

    conv14 = conv3d(conv14,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv14_2')

    conv14 = tf.layers.batch_normalization(conv14, training=is_training)

    conv14 = prelu(conv14,'prelu14')

    conv15 = conv3d(conv14,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv15_1')

    conv15 = conv3d(conv15,256,256,[3,3,3],[1,1,1,1,1],'SAME','Conv15_2')

    conv15 = tf.layers.batch_normalization(conv15, training=is_training)

    conv15 = prelu(conv15,'prelu15')

    res6 = tf.add(concat1,conv15)

    upsamp2 = upsamp(res6,64,'Upsampling2')

    upsamp2 = prelu(upsamp2,'prelu_upsamp2')

    concat2 = tf.concat([upsamp2,features['res3']],axis=4)

    conv16 = conv3d(concat2,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv16_1')

    conv16 = conv3d(conv16,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv16_2')

    conv16 = tf.layers.batch_normalization(conv16, training=is_training)

    conv16 = prelu(conv16,'prelu16')

    conv17 = conv3d(conv16,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv17_1')

    conv17 = conv3d(conv17,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv17_2')

    conv17 = tf.layers.batch_normalization(conv17, training=is_training)

    conv17 = prelu(conv17,'prelu17')

    conv18 = conv3d(conv17,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv18_1')

    conv18 = conv3d(conv18,128,128,[3,3,3],[1,1,1,1,1],'SAME','Conv18_2')

    conv18 = tf.layers.batch_normalization(conv18, training=is_training)

    conv18 = prelu(conv18,'prelu18')

    res7 = tf.add(concat2,conv18)

    upsamp3 = upsamp(res7,32,'Upsampling3')

    upsamp3 = prelu(upsamp3,'prelu_upsamp3')

    concat3 = tf.concat([upsamp3,features['res2']],axis=4)

    conv19 = conv3d(concat3,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv19_1')

    conv19 = conv3d(conv19,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv19_2')

    conv19 = tf.layers.batch_normalization(conv19, training=is_training)

    conv19 = prelu(conv19,'prelu19')

    conv20 = conv3d(conv19,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv20_1')

    conv20 = conv3d(conv20,64,64,[3,3,3],[1,1,1,1,1],'SAME','Conv20_2')

    conv20 = tf.layers.batch_normalization(conv20, training=is_training)

    conv20 = prelu(conv20,'prelu20')

    res8 = tf.add(concat3,conv20)

    upsamp4 = upsamp(res8,16,'Upsampling4')

    upsamp4 = prelu(upsamp4,'prelu_upsamp4')

    concat4 = tf.concat([upsamp4,features['res1']],axis=4)

    conv21 = conv3d(concat4,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv21_1')

    conv21 = conv3d(conv21,32,32,[3,3,3],[1,1,1,1,1],'SAME','Conv21_2')

    conv21 = tf.layers.batch_normalization(conv21, training=is_training)

    conv21 = prelu(conv21,'prelu21')

    res9 = tf.add(concat4,conv21)

    conv22 = conv3d(res9,32,1,[1,1,1],[1,1,1,1,1],'SAME','Conv22')

    conv22 = tf.nn.sigmoid(conv22,'sigmoid')

    return conv22

# loss and optimizer

def model_fn():

    features, labels = next_item

    features = tf.reshape(features, [-1, 128, 128, 64, 1])

    labels = tf.cast(tf.reshape(labels, [-1, 128, 128, 64, 1]), dtype=tf.float32)

    # writing summaries to tensorboard

    tf.summary.image('features', features[:, :, :, 32:33, 0], max_outputs=2,collections=['val'])

    tf.summary.image('labels', labels[:, :, :, 32:33, 0], max_outputs=2,collections=['val'])

    labels = tf.reshape(labels, [-1,128*128*64])

    encoded = graph_encoder(features)

    decoded = graph_decoder(encoded)

    decoded = tf.reshape(decoded, [-1, 128, 128, 64])

    tf.summary.image('segmentation', decoded[:, :, :, 32:33], max_outputs=2, collections=['val'])

    output =  tf.reshape(decoded, [-1,128*128*64])

    # dice loss

    cost = tf.reduce_mean(tf.divide(smoothing + tf.multiply(2.0, tf.reduce_sum(output * labels, axis=-1)),

                 tf.add(tf.reduce_sum(output, axis=-1), tf.reduce_sum(labels, axis=-1))))

    tf.summary.scalar('training_loss', cost)

    tf.summary.scalar('val_loss', cost,collections=['val'])

    # for batchnorm

    with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):

        opt = tf.train.AdamOptimizer(learning_rate)

        grads = tf.gradients(1-cost, tf.trainable_variables())

        grads = list(zip(grads, tf.trainable_variables()))

        training_operation = opt.apply_gradients(grads_and_vars=grads)

    for grad, var in grads:

        tf.summary.histogram(var.name.replace(':',"_") + '/gradient', grad)

        tf.summary.histogram(var.name.replace(':',"_") , var)

    return cost, training_operation

# running the session 

def train():

    with tf.Session() as sess:

        cost, opt = model_fn()

        sess.run(tf.global_variables_initializer())

        # merging tensorflow summaries

        merged = tf.summary.merge_all()

        merged_val = tf.summary.merge_all(key = 'val')

        train_writer = tf.summary.FileWriter('event/train',sess.graph)

        val_writer = tf.summary.FileWriter('event/val')

        saver = tf.train.Saver()

        for epoch in range(1, no_of_epochs+1):

            start_time = time.time()

            train_loss = []

            examples = 0

            # initializing iterator with training dataset

            sess.run([training_init_op])

            while(True):

                try:

                    # training procedure

                    examples += 1

                    loss, _, summary = sess.run([cost, opt, merged], feed_dict={is_training: True})

                    train_writer.add_summary(summary,epoch)

                    train_loss.append(loss)

                    print('Epoch: {} - ex: {} - loss: {:.6f}'.format(epoch, examples*batch_size, sum(train_loss)/len(train_loss)), end="\r")

                except tf.errors.OutOfRangeError:

                    val_loss = []

                    val_example = 0

                    # initializing iterator with validation dataset

                    sess.run([validation_init_op])

                    while(True):

                        try:

                            val_example += 1

                            loss, summary_l = sess.run([cost, merged_val], feed_dict={is_training: False})

                            val_writer.add_summary(summary_l,epoch)

                            val_loss.append(loss)

                            print('Epoch: {} - ex: {} - val_loss: {:.6f}'.format(epoch, val_example*batch_size, sum(val_loss)/len(val_loss)), end="\r")

                        except tf.errors.OutOfRangeError:

                            break

                    break

            print('Epoch: {}/{} - loss: {:.6f} - val_loss: {:.6f} - time: {:.4f}'.format(epoch, no_of_epochs, 

                sum(train_loss)/len(train_loss), sum(val_loss)/len(val_loss), time.time()-start_time))

            # saving weights

            if epoch%20==0:

                saver.save(sess, '/temp/weights_epoch_{0}.ckpt'.format(epoch))

if __name__ == '__main__':

    train()