from __future__ import print_function

import numpy as np

import sys

import tensorflow as tf

from tensorflow.image import resize_images

# try:

#     import special_grads

# except KeyError as e:

#     print('WARN: Cannot define MaxPoolGrad, likely already defined for this version of tensorflow: %s' % e, file=sys.stderr)

from tensorflow.python.platform import flags

from layer import conv_block, deconv_block, fc, max_pool, concat2d

from utils import xent, kd, _get_segmentation_cost, _get_compactness_cost

class SAML:

    def __init__(self, args):

        """ Call construct_model_*() after initializing MASF"""

        self.args = args

        self.batch_size = args.meta_batch_size

        self.test_batch_size = args.test_batch_size

        self.volume_size = args.volume_size

        self.n_class = args.n_class

        self.compactness_loss_weight = args.compactness_loss_weight

        self.smoothness_loss_weight = args.smoothness_loss_weight

        self.margin = args.margin

        self.forward = self.forward_unet

        self.construct_weights = self.construct_unet_weights

        self.seg_loss = _get_segmentation_cost

        self.get_compactness_cost = _get_compactness_cost

    def construct_model_train(self, prefix='metatrain_'):

        # a: meta-train for inner update, b: meta-test for meta loss

        self.inputa = tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], self.volume_size[2]])

        self.labela = tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], self.n_class])

        self.inputa1= tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], self.volume_size[2]])

        self.labela1= tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], self.n_class])

        self.inputb = tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], self.volume_size[2]])

        self.labelb = tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], self.n_class])

        self.input_group = tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], self.volume_size[2]])

        self.label_group = tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], self.n_class])

        self.contour_group = tf.placeholder("float", shape=[self.batch_size, self.volume_size[0], self.volume_size[1], 1])

        self.metric_label_group = tf.placeholder(tf.int32, shape=[self.batch_size, 1])

        self.training_mode = tf.placeholder_with_default(True, shape = None, name = "training_mode_for_bn_moving")

        self.clip_value = self.args.gradients_clip_value

        self.KEEP_PROB = tf.placeholder(tf.float32)

        with tf.variable_scope('model', reuse=None) as training_scope:

            if 'weights' in dir(self):

                print('weights already defined')

                training_scope.reuse_variables()

                weights = self.weights

            else:

                # Define the weights

                self.weights = weights = self.construct_weights()

            def task_metalearn(inp, reuse=True):

                # Function to perform meta learning update """

                inputa, inputa1, inputb, labela, labela1, labelb, input_group, contour_group, metric_label_group = inp

                # Obtaining the conventional task loss on meta-train

                task_outputa, _, _ = self.forward(inputa, weights, is_training=self.training_mode)

                task_lossa = self.seg_loss(task_outputa, labela)

                task_outputa1, _, _ = self.forward(inputa1, weights, is_training=self.training_mode)

                task_lossa1 = self.seg_loss(task_outputa1, labela1)

                ## perform inner update with plain gradient descent on meta-train

                grads = tf.gradients((task_lossa + task_lossa1)/2.0, list(weights.values()))

                grads = [tf.stop_gradient(grad) for grad in grads] # first-order gradients approximation

                gradients = dict(zip(weights.keys(), grads))

                # fast_weights = dict(zip(weights.keys(), [weights[key] - self.inner_lr * gradients[key] for key in weights.keys()]))

                fast_weights = dict(zip(weights.keys(), [weights[key] - self.inner_lr * tf.clip_by_norm(gradients[key], clip_norm=self.clip_value) for key in weights.keys()]))

                ## compute compactness loss

                task_outputb, task_predmaskb, _ = self.forward(inputb, fast_weights, is_training=self.training_mode)

                task_lossb = self.seg_loss(task_outputb, labelb)

                compactness_loss_b, length, area, boundary_b = self.get_compactness_cost(task_outputb, labelb)

                compactness_loss_b = self.compactness_loss_weight * compactness_loss_b

                # compute smoothness loss

                _, _, embeddings = self.forward(input_group, fast_weights, is_training=self.training_mode)

                coutour_embeddings = self.extract_coutour_embedding(contour_group, embeddings)

                metric_embeddings = self.forward_metric_net(coutour_embeddings)

                print (metric_label_group.shape)

                print (metric_embeddings.shape)

                smoothness_loss_b = tf.contrib.losses.metric_learning.triplet_semihard_loss(labels=metric_label_group[..., 0], embeddings=metric_embeddings, margin=self.margin)

                smoothness_loss_b = self.smoothness_loss_weight * smoothness_loss_b

                task_output = [task_lossb, compactness_loss_b, smoothness_loss_b, task_predmaskb, boundary_b, length, area, task_lossa, task_lossa1]

                return task_output

            self.global_step = tf.Variable(0, trainable=False)

            # self.inner_lr = tf.train.exponential_decay(learning_rate=self.args.inner_lr, global_step=self.global_step, decay_steps=self.args.decay_step, decay_rate=self.args.decay_rate)

            # self.outer_lr = tf.train.exponential_decay(learning_rate=self.args.outer_lr, global_step=self.global_step, decay_steps=self.args.decay_step, decay_rate=self.args.decay_rate)

            self.inner_lr = tf.Variable(self.args.inner_lr, trainable=False)

            self.outer_lr = tf.Variable(self.args.outer_lr, trainable=False)

            self.metric_lr = tf.Variable(self.args.metric_lr, trainable=False)

            input_tensors = (self.inputa, self.inputa1, self.inputb, self.labela, self.labela1, self.labelb, self.input_group, self.contour_group, self.metric_label_group)

            result = task_metalearn(inp=input_tensors)

            self.seg_loss_b, self.compactness_loss_b, self.smoothness_loss_b, self.task_predmaskb, self.boundary_b, self.length, self.area, self.seg_loss_a, self.seg_loss_a1= result

        ## Performance & Optimization

        if 'train' in prefix:

            self.source_loss = (self.seg_loss_a + self.seg_loss_a1) / 2.0

            self.target_loss = self.seg_loss_b + self.compactness_loss_b + self.smoothness_loss_b

            var_list_segmentor = [v for v in tf.trainable_variables() if 'metric' not in v.name.split('/')]

            var_list_metric = [v for v in tf.trainable_variables() if 'metric' in v.name.split('/')]

            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)

            with tf.control_dependencies(update_ops):

                self.task_train_op = tf.train.AdamOptimizer(learning_rate=self.inner_lr).minimize(self.source_loss, global_step=self.global_step)

            optimizer = tf.train.AdamOptimizer(self.outer_lr)

            gvs = optimizer.compute_gradients(self.target_loss, var_list=var_list_segmentor)

            # observe stability of gradients for meta loss

            # l2_norm = lambda t: tf.sqrt(tf.reduce_sum(tf.pow(t, 2)))

            # for grad, var in gvs:

            #     tf.summary.histogram("gradients_norm/" + var.name, l2_norm(grad))

            #     tf.summary.histogram("feature_extractor_var_norm/" + var.name, l2_norm(var))

            #     tf.summary.histogram('gradients/' + var.name, var)

            #     tf.summary.histogram("feature_extractor_var/" + var.name, var)

            # gvs = [(grad, var) for grad, var in gvs]

            gvs = [(tf.clip_by_norm(grad, clip_norm=self.clip_value), var) for grad, var in gvs]

            self.meta_train_op = optimizer.apply_gradients(gvs)

            # for grad, var in gvs:

            #     tf.summary.histogram("gradients_norm_clipped/" + var.name, l2_norm(grad))

            #     tf.summary.histogram('gradients_clipped/' + var.name, var)

            self.metric_train_op = tf.train.AdamOptimizer(self.metric_lr).minimize(self.smoothness_loss_b, var_list=var_list_metric)

        ## Summaries

        # scalar_summaries = []

        # train_images = []

        # val_images = []

        tf.summary.scalar(prefix+'source_1 loss', self.seg_loss_a)

        tf.summary.scalar(prefix+'source_2 loss', self.seg_loss_a1)

        tf.summary.scalar(prefix+'target_loss', self.seg_loss_b)

        tf.summary.scalar(prefix+'target_coutour_loss', self.compactness_loss_b)

        tf.summary.scalar(prefix+'target_length', self.length)

        tf.summary.scalar(prefix+'target_area', self.area)

        tf.summary.image("meta_test_mask", tf.expand_dims(tf.cast(self.task_predmaskb, tf.float32), 3))

        tf.summary.image("meta_test_gth", tf.expand_dims(tf.cast(self.labelb[:,:,:,1], tf.float32), 3))

        tf.summary.image("meta_test_image", tf.expand_dims(tf.cast(self.inputb[:,:,:,1], tf.float32), 3))

        tf.summary.image("meta_test_boundary", tf.expand_dims(tf.cast(self.boundary_b[:,:,:], tf.float32), 3))

        tf.summary.image("meta_test_ct_bg_sample", tf.expand_dims(tf.cast(self.contour_group[:,:,:, 0], tf.float32), 3))

        tf.summary.image("meta_input_group", tf.expand_dims(tf.cast(self.input_group[:,:,:, 1], tf.float32), 3))

        tf.summary.image("label_group", tf.expand_dims(tf.cast(self.label_group[:,:,:, 1], tf.float32), 3))

    def extract_coutour_embedding(self, coutour, embeddings):

        coutour_embeddings = coutour * embeddings

        average_embeddings = tf.reduce_sum(coutour_embeddings, [1,2])/tf.reduce_sum(coutour, [1,2])

        # print (coutour.shape)

        # print (embeddings.shape)

        # print (coutour_embeddings.shape)

        # print (average_embeddings.shape)

        return average_embeddings

    def construct_model_test(self, prefix='test'):

        self.test_input = tf.placeholder("float", shape=[self.test_batch_size, self.volume_size[0], self.volume_size[1], self.volume_size[2]])

        self.test_label = tf.placeholder("float", shape=[self.test_batch_size, self.volume_size[0], self.volume_size[1], self.n_class])

        with tf.variable_scope('model', reuse=None) as testing_scope:

            if 'weights' in dir(self):

                testing_scope.reuse_variables()

                weights = self.weights

            else:

                raise ValueError('Weights not initilized. Create training model before testing model')

            outputs, mask, _ = self.forward(self.test_input, weights)

            losses = self.seg_loss(outputs, self.test_label)

            # self.pred_prob = tf.nn.softmax(outputs)

            self.outputs = mask

        self.test_loss = losses

        # self.test_acc = accuracies

    def forward_metric_net(self, x):

        with tf.variable_scope('metric', reuse=tf.AUTO_REUSE) as scope:

            w1 = tf.get_variable('w1', shape=[48,24])

            b1 = tf.get_variable('b1', shape=[24])

            out = fc(x, w1, b1, activation='leaky_relu')

            w2 = tf.get_variable('w2', shape=[24,16])

            b2 = tf.get_variable('b2', shape=[16])

            out = fc(out, w2, b2, activation='leaky_relu')

        return out

    def construct_unet_weights(self):

        weights = {}

        conv_initializer = tf.contrib.layers.xavier_initializer_conv2d(dtype=tf.float32)

        with tf.variable_scope('conv1') as scope:

            weights['conv11_weights'] = tf.get_variable('weights', shape=[5, 5, 3, 16], initializer=conv_initializer)

            weights['conv11_biases'] = tf.get_variable('biases', [16])

            weights['conv12_weights'] = tf.get_variable('weights2', shape=[5, 5, 16, 16], initializer=conv_initializer)

            weights['conv12_biases'] = tf.get_variable('biases2', [16])

        with tf.variable_scope('conv2') as scope:

            weights['conv21_weights'] = tf.get_variable('weights', shape=[5, 5, 16, 32], initializer=conv_initializer)

            weights['conv21_biases'] = tf.get_variable('biases', [32])

            weights['conv22_weights'] = tf.get_variable('weights2', shape=[5, 5, 32, 32], initializer=conv_initializer)

            weights['conv22_biases'] = tf.get_variable('biases2', [32])

        ## Network has downsample here

        with tf.variable_scope('conv3') as scope:

            weights['conv31_weights'] = tf.get_variable('weights', shape=[3, 3, 32, 64], initializer=conv_initializer)

            weights['conv31_biases'] = tf.get_variable('biases', [64])

            weights['conv32_weights'] = tf.get_variable('weights2', shape=[3, 3, 64, 64], initializer=conv_initializer)

            weights['conv32_biases'] = tf.get_variable('biases2', [64])

        with tf.variable_scope('conv4') as scope:

            weights['conv41_weights'] = tf.get_variable('weights', shape=[3, 3, 64, 128], initializer=conv_initializer)

            weights['conv41_biases'] = tf.get_variable('biases', [128])

            weights['conv42_weights'] = tf.get_variable('weights2', shape=[3, 3, 128, 128], initializer=conv_initializer)

            weights['conv42_biases'] = tf.get_variable('biases2', [128])

        ## Network has downsample here

        with tf.variable_scope('conv5') as scope:

            weights['conv51_weights'] = tf.get_variable('weights', shape=[3, 3, 128, 256], initializer=conv_initializer)

            weights['conv51_biases'] = tf.get_variable('biases', [256])

            weights['conv52_weights'] = tf.get_variable('weights2', shape=[3, 3, 256, 256], initializer=conv_initializer)

            weights['conv52_biases'] = tf.get_variable('biases2', [256])

        with tf.variable_scope('deconv6') as scope:

            weights['deconv6_weights'] = tf.get_variable('weights0', shape=[3, 3, 128, 256], initializer=conv_initializer)

            weights['deconv6_biases'] = tf.get_variable('biases0', shape=[128], initializer=conv_initializer)

            weights['conv61_weights'] = tf.get_variable('weights', shape=[3, 3, 256, 128], initializer=conv_initializer)

            weights['conv61_biases'] = tf.get_variable('biases', [128])

            weights['conv62_weights'] = tf.get_variable('weights2', shape=[3, 3, 128, 128], initializer=conv_initializer)

            weights['conv62_biases'] = tf.get_variable('biases2', [128])

        with tf.variable_scope('deconv7') as scope:

            weights['deconv7_weights'] = tf.get_variable('weights0', shape=[3, 3, 64, 128], initializer=conv_initializer)

            weights['deconv7_biases'] = tf.get_variable('biases0', shape=[64], initializer=conv_initializer)

            weights['conv71_weights'] = tf.get_variable('weights', shape=[3, 3, 128, 64], initializer=conv_initializer)

            weights['conv71_biases'] = tf.get_variable('biases', [64])

            weights['conv72_weights'] = tf.get_variable('weights2', shape=[3, 3, 64, 64], initializer=conv_initializer)

            weights['conv72_biases'] = tf.get_variable('biases2', [64])

        with tf.variable_scope('deconv8') as scope:

            weights['deconv8_weights'] = tf.get_variable('weights0', shape=[3, 3, 32, 64], initializer=conv_initializer)

            weights['deconv8_biases'] = tf.get_variable('biases0', shape=[32], initializer=conv_initializer)

            weights['conv81_weights'] = tf.get_variable('weights', shape=[3, 3, 64, 32], initializer=conv_initializer)

            weights['conv81_biases'] = tf.get_variable('biases', [32])

            weights['conv82_weights'] = tf.get_variable('weights2', shape=[3, 3, 32, 32], initializer=conv_initializer)

            weights['conv82_biases'] = tf.get_variable('biases2', [32])

        with tf.variable_scope('deconv9') as scope:

            weights['deconv9_weights'] = tf.get_variable('weights0', shape=[3, 3, 16, 32], initializer=conv_initializer)

            weights['deconv9_biases'] = tf.get_variable('biases0', shape=[16], initializer=conv_initializer)

            weights['conv91_weights'] = tf.get_variable('weights', shape=[3, 3, 32, 16], initializer=conv_initializer)

            weights['conv91_biases'] = tf.get_variable('biases', [16])

            weights['conv92_weights'] = tf.get_variable('weights2', shape=[3, 3, 16, 16], initializer=conv_initializer)

            weights['conv92_biases'] = tf.get_variable('biases2', [16])

        with tf.variable_scope('output') as scope:

            weights['output_weights'] = tf.get_variable('weights', shape=[3, 3, 16, 2], initializer=conv_initializer)

            weights['output_biases'] = tf.get_variable('biases', [2])

        return weights

    def forward_unet(self, inp, weights, is_training=True):

        self.conv11 = conv_block(inp, weights['conv11_weights'], weights['conv11_biases'], scope='conv1/bn1', bn=False, is_training=is_training)

        self.conv12 = conv_block(self.conv11, weights['conv12_weights'], weights['conv12_biases'], scope='conv1/bn2', is_training=is_training)

        self.pool11 = max_pool(self.conv12, 2, 2, 2, 2, padding='VALID')

        # 192x192x16

        self.conv21 = conv_block(self.pool11, weights['conv21_weights'], weights['conv21_biases'], scope='conv2/bn1', is_training=is_training)

        self.conv22 = conv_block(self.conv21, weights['conv22_weights'], weights['conv22_biases'], scope='conv2/bn2', is_training=is_training)

        self.pool21 = max_pool(self.conv22, 2, 2, 2, 2, padding='VALID')

        # 96x96x32

        self.conv31 = conv_block(self.pool21, weights['conv31_weights'], weights['conv31_biases'], scope='conv3/bn1', is_training=is_training)

        self.conv32 = conv_block(self.conv31, weights['conv32_weights'], weights['conv32_biases'], scope='conv3/bn2', is_training=is_training)

        self.pool31 = max_pool(self.conv32, 2, 2, 2, 2, padding='VALID')

        # 48x48x64

        self.conv41 = conv_block(self.pool31, weights['conv41_weights'], weights['conv41_biases'], scope='conv4/bn1', is_training=is_training)

        self.conv42 = conv_block(self.conv41, weights['conv42_weights'], weights['conv42_biases'], scope='conv4/bn2', is_training=is_training)

        self.pool41 = max_pool(self.conv42, 2, 2, 2, 2, padding='VALID')

        # 24x24x128

        self.conv51 = conv_block(self.pool41, weights['conv51_weights'], weights['conv51_biases'], scope='conv5/bn1', is_training=is_training)

        self.conv52 = conv_block(self.conv51, weights['conv52_weights'], weights['conv52_biases'], scope='conv5/bn2', is_training=is_training)

        # 24x24x256

        ## add upsampling, meanwhile, channel number is reduced to half

        self.deconv6 = deconv_block(self.conv52, weights['deconv6_weights'], weights['deconv6_biases'], scope='deconv/bn6', is_training=is_training)

        # 48x48x128

        self.sum6 = concat2d(self.deconv6, self.deconv6)

        self.conv61 = conv_block(self.sum6, weights['conv61_weights'], weights['conv61_biases'], scope='conv6/bn1', is_training=is_training)

        self.conv62 = conv_block(self.conv61, weights['conv62_weights'], weights['conv62_biases'], scope='conv6/bn2', is_training=is_training)

        # 48x48x128

        self.deconv7 = deconv_block(self.conv62, weights['deconv7_weights'], weights['deconv7_biases'], scope='deconv/bn7', is_training=is_training)

        # 96x96x64

        self.sum7 = concat2d(self.deconv7, self.deconv7)

        self.conv71 = conv_block(self.sum7, weights['conv71_weights'], weights['conv71_biases'], scope='conv7/bn1', is_training=is_training)

        self.conv72 = conv_block(self.conv71, weights['conv72_weights'], weights['conv72_biases'], scope='conv7/bn2', is_training=is_training)

        # 96x96x64

        self.deconv8 = deconv_block(self.conv72, weights['deconv8_weights'], weights['deconv8_biases'], scope='deconv/bn8', is_training=is_training)

        # 192x192x32

        self.sum8 = concat2d(self.deconv8, self.deconv8)

        self.conv81 = conv_block(self.sum8, weights['conv81_weights'], weights['conv81_biases'], scope='conv8/bn1', is_training=is_training)

        self.conv82 = conv_block(self.conv81, weights['conv82_weights'], weights['conv82_biases'], scope='conv8/bn2', is_training=is_training)

        self.conv82_resize = tf.image.resize_images(self.conv82, [384, 384], method=tf.image.ResizeMethod.BILINEAR, align_corners=False)

        # 192x192x32

        self.deconv9 = deconv_block(self.conv82, weights['deconv9_weights'], weights['deconv9_biases'], scope='deconv/bn9', is_training=is_training)

        # 384x384x16

        self.sum9 = concat2d(self.deconv9, self.deconv9)

        self.conv91 = conv_block(self.sum9, weights['conv91_weights'], weights['conv91_biases'], scope='conv9/bn1', is_training=is_training)

        self.conv92 = conv_block(self.conv91, weights['conv92_weights'], weights['conv92_biases'], scope='conv9/bn2', is_training=is_training)

        # 384x384x16

        self.logits = conv_block(self.conv92, weights['output_weights'], weights['output_biases'], scope='outpu/bn', bn=False, is_training=is_training)

        #384x384x2

        self.pred_prob = tf.nn.softmax(self.logits) # shape [batch, w, h, num_classes]

        self.pred_compact = tf.argmax(self.pred_prob, axis=-1) # shape [batch, w, h]

        self.embeddings = concat2d(self.conv82_resize, self.conv92)

        return self.pred_prob, self.pred_compact, self.embeddings