# ---------------------------------------------------------
# Tensorflow MPC-GAN Implementation
# Licensed under The MIT License [see LICENSE for details]
# Written by Hulin Kuang
# ---------------------------------------------------------
import tensorflow as tf
# noinspection PyPep8Naming
import TensorFlow_utils as tf_utils
import numpy as np
import utils as utils

class MPCGAN(object):
    def __init__(self, sess, flags, image_size,channels):
        self.sess = sess
        self.flags = flags
        self.image_size = image_size
        self.channels = channels

        self.alpha_recip = 1. / self.flags.ratio_gan2seg if self.flags.ratio_gan2seg > 0 else 0
        self._gen_train_ops, self._dis_train_ops = [], []
        self.gen_c, self.dis_c = 32, 32

        self._build_net()
        self._init_assign_op()  # initialize assign operations

        print('Initialized MPCGAN SUCCESS!\n')

    def _build_net(self):
        self.X = tf.placeholder(tf.float32, shape=[None, *self.image_size, 1], name='image')
        self.X1 = tf.placeholder(tf.float32, shape=[None, *self.image_size, 1], name='image1')
        self.X2 = tf.placeholder(tf.float32, shape=[None, *self.image_size, 1], name='image2')
        self.X3 = tf.placeholder(tf.float32, shape=[None, *self.image_size, 1], name='image3')
        self.Y = tf.placeholder(tf.float32, shape=[None, *self.image_size, 1], name='vessel')

        # self.g_samples_mirror = self.generator(self.X)
        self.g_samples = self.generator(self.X, self.X1, self.X2, self.X3)

        self.real_pair = tf.concat([self.X, self.X1, self.X2, self.X3, self.Y], axis=3)
        self.fake_pair = tf.concat([self.X, self.X1, self.X2, self.X3, self.g_samples], axis=3)
        # self.real_pair = tf.concat([self.X, self.g_samples_mirror, self.Y], axis=3)
        # self.fake_pair = tf.concat([self.X, self.g_samples_mirror, self.g_samples], axis=3)

        d_real, d_logit_real = self.discriminator(self.real_pair)
        d_fake, d_logit_fake = self.discriminator(self.fake_pair, is_reuse=True)#True

        # discrminator loss
        self.d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
            logits=d_logit_real, labels=tf.ones_like(d_real)))
        self.d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
            logits=d_logit_fake, labels=tf.zeros_like(d_logit_fake)))
        self.d_loss = self.d_loss_real + self.d_loss_fake

        # generator loss
        gan_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
            logits=d_logit_fake, labels=tf.ones_like(d_logit_fake)))
        seg_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.g_samples, labels=self.Y))
        #seg_loss = -dice_coeff
        self.g_loss = self.alpha_recip * gan_loss + seg_loss

        t_vars = tf.trainable_variables()
        d_vars = [var for var in t_vars if 'd_' in var.name]
        g_vars = [var for var in t_vars if 'g_' in var.name]

        dis_op = tf.train.AdamOptimizer(learning_rate=self.flags.learning_rate, beta1=self.flags.beta1)\
            .minimize(self.d_loss, var_list=d_vars)
        dis_ops = [dis_op] + self._dis_train_ops
        self.dis_optim = tf.group(*dis_ops)

        gen_op = tf.train.AdamOptimizer(learning_rate=self.flags.learning_rate, beta1=self.flags.beta1)\
            .minimize(self.g_loss, var_list=g_vars)
        gen_ops = [gen_op] + self._gen_train_ops
        self.gen_optim = tf.group(*gen_ops)

    def _init_assign_op(self):
        self.best_auc_sum_placeholder = tf.placeholder(tf.float32, name='best_auc_sum_placeholder')
        self.auc_pr_placeholder = tf.placeholder(tf.float32, name='auc_pr_placeholder')
        self.auc_roc_placeholder = tf.placeholder(tf.float32, name='auc_roc_placeholder')
        self.dice_coeff_placeholder = tf.placeholder(tf.float32, name='dice_coeff_placeholder')
        self.acc_placeholder = tf.placeholder(tf.float32, name='acc_placeholder')
        self.sensitivity_placeholder = tf.placeholder(tf.float32, name='sensitivity_placeholder')
        self.specificity_placeholder = tf.placeholder(tf.float32, name='specificity_placeholder')
        self.score_placeholder = tf.placeholder(tf.float32, name='score_placeholder')
        self.gloss_save_placeholder = tf.placeholder(tf.float32, name='gloss_placeholder')
        self.dloss_save_placeholder = tf.placeholder(tf.float32, name='dloss_placeholder')

        self.best_auc_sum = tf.Variable(0., trainable=False, dtype=tf.float32, name='best_auc_sum')
        auc_pr = tf.Variable(0., trainable=False, dtype=tf.float32, name='auc_pr')
        auc_roc = tf.Variable(0., trainable=False, dtype=tf.float32, name='auc_roc')
        dice_coeff = tf.Variable(0., trainable=False, dtype=tf.float32, name='dice_coeff')
        acc = tf.Variable(0., trainable=False, dtype=tf.float32, name='acc')
        sensitivity = tf.Variable(0., trainable=False, dtype=tf.float32, name='sensitivity')
        specificity = tf.Variable(0., trainable=False, dtype=tf.float32, name='specificity')
        score = tf.Variable(0., trainable=False, dtype=tf.float32, name='score')
        dloss_save = tf.Variable(0., trainable=False, dtype=tf.float32, name='dloss_save')
        gloss_save = tf.Variable(0., trainable=False, dtype=tf.float32, name='gloss_save')

        self.best_auc_sum_assign_op = self.best_auc_sum.assign(self.best_auc_sum_placeholder)
        auc_pr_assign_op = auc_pr.assign(self.auc_pr_placeholder)
        auc_roc_assign_op = auc_roc.assign(self.auc_roc_placeholder)
        dice_coeff_assign_op = dice_coeff.assign(self.dice_coeff_placeholder)
        acc_assign_op = acc.assign(self.acc_placeholder)
        sensitivity_assign_op = sensitivity.assign(self.sensitivity_placeholder)
        specificity_assign_op = specificity.assign(self.specificity_placeholder)
        dloss_save_op = dloss_save.assign(self.dloss_save_placeholder)
        gloss_save_op = gloss_save.assign(self.gloss_save_placeholder)
        score_assign_op = dloss_save.assign(self.score_placeholder)

        self.measure_assign_op = tf.group(auc_pr_assign_op, auc_roc_assign_op, dice_coeff_assign_op,
                                          acc_assign_op, sensitivity_assign_op, specificity_assign_op,
                                          score_assign_op)

        self.measure_loss_op = tf.group(dloss_save_op, gloss_save_op)

        # for tensorboard
        if not self.flags.is_test:
            self.writer = tf.summary.FileWriter("{}/logs/{}_{}_{}".format(
                self.flags.dataset, self.flags.discriminator, self.flags.train_interval, self.flags.batch_size))
            # self.writer_loss = tf.summary.FileWriter("{}/logs_loss/{}_{}_{}".format(
            #     self.flags.dataset, self.flags.discriminator, self.flags.train_interval, self.flags.batch_size))

        auc_pr_summ = tf.summary.scalar("auc_pr_summary", auc_pr)
        auc_roc_summ = tf.summary.scalar("auc_roc_summary", auc_roc)
        dice_coeff_summ = tf.summary.scalar("dice_coeff_summary", dice_coeff)
        acc_summ = tf.summary.scalar("acc_summary", acc)
        sensitivity_summ = tf.summary.scalar("sensitivity_summary", sensitivity)
        specificity_summ = tf.summary.scalar("specificity_summary", specificity)
        score_summ = tf.summary.scalar("score_summary", score)
        gloss_summ = tf.summary.scalar("gloss_summary", gloss_save)
        dloss_summ = tf.summary.scalar("dloss_summary", dloss_save)

        self.measure_summary = tf.summary.merge([auc_pr_summ, auc_roc_summ, dice_coeff_summ, acc_summ,
                                                 sensitivity_summ, specificity_summ, score_summ])

        self.measure_loss_summary = tf.summary.merge([ gloss_summ, dloss_summ])

    def generator(self, data, data1, data2, data3, name='g_'):
        with tf.variable_scope(name):
            # difference path
            # conv1_1: (N, 512, 512, 1) -> (N, 320, 320, 32)
            conv1_1 = tf_utils.conv2d(data1, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_1_conv1_1')
            conv1_1 = tf_utils.batch_norm(conv1_1, name='conv1_1_batch1', _ops=self._gen_train_ops)
            conv1_1 = tf.nn.relu(conv1_1, name='conv1_1_relu1')
            conv1_1 = tf_utils.conv2d(conv1_1, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_1_conv1_2')
            conv1_1 = tf_utils.batch_norm(conv1_1, name='conv1_1_batch2', _ops=self._gen_train_ops)
            conv1_1 = tf.nn.relu(conv1_1, name='conv1_1_relu2')
            pool1 = tf_utils.max_pool_2x2(conv1_1, name='maxpool1')

            # conv1_2: (N, 320, 320, 32) -> (N, 160, 160, 64)
            conv1_2 = tf_utils.conv2d(pool1, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_2_conv1_1')
            conv1_2 = tf_utils.batch_norm(conv1_2, name='conv1_2_batch1', _ops=self._gen_train_ops)
            conv1_2 = tf.nn.relu(conv1_2, name='conv1_2_relu1')
            conv1_2 = tf_utils.conv2d(conv1_2, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_2_conv1_2')
            conv1_2 = tf_utils.batch_norm(conv1_2, name='conv1_2-batch2', _ops=self._gen_train_ops)
            conv1_2 = tf.nn.relu(conv1_2, name='conv1_2_relu2')
            pool2 = tf_utils.max_pool_2x2(conv1_2, name='maxpool2')

            # conv1_3: (N, 160, 160, 64) -> (N, 80, 80, 128)
            conv1_3 = tf_utils.conv2d(pool2, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_3_conv1_1')
            conv1_3 = tf_utils.batch_norm(conv1_3, name='conv1_3_batch1', _ops=self._gen_train_ops)
            conv1_3 = tf.nn.relu(conv1_3, name='conv1_3_relu1')
            conv1_3 = tf_utils.conv2d(conv1_3, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_3_conv1_2')
            conv1_3 = tf_utils.batch_norm(conv1_3, name='conv1_3_batch2', _ops=self._gen_train_ops)
            conv1_3 = tf.nn.relu(conv1_3, name='conv1_3_relu2')
            pool3 = tf_utils.max_pool_2x2(conv1_3, name='maxpool3')

            # conv1_4: (N, 80, 80, 128) -> (N, 40, 40, 256)
            conv1_4 = tf_utils.conv2d(pool3, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_4_conv1_1')
            conv1_4 = tf_utils.batch_norm(conv1_4, name='conv1_4_batch1', _ops=self._gen_train_ops)
            conv1_4 = tf.nn.relu(conv1_4, name='conv1_4_relu1')
            conv1_4 = tf_utils.conv2d(conv1_4, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_4_conv1_2')
            conv1_4 = tf_utils.batch_norm(conv1_4, name='conv1_4_batch2', _ops=self._gen_train_ops)
            conv1_4 = tf.nn.relu(conv1_4, name='conv1_4_relu2')
            pool4 = tf_utils.max_pool_2x2(conv1_4, name='maxpool4')

            # conv1_5: (N, 40, 40, 256) -> (N, 40, 40, 512)
            conv1_5 = tf_utils.conv2d(pool4, 16 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_5_conv1_1')
            conv1_5 = tf_utils.batch_norm(conv1_5, name='conv1_5_batch1', _ops=self._gen_train_ops)
            conv1_5 = tf.nn.relu(conv1_5, name='conv1_5_relu1')
            conv1_5 = tf_utils.conv2d(conv1_5, 16 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_5_conv1_2')
            conv1_5 = tf_utils.batch_norm(conv1_5, name='conv1_5_batch2', _ops=self._gen_train_ops)
            conv1_5 = tf.nn.relu(conv1_5, name='conv1_5_relu2')

            # conv1_6: (N, 40, 40, 512) -> (N, 80, 80, 256)
            up1 = tf_utils.upsampling2d(conv1_5, size=(2, 2), name='conv1_6_up')
            conv1_6 = tf.concat([up1, conv1_4], axis=3, name='conv1_6_concat')
            conv1_6 = tf_utils.conv2d(conv1_6, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_6_conv1_1')
            conv1_6 = tf_utils.batch_norm(conv1_6, name='conv1_6_batch1', _ops=self._gen_train_ops)
            conv1_6 = tf.nn.relu(conv1_6, name='conv1_6_relu1')
            conv1_6 = tf_utils.conv2d(conv1_6, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_6_conv1_2')
            conv1_6 = tf_utils.batch_norm(conv1_6, name='conv1_6_batch2', _ops=self._gen_train_ops)
            conv1_6 = tf.nn.relu(conv1_6, name='conv1_6_relu2')

            # conv1_7: (N, 80, 80, 256) -> (N, 160, 160, 128)
            up2 = tf_utils.upsampling2d(conv1_6, size=(2, 2), name='conv1_7_up')
            conv1_7 = tf.concat([up2, conv1_3], axis=3, name='conv1_7_concat')
            conv1_7 = tf_utils.conv2d(conv1_7, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_7_conv1_1')
            conv1_7 = tf_utils.batch_norm(conv1_7, name='conv1_7_batch1', _ops=self._gen_train_ops)
            conv1_7 = tf.nn.relu(conv1_7, name='conv1_7_relu1')
            conv1_7 = tf_utils.conv2d(conv1_7, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_7_conv1_2')
            conv1_7 = tf_utils.batch_norm(conv1_7, name='conv1_7_batch2', _ops=self._gen_train_ops)
            conv1_7 = tf.nn.relu(conv1_7, name='conv1_7_relu2')

            # conv1_8: (N, 160, 160, 128) -> (N, 320, 320, 64)
            up3 = tf_utils.upsampling2d(conv1_7, size=(2, 2), name='conv1_8_up')
            conv1_8 = tf.concat([up3, conv1_2], axis=3, name='conv1_8_concat')
            conv1_8 = tf_utils.conv2d(conv1_8, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_8_conv1_1')
            conv1_8 = tf_utils.batch_norm(conv1_8, name='conv1_8_batch1', _ops=self._gen_train_ops)
            conv1_8 = tf.nn.relu(conv1_8, name='conv1_8_relu1')
            conv1_8 = tf_utils.conv2d(conv1_8, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_8_conv1_2')
            conv1_8 = tf_utils.batch_norm(conv1_8, name='conv1_8_batch2', _ops=self._gen_train_ops)
            conv1_8 = tf.nn.relu(conv1_8, name='conv1_8_relu2')

            # conv1_9: (N, 320, 320, 64) -> (N, 512, 512, 32)
            up4 = tf_utils.upsampling2d(conv1_8, size=(2, 2), name='conv1_9_up')
            conv1_9 = tf.concat([up4, conv1_1], axis=3, name='conv1_9_concat')
            conv1_9 = tf_utils.conv2d(conv1_9, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_9_conv1_1')
            conv1_9 = tf_utils.batch_norm(conv1_9, name='conv1_9_batch1', _ops=self._gen_train_ops)
            conv1_9 = tf.nn.relu(conv1_9, name='conv1_9_relu1')
            conv1_9 = tf_utils.conv2d(conv1_9, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_9_conv1_2')
            conv1_9 = tf_utils.batch_norm(conv1_9, name='conv1_9_batch2', _ops=self._gen_train_ops)
            conv1_9 = tf.nn.relu(conv1_9, name='conv1_9_relu2')

            # distance path
            # conv2_1: (N, 512, 512, 1) -> (N, 320, 320, 32)
            conv2_1 = tf_utils.conv2d(data2, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_1_conv2_1')
            conv2_1 = tf_utils.batch_norm(conv2_1, name='conv2_1_batch1', _ops=self._gen_train_ops)
            conv2_1 = tf.nn.relu(conv2_1, name='conv2_1_relu1')
            conv2_1 = tf_utils.conv2d(conv2_1, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_1_conv2_2')
            conv2_1 = tf_utils.batch_norm(conv2_1, name='conv2_1_batch2', _ops=self._gen_train_ops)
            conv2_1 = tf.nn.relu(conv2_1, name='conv2_1_relu2')
            pool1 = tf_utils.max_pool_2x2(conv2_1, name='maxpool1')

            # conv2_2: (N, 320, 320, 32) -> (N, 160, 160, 64)
            conv2_2 = tf_utils.conv2d(pool1, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_2_conv2_1')
            conv2_2 = tf_utils.batch_norm(conv2_2, name='conv2_2_batch1', _ops=self._gen_train_ops)
            conv2_2 = tf.nn.relu(conv2_2, name='conv2_2_relu1')
            conv2_2 = tf_utils.conv2d(conv2_2, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_2_conv2_2')
            conv2_2 = tf_utils.batch_norm(conv2_2, name='conv2_2-batch2', _ops=self._gen_train_ops)
            conv2_2 = tf.nn.relu(conv2_2, name='conv2_2_relu2')
            pool2 = tf_utils.max_pool_2x2(conv2_2, name='maxpool2')

            # conv2_3: (N, 160, 160, 64) -> (N, 80, 80, 128)
            conv2_3 = tf_utils.conv2d(pool2, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_3_conv2_1')
            conv2_3 = tf_utils.batch_norm(conv2_3, name='conv2_3_batch1', _ops=self._gen_train_ops)
            conv2_3 = tf.nn.relu(conv2_3, name='conv2_3_relu1')
            conv2_3 = tf_utils.conv2d(conv2_3, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_3_conv2_2')
            conv2_3 = tf_utils.batch_norm(conv2_3, name='conv2_3_batch2', _ops=self._gen_train_ops)
            conv2_3 = tf.nn.relu(conv2_3, name='conv2_3_relu2')
            pool3 = tf_utils.max_pool_2x2(conv2_3, name='maxpool3')

            # conv2_4: (N, 80, 80, 128) -> (N, 40, 40, 256)
            conv2_4 = tf_utils.conv2d(pool3, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_4_conv2_1')
            conv2_4 = tf_utils.batch_norm(conv2_4, name='conv2_4_batch1', _ops=self._gen_train_ops)
            conv2_4 = tf.nn.relu(conv2_4, name='conv2_4_relu1')
            conv2_4 = tf_utils.conv2d(conv2_4, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_4_conv2_2')
            conv2_4 = tf_utils.batch_norm(conv2_4, name='conv2_4_batch2', _ops=self._gen_train_ops)
            conv2_4 = tf.nn.relu(conv2_4, name='conv2_4_relu2')
            pool4 = tf_utils.max_pool_2x2(conv2_4, name='maxpool4')

            # conv2_5: (N, 40, 40, 256) -> (N, 40, 40, 512)
            conv2_5 = tf_utils.conv2d(pool4, 16 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_5_conv2_1')
            conv2_5 = tf_utils.batch_norm(conv2_5, name='conv2_5_batch1', _ops=self._gen_train_ops)
            conv2_5 = tf.nn.relu(conv2_5, name='conv2_5_relu1')
            conv2_5 = tf_utils.conv2d(conv2_5, 16 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_5_conv2_2')
            conv2_5 = tf_utils.batch_norm(conv2_5, name='conv2_5_batch2', _ops=self._gen_train_ops)
            conv2_5 = tf.nn.relu(conv2_5, name='conv2_5_relu2')

            # conv2_6: (N, 40, 40, 512) -> (N, 80, 80, 256)
            up1 = tf_utils.upsampling2d(conv2_5, size=(2, 2), name='conv2_6_up')
            conv2_6 = tf.concat([up1, conv2_4], axis=3, name='conv2_6_concat')
            conv2_6 = tf_utils.conv2d(conv2_6, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_6_conv2_1')
            conv2_6 = tf_utils.batch_norm(conv2_6, name='conv2_6_batch1', _ops=self._gen_train_ops)
            conv2_6 = tf.nn.relu(conv2_6, name='conv2_6_relu1')
            conv2_6 = tf_utils.conv2d(conv2_6, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_6_conv2_2')
            conv2_6 = tf_utils.batch_norm(conv2_6, name='conv2_6_batch2', _ops=self._gen_train_ops)
            conv2_6 = tf.nn.relu(conv2_6, name='conv2_6_relu2')

            # conv2_7: (N, 80, 80, 256) -> (N, 160, 160, 128)
            up2 = tf_utils.upsampling2d(conv2_6, size=(2, 2), name='conv2_7_up')
            conv2_7 = tf.concat([up2, conv2_3], axis=3, name='conv2_7_concat')
            conv2_7 = tf_utils.conv2d(conv2_7, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_7_conv2_1')
            conv2_7 = tf_utils.batch_norm(conv2_7, name='conv2_7_batch1', _ops=self._gen_train_ops)
            conv2_7 = tf.nn.relu(conv2_7, name='conv2_7_relu1')
            conv2_7 = tf_utils.conv2d(conv2_7, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_7_conv2_2')
            conv2_7 = tf_utils.batch_norm(conv2_7, name='conv2_7_batch2', _ops=self._gen_train_ops)
            conv2_7 = tf.nn.relu(conv2_7, name='conv2_7_relu2')

            # conv2_8: (N, 160, 160, 128) -> (N, 320, 320, 64)
            up3 = tf_utils.upsampling2d(conv2_7, size=(2, 2), name='conv2_8_up')
            conv2_8 = tf.concat([up3, conv2_2], axis=3, name='conv2_8_concat')
            conv2_8 = tf_utils.conv2d(conv2_8, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_8_conv2_1')
            conv2_8 = tf_utils.batch_norm(conv2_8, name='conv2_8_batch1', _ops=self._gen_train_ops)
            conv2_8 = tf.nn.relu(conv2_8, name='conv2_8_relu1')
            conv2_8 = tf_utils.conv2d(conv2_8, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_8_conv2_2')
            conv2_8 = tf_utils.batch_norm(conv2_8, name='conv2_8_batch2', _ops=self._gen_train_ops)
            conv2_8 = tf.nn.relu(conv2_8, name='conv2_8_relu2')

            # conv2_9: (N, 320, 320, 64) -> (N, 512, 512, 32)
            up4 = tf_utils.upsampling2d(conv2_8, size=(2, 2), name='conv2_9_up')
            conv2_9 = tf.concat([up4, conv2_1], axis=3, name='conv2_9_concat')
            conv2_9 = tf_utils.conv2d(conv2_9, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_9_conv2_1')
            conv2_9 = tf_utils.batch_norm(conv2_9, name='conv2_9_batch1', _ops=self._gen_train_ops)
            conv2_9 = tf.nn.relu(conv2_9, name='conv2_9_relu1')
            conv2_9 = tf_utils.conv2d(conv2_9, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_9_conv2_2')
            conv2_9 = tf_utils.batch_norm(conv2_9, name='conv2_9_batch2', _ops=self._gen_train_ops)
            conv2_9 = tf.nn.relu(conv2_9, name='conv2_9_relu2')

            # location path
            # conv3_1: (N, 512, 512, 1) -> (N, 320, 320, 32)
            conv3_1 = tf_utils.conv2d(data3, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_1_conv3_1')
            conv3_1 = tf_utils.batch_norm(conv3_1, name='conv3_1_batch1', _ops=self._gen_train_ops)
            conv3_1 = tf.nn.relu(conv3_1, name='conv3_1_relu1')
            conv3_1 = tf_utils.conv2d(conv3_1, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_1_conv3_2')
            conv3_1 = tf_utils.batch_norm(conv3_1, name='conv3_1_batch2', _ops=self._gen_train_ops)
            conv3_1 = tf.nn.relu(conv3_1, name='conv3_1_relu2')
            pool1 = tf_utils.max_pool_2x2(conv3_1, name='maxpool1')

            # conv3_2: (N, 320, 320, 32) -> (N, 160, 160, 64)
            conv3_2 = tf_utils.conv2d(pool1, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_2_conv3_1')
            conv3_2 = tf_utils.batch_norm(conv3_2, name='conv3_2_batch1', _ops=self._gen_train_ops)
            conv3_2 = tf.nn.relu(conv3_2, name='conv3_2_relu1')
            conv3_2 = tf_utils.conv2d(conv3_2, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_2_conv3_2')
            conv3_2 = tf_utils.batch_norm(conv3_2, name='conv3_2-batch2', _ops=self._gen_train_ops)
            conv3_2 = tf.nn.relu(conv3_2, name='conv3_2_relu2')
            pool2 = tf_utils.max_pool_2x2(conv3_2, name='maxpool2')

            # conv3_3: (N, 160, 160, 64) -> (N, 80, 80, 128)
            conv3_3 = tf_utils.conv2d(pool2, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_3_conv3_1')
            conv3_3 = tf_utils.batch_norm(conv3_3, name='conv3_3_batch1', _ops=self._gen_train_ops)
            conv3_3 = tf.nn.relu(conv3_3, name='conv3_3_relu1')
            conv3_3 = tf_utils.conv2d(conv3_3, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_3_conv3_2')
            conv3_3 = tf_utils.batch_norm(conv3_3, name='conv3_3_batch2', _ops=self._gen_train_ops)
            conv3_3 = tf.nn.relu(conv3_3, name='conv3_3_relu2')
            pool3 = tf_utils.max_pool_2x2(conv3_3, name='maxpool3')

            # conv3_4: (N, 80, 80, 128) -> (N, 40, 40, 256)
            conv3_4 = tf_utils.conv2d(pool3, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_4_conv3_1')
            conv3_4 = tf_utils.batch_norm(conv3_4, name='conv3_4_batch1', _ops=self._gen_train_ops)
            conv3_4 = tf.nn.relu(conv3_4, name='conv3_4_relu1')
            conv3_4 = tf_utils.conv2d(conv3_4, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_4_conv3_2')
            conv3_4 = tf_utils.batch_norm(conv3_4, name='conv3_4_batch2', _ops=self._gen_train_ops)
            conv3_4 = tf.nn.relu(conv3_4, name='conv3_4_relu2')
            pool4 = tf_utils.max_pool_2x2(conv3_4, name='maxpool4')

            # conv3_5: (N, 40, 40, 256) -> (N, 40, 40, 512)
            conv3_5 = tf_utils.conv2d(pool4, 16 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_5_conv3_1')
            conv3_5 = tf_utils.batch_norm(conv3_5, name='conv3_5_batch1', _ops=self._gen_train_ops)
            conv3_5 = tf.nn.relu(conv3_5, name='conv3_5_relu1')
            conv3_5 = tf_utils.conv2d(conv3_5, 16 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_5_conv3_2')
            conv3_5 = tf_utils.batch_norm(conv3_5, name='conv3_5_batch2', _ops=self._gen_train_ops)
            conv3_5 = tf.nn.relu(conv3_5, name='conv3_5_relu2')

            # conv3_6: (N, 40, 40, 512) -> (N, 80, 80, 256)
            up1 = tf_utils.upsampling2d(conv3_5, size=(2, 2), name='conv3_6_up')
            conv3_6 = tf.concat([up1, conv3_4], axis=3, name='conv3_6_concat')
            conv3_6 = tf_utils.conv2d(conv3_6, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_6_conv3_1')
            conv3_6 = tf_utils.batch_norm(conv3_6, name='conv3_6_batch1', _ops=self._gen_train_ops)
            conv3_6 = tf.nn.relu(conv3_6, name='conv3_6_relu1')
            conv3_6 = tf_utils.conv2d(conv3_6, 8 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_6_conv3_2')
            conv3_6 = tf_utils.batch_norm(conv3_6, name='conv3_6_batch2', _ops=self._gen_train_ops)
            conv3_6 = tf.nn.relu(conv3_6, name='conv3_6_relu2')

            # conv3_7: (N, 80, 80, 256) -> (N, 160, 160, 128)
            up2 = tf_utils.upsampling2d(conv3_6, size=(2, 2), name='conv3_7_up')
            conv3_7 = tf.concat([up2, conv3_3], axis=3, name='conv3_7_concat')
            conv3_7 = tf_utils.conv2d(conv3_7, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_7_conv3_1')
            conv3_7 = tf_utils.batch_norm(conv3_7, name='conv3_7_batch1', _ops=self._gen_train_ops)
            conv3_7 = tf.nn.relu(conv3_7, name='conv3_7_relu1')
            conv3_7 = tf_utils.conv2d(conv3_7, 4 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_7_conv3_2')
            conv3_7 = tf_utils.batch_norm(conv3_7, name='conv3_7_batch2', _ops=self._gen_train_ops)
            conv3_7 = tf.nn.relu(conv3_7, name='conv3_7_relu2')

            # conv3_8: (N, 160, 160, 128) -> (N, 320, 320, 64)
            up3 = tf_utils.upsampling2d(conv3_7, size=(2, 2), name='conv3_8_up')
            conv3_8 = tf.concat([up3, conv3_2], axis=3, name='conv3_8_concat')
            conv3_8 = tf_utils.conv2d(conv3_8, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_8_conv3_1')
            conv3_8 = tf_utils.batch_norm(conv3_8, name='conv3_8_batch1', _ops=self._gen_train_ops)
            conv3_8 = tf.nn.relu(conv3_8, name='conv3_8_relu1')
            conv3_8 = tf_utils.conv2d(conv3_8, 2 * self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_8_conv3_2')
            conv3_8 = tf_utils.batch_norm(conv3_8, name='conv3_8_batch2', _ops=self._gen_train_ops)
            conv3_8 = tf.nn.relu(conv3_8, name='conv3_8_relu2')

            # conv3_9: (N, 320, 320, 64) -> (N, 512, 512, 32)
            up4 = tf_utils.upsampling2d(conv3_8, size=(2, 2), name='conv3_9_up')
            conv3_9 = tf.concat([up4, conv3_1], axis=3, name='conv3_9_concat')
            conv3_9 = tf_utils.conv2d(conv3_9, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_9_conv3_1')
            conv3_9 = tf_utils.batch_norm(conv3_9, name='conv3_9_batch1', _ops=self._gen_train_ops)
            conv3_9 = tf.nn.relu(conv3_9, name='conv3_9_relu1')
            conv3_9 = tf_utils.conv2d(conv3_9, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_9_conv3_2')
            conv3_9 = tf_utils.batch_norm(conv3_9, name='conv3_9_batch2', _ops=self._gen_train_ops)
            conv3_9 = tf.nn.relu(conv3_9, name='conv3_9_relu2')

            # intensity path
            # conv1: (N, 512, 512, 1) -> (N, 320, 320, 32)
            conv1 = tf_utils.conv2d(data, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_conv1')
            conv1 = tf_utils.batch_norm(conv1, name='conv1_batch1', _ops=self._gen_train_ops)
            conv1 = tf.nn.relu(conv1, name='conv1_relu1')
            conv1 = tf_utils.conv2d(conv1, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_conv2')
            conv1 = tf_utils.batch_norm(conv1, name='conv1_batch2', _ops=self._gen_train_ops)
            conv1 = tf.nn.relu(conv1, name='conv1_relu2')
            conv1 = tf.concat([conv1, conv1_1, conv2_1, conv3_1], axis=3, name='conv1_concat')

            pool1 = tf_utils.max_pool_2x2(conv1, name='maxpool1')

            # conv2: (N, 320, 320, 32) -> (N, 160, 160, 64)
            conv2 = tf_utils.conv2d(pool1, 2*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_conv1')
            conv2 = tf_utils.batch_norm(conv2, name='conv2_batch1', _ops=self._gen_train_ops)
            conv2 = tf.nn.relu(conv2, name='conv2_relu1')
            conv2 = tf_utils.conv2d(conv2, 2*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_conv2')
            conv2 = tf_utils.batch_norm(conv2, name='conv2-batch2', _ops=self._gen_train_ops)
            conv2 = tf.nn.relu(conv2, name='conv2_relu2')
            conv2 = tf.concat([conv2, conv1_2, conv2_2, conv3_2], axis=3, name='conv2_concat')
            pool2 = tf_utils.max_pool_2x2(conv2, name='maxpool2')

            # conv3: (N, 160, 160, 64) -> (N, 80, 80, 128)
            conv3 = tf_utils.conv2d(pool2, 4*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_conv1')
            conv3 = tf_utils.batch_norm(conv3, name='conv3_batch1', _ops=self._gen_train_ops)
            conv3 = tf.nn.relu(conv3, name='conv3_relu1')
            conv3 = tf_utils.conv2d(conv3, 4*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_conv2')
            conv3 = tf_utils.batch_norm(conv3, name='conv3_batch2', _ops=self._gen_train_ops)
            conv3 = tf.nn.relu(conv3, name='conv3_relu2')
            conv3 = tf.concat([conv3, conv1_3, conv2_3, conv3_3], axis=3, name='conv3_concat')
            pool3 = tf_utils.max_pool_2x2(conv3, name='maxpool3')

            # conv4: (N, 80, 80, 128) -> (N, 40, 40, 256)
            conv4 = tf_utils.conv2d(pool3, 8*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv4_conv1')
            conv4 = tf_utils.batch_norm(conv4, name='conv4_batch1', _ops=self._gen_train_ops)
            conv4 = tf.nn.relu(conv4, name='conv4_relu1')
            conv4 = tf_utils.conv2d(conv4, 8*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv4_conv2')
            conv4 = tf_utils.batch_norm(conv4, name='conv4_batch2', _ops=self._gen_train_ops)
            conv4 = tf.nn.relu(conv4, name='conv4_relu2')
            conv4 = tf.concat([conv4, conv1_4, conv2_4, conv3_4], axis=3, name='conv4_concat')
            pool4 = tf_utils.max_pool_2x2(conv4, name='maxpool4')

            # conv5: (N, 40, 40, 256) -> (N, 40, 40, 512)
            conv5 = tf_utils.conv2d(pool4, 16*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv5_conv1')
            conv5 = tf_utils.batch_norm(conv5, name='conv5_batch1', _ops=self._gen_train_ops)
            conv5 = tf.nn.relu(conv5, name='conv5_relu1')
            conv5 = tf_utils.conv2d(conv5, 16*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv5_conv2')
            conv5 = tf_utils.batch_norm(conv5, name='conv5_batch2', _ops=self._gen_train_ops)
            conv5 = tf.nn.relu(conv5, name='conv5_relu2')

            # conv6: (N, 40, 40, 512) -> (N, 80, 80, 256)
            up1 = tf_utils.upsampling2d(conv5, size=(2, 2), name='conv6_up')
            conv6 = tf.concat([up1, conv4], axis=3, name='conv6_concat')
            conv6 = tf_utils.conv2d(conv6, 8*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv6_conv1')
            conv6 = tf_utils.batch_norm(conv6, name='conv6_batch1', _ops=self._gen_train_ops)
            conv6 = tf.nn.relu(conv6, name='conv6_relu1')
            conv6 = tf_utils.conv2d(conv6, 8*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv6_conv2')
            conv6 = tf_utils.batch_norm(conv6, name='conv6_batch2', _ops=self._gen_train_ops)
            conv6 = tf.nn.relu(conv6, name='conv6_relu2')
            conv6 = tf.concat([conv6, conv1_6, conv2_6, conv3_6], axis=3, name='conv6_concat')

            # conv7: (N, 80, 80, 256) -> (N, 160, 160, 128)
            up2 = tf_utils.upsampling2d(conv6, size=(2, 2), name='conv7_up')
            conv7 = tf.concat([up2, conv3], axis=3, name='conv7_concat')
            conv7 = tf_utils.conv2d(conv7, 4*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv7_conv1')
            conv7 = tf_utils.batch_norm(conv7, name='conv7_batch1', _ops=self._gen_train_ops)
            conv7 = tf.nn.relu(conv7, name='conv7_relu1')
            conv7 = tf_utils.conv2d(conv7, 4*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv7_conv2')
            conv7 = tf_utils.batch_norm(conv7, name='conv7_batch2', _ops=self._gen_train_ops)
            conv7 = tf.nn.relu(conv7, name='conv7_relu2')
            conv7 = tf.concat([conv7, conv1_7, conv2_7, conv3_7], axis=3, name='conv7_concat')

            # conv8: (N, 160, 160, 128) -> (N, 320, 320, 64)
            up3 = tf_utils.upsampling2d(conv7, size=(2, 2), name='conv8_up')
            conv8 = tf.concat([up3, conv2], axis=3, name='conv8_concat')
            conv8 = tf_utils.conv2d(conv8, 2*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv8_conv1')
            conv8 = tf_utils.batch_norm(conv8, name='conv8_batch1', _ops=self._gen_train_ops)
            conv8 = tf.nn.relu(conv8, name='conv8_relu1')
            conv8 = tf_utils.conv2d(conv8, 2*self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv8_conv2')
            conv8 = tf_utils.batch_norm(conv8, name='conv8_batch2', _ops=self._gen_train_ops)
            conv8 = tf.nn.relu(conv8, name='conv8_relu2')
            conv8 = tf.concat([conv8, conv1_8, conv2_8, conv3_8], axis=3, name='conv8_concat')

            # conv9: (N, 320, 320, 64) -> (N, 512, 512, 32)
            up4 = tf_utils.upsampling2d(conv8, size=(2, 2), name='conv9_up')
            conv9 = tf.concat([up4, conv1], axis=3, name='conv9_concat')
            conv9 = tf_utils.conv2d(conv9, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv9_conv1')
            conv9 = tf_utils.batch_norm(conv9, name='conv9_batch1', _ops=self._gen_train_ops)
            conv9 = tf.nn.relu(conv9, name='conv9_relu1')
            conv9 = tf_utils.conv2d(conv9, self.gen_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv9_conv2')
            conv9 = tf_utils.batch_norm(conv9, name='conv9_batch2', _ops=self._gen_train_ops)
            conv9 = tf.nn.relu(conv9, name='conv9_relu2')
            conv9 = tf.concat([conv9, conv1_9, conv2_9, conv3_9], axis=3, name='conv1_concat')

            # output layer: (N, 512, 512, 32) -> (N, 512, 512, 1)
            output = tf_utils.conv2d(conv9, 1, k_h=1, k_w=1, d_h=1, d_w=1, name='conv_output')

            return tf.nn.sigmoid(output)

    def discriminator(self, data, name='d_', is_reuse=False):
        with tf.variable_scope(name) as scope:
            if is_reuse is True:
                scope.reuse_variables()

            # conv1: (N, 512, 512, 4) -> (N,, 512, 512, 32)
            conv1 = tf_utils.conv2d(data, self.dis_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv1_conv1')
            conv1 = tf_utils.lrelu(conv1, name='conv1_lrelu1')

            # conv2: (N, 512, 512, 32) -> (N, 512, 512, 64)
            conv2 = tf_utils.conv2d(conv1, 2*self.dis_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv2_conv1')
            conv2 = tf_utils.lrelu(conv2)

            # conv3: (N, 512, 512, 64) -> (N, 512, 512, 128)
            conv3 = tf_utils.conv2d(conv2, 4*self.dis_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv3_conv1')
            conv3 = tf_utils.lrelu(conv3)

            # conv4: (N, 512, 512, 128) -> (N, 512, 512, 256)
            conv4 = tf_utils.conv2d(conv3, 8*self.dis_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv4_conv1')
            conv4 = tf_utils.lrelu(conv4)

            # conv5: (N, 512, 512, 256) -> (N, 512, 512, 512)
            conv5 = tf_utils.conv2d(conv4, 16*self.dis_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv5_conv1')
            conv5 = tf_utils.lrelu(conv5)

            # conv6: (N, 512, 512, 512) -> (N, 512, 512, 1024)
            conv6 = tf_utils.conv2d(conv5, 32*self.dis_c, k_h=3, k_w=3, d_h=1, d_w=1, name='conv6_conv1')
            conv6 = tf_utils.lrelu(conv6)

            # output layer: (N, 512, 512, 1024) -> (N, 512, 512, 1)
            output = tf_utils.conv2d(conv3, 1, k_h=1, k_w=1, d_h=1, d_w=1, name='conv_output')

            return tf.nn.sigmoid(output), output

    def train_dis(self, x_data, y_data):
        feed_dict = {self.X: x_data, self.Y: y_data}
        # run discriminator
        _, d_loss = self.sess.run([self.dis_optim, self.d_loss], feed_dict=feed_dict)

        return d_loss

    def train_gen(self, x_data, y_data):
        x_data0 = x_data[:,:,:,0]
        x_data0 = np.expand_dims(x_data0, axis=3)
        x_data1 = x_data[:,:,:,1]
        x_data1 = np.expand_dims(x_data1, axis=3)
        x_data2 = x_data[:,:,:,2]
        x_data2 = np.expand_dims(x_data2, axis=3)
        x_data3 = x_data[:,:,:,3]
        x_data3 = np.expand_dims(x_data3, axis=3)
        feed_dict = {self.X: x_data0, self.X1: x_data1,self.X2: x_data2,self.X3: x_data3,self.Y: y_data}
        # feed_dict = {self.X: x_data, self.Y: y_data}
        # run generator
        _, g_loss = self.sess.run([self.gen_optim, self.g_loss], feed_dict=feed_dict)

        return g_loss

    def measure_assign(self, auc_pr, auc_roc, dice_coeff, acc, sensitivity, specificity, score, iter_time):
        feed_dict = {self.auc_pr_placeholder: auc_pr,
                     self.auc_roc_placeholder: auc_roc,
                     self.dice_coeff_placeholder: dice_coeff,
                     self.acc_placeholder: acc,
                     self.sensitivity_placeholder: sensitivity,
                     self.specificity_placeholder: specificity,
                     self.score_placeholder: score}

        self.sess.run(self.measure_assign_op, feed_dict=feed_dict)

        summary = self.sess.run(self.measure_summary)
        self.writer.add_summary(summary, iter_time)

    def best_auc_sum_assign(self, auc_sum):
        self.sess.run(self.best_auc_sum_assign_op, feed_dict={self.best_auc_sum_placeholder: auc_sum})

    def sample_imgs(self, x_data):
        x_data0 = x_data[:,:,:,0]
        x_data0 = np.expand_dims(x_data0, axis=3)
        x_data1 = x_data[:,:,:,1]
        x_data1 = np.expand_dims(x_data1, axis=3)
        x_data2 = x_data[:,:,:,2]
        x_data2 = np.expand_dims(x_data2, axis=3)
        x_data3 = x_data[:,:,:,3]
        x_data3 = np.expand_dims(x_data3, axis=3)

        return self.sess.run(self.g_samples, feed_dict={self.X: x_data0, self.X1: x_data1,self.X2: x_data2,self.X3: x_data3})

    def measure_loss(self, gloss, dloss, iter_time):
        feed_dict = {self.gloss_save_placeholder: gloss,
                     self.dloss_save_placeholder: dloss}

        self.sess.run(self.measure_loss_op, feed_dict=feed_dict)
        summary = self.sess.run(self.measure_loss_summary)
        self.writer.add_summary(summary, iter_time)


#! -*- coding: utf-8 -*-
# 自定义层

import numpy as np
import tensorflow as tf
from bert4keras.backend import keras, K
from bert4keras.backend import search_layer
from bert4keras.backend import sequence_masking
from bert4keras.backend import pool1d
from bert4keras.backend import divisible_temporal_padding
from bert4keras.snippets import is_string
from keras import initializers, activations
from keras.layers import *

if keras.__version__[-2:] != 'tf' and keras.__version__ < '2.3':

    class Layer(keras.layers.Layer):
        """重新定义Layer，赋予“层中层”功能
        （仅keras 2.3以下版本需要）
        """
        def __setattr__(self, name, value):
            if isinstance(value, keras.layers.Layer):
                if not hasattr(self, '_layers'):
                    self._layers = []
                if value not in self._layers:
                    self._layers.append(value)
            super(Layer, self).__setattr__(name, value)

        @property
        def trainable_weights(self):
            trainable = getattr(self, 'trainable', True)
            if trainable:
                trainable_weights = super(Layer, self).trainable_weights[:]
                for l in getattr(self, '_layers', []):
                    trainable_weights += l.trainable_weights
                return trainable_weights
            else:
                return []

        @property
        def non_trainable_weights(self):
            trainable = getattr(self, 'trainable', True)
            non_trainable_weights = super(Layer, self).non_trainable_weights[:]
            for l in getattr(self, '_layers', []):
                if trainable:
                    non_trainable_weights += l.non_trainable_weights
                else:
                    non_trainable_weights += l.weights
            return non_trainable_weights

class ZeroMasking(Layer):
    """啥都不做，就是加上mask
    """
    def call(self, inputs):
        self._output_mask = K.cast(K.greater(inputs, 0), K.floatx())
        return inputs

    @property
    def output_mask(self):
        return self._output_mask

class MultiHeadAttention(Layer):
    """多头注意力机制
    """
    def __init__(self,
                 heads,
                 head_size,
                 key_size=None,
                 pool_size=None,
                 kernel_initializer='glorot_uniform',
                 max_relative_position=None,
                 **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.heads = heads
        self.head_size = head_size
        self.out_dim = heads * head_size
        self.key_size = key_size or head_size
        self.pool_size = pool_size or 1
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.max_relative_position = max_relative_position

    def build(self, input_shape):
        super(MultiHeadAttention, self).build(input_shape)
        self.q_dense = Dense(units=self.key_size * self.heads,
                             kernel_initializer=self.kernel_initializer)
        self.k_dense = Dense(units=self.key_size * self.heads,
                             kernel_initializer=self.kernel_initializer)
        self.v_dense = Dense(units=self.out_dim,
                             kernel_initializer=self.kernel_initializer)
        self.o_dense = Dense(units=self.out_dim,
                             kernel_initializer=self.kernel_initializer)

        if self.max_relative_position is not None:
            if self.head_size != self.key_size:
                raise ValueError('head_size must be equal to key_size ' +
                                 'while use relative position embeddings.')

            def initializer(shape, dtype=None):
                vocab_size, depth = shape
                embeddings = np.zeros(shape)
                for pos in range(vocab_size):
                    for i in range(depth // 2):
                        theta = pos / np.power(10000, 2. * i / depth)
                        embeddings[pos, 2 * i] = np.sin(theta)
                        embeddings[pos, 2 * i + 1] = np.cos(theta)
                return embeddings

            shape = (2 * self.max_relative_position + 1, self.head_size)
            self.relative_embeddings = self.add_weight(name='relative_embeddings',
                                                       shape=shape,
                                                       initializer=initializer,
                                                       trainable=False)

    def call(self, inputs, q_mask=None, v_mask=None, a_mask=None):
        """实现多头注意力
        q_mask: 对输入的query序列的mask。
                主要是将输出结果的padding部分置0。
        v_mask: 对输入的value序列的mask。
                主要是防止attention读取到padding信息。
        a_mask: 对attention矩阵的mask。
                不同的attention mask对应不同的应用。
        """
        q, k, v = inputs[:3]
        if a_mask:
            if len(inputs) == 3:
                a_mask = 'history_only'
            else:
                a_mask = inputs[3]
        if q_mask is not None:
            if not hasattr(self, 'q_mask_layer'):
                self.q_mask_layer = search_layer(q, q_mask)
            q_mask = self.q_mask_layer.output_mask
        if v_mask is not None:
            if not hasattr(self, 'v_mask_layer'):
                self.v_mask_layer = search_layer(v, v_mask)
            v_mask = self.v_mask_layer.output_mask
        # Pooling
        if self.pool_size > 1:
            is_self_attention = (q is k is v)
            q_in_len = K.shape(q)[1]
            q = sequence_masking(q, q_mask, 0)
            q = divisible_temporal_padding(q, self.pool_size)
            q = pool1d(q, self.pool_size, self.pool_size, pool_mode='avg')
            if is_self_attention:
                k = v = q
            else:
                k = sequence_masking(k, v_mask, 0)
                k = divisible_temporal_padding(k, self.pool_size)
                k = pool1d(k, self.pool_size, self.pool_size, pool_mode='avg')
                v = sequence_masking(v, v_mask, 0)
                v = divisible_temporal_padding(v, self.pool_size)
                v = pool1d(v, self.pool_size, self.pool_size, pool_mode='avg')
            if v_mask is not None:
                v_mask = v_mask[:, ::self.pool_size]
            if a_mask is not None and not is_string(a_mask):
                a_mask = a_mask[..., ::self.pool_size, ::self.pool_size]
        # 线性变换
        qw = self.q_dense(q)
        kw = self.k_dense(k)
        vw = self.v_dense(v)
        # 形状变换
        qw = K.reshape(qw, (-1, K.shape(q)[1], self.heads, self.key_size))
        kw = K.reshape(kw, (-1, K.shape(k)[1], self.heads, self.key_size))
        vw = K.reshape(vw, (-1, K.shape(v)[1], self.heads, self.head_size))
        # Attention
        a = tf.einsum('bjhd,bkhd->bhjk', qw, kw)
        # 相对位置编码
        if self.max_relative_position is not None:
            q_idxs = K.arange(0, K.shape(q)[1], dtype='int32')
            q_idxs = K.expand_dims(q_idxs, 1)
            v_idxs = K.arange(0, K.shape(v)[1], dtype='int32')
            v_idxs = K.expand_dims(v_idxs, 0)
            pos_ids = v_idxs - q_idxs
            pos_ids = K.clip(pos_ids, -self.max_relative_position,
                             self.max_relative_position)
            pos_ids = pos_ids + self.max_relative_position
            pos_embeddings = K.gather(self.relative_embeddings, pos_ids)
            a = a + tf.einsum('bjhd,jkd->bhjk', qw, pos_embeddings)
        # Attention（续）
        a = a / self.key_size**0.5
        a = sequence_masking(a, v_mask, 1, -1)
        if a_mask is not None:
            if is_string(a_mask):
                ones = K.ones_like(a[:1, :1])
                a_mask = (ones - tf.linalg.band_part(ones, -1, 0)) * 1e12
                a = a - a_mask
            else:
                a = a - (1 - a_mask) * 1e12
        a = K.softmax(a)
        # 完成输出
        o = tf.einsum('bhjk,bkhd->bjhd', a, vw)
        if self.max_relative_position is not None:
            o = o + tf.einsum('bhjk,jkd->bjhd', a, pos_embeddings)
        o = K.reshape(o, (-1, K.shape(o)[1], self.out_dim))
        o = self.o_dense(o)
        # 恢复长度
        if self.pool_size > 1:
            o = K.repeat_elements(o, self.pool_size, 1)[:, :q_in_len]
        # 返回结果
        o = sequence_masking(o, q_mask, 0)
        return o

    def compute_output_shape(self, input_shape):
        return (input_shape[0][0], input_shape[0][1], self.out_dim)

    def get_config(self):
        config = {
            'heads': self.heads,
            'head_size': self.head_size,
            'key_size': self.key_size,
            'pool_size': self.pool_size,
            'kernel_initializer': initializers.serialize(self.kernel_initializer),
            'max_relative_position': self.max_relative_position,
        }
        base_config = super(MultiHeadAttention, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class LayerNormalization(Layer):
    """(Conditional) Layer Normalization
    hidden_*系列参数仅为有条件输入时(conditional=True)使用
    """
    def __init__(self,
                 conditional=False,
                 hidden_units=None,
                 hidden_activation='linear',
                 hidden_initializer='glorot_uniform',
                 **kwargs):
        super(LayerNormalization, self).__init__(**kwargs)
        self.conditional = conditional
        self.hidden_units = hidden_units
        self.hidden_activation = activations.get(hidden_activation)
        self.hidden_initializer = initializers.get(hidden_initializer)
        self.epsilon = K.epsilon() * K.epsilon()

    def build(self, input_shape):
        super(LayerNormalization, self).build(input_shape)

        if self.conditional:
            shape = (input_shape[0][-1], )
        else:
            shape = (input_shape[-1], )

        self.gamma = self.add_weight(shape=shape,
                                     initializer='ones',
                                     name='gamma')
        self.beta = self.add_weight(shape=shape,
                                    initializer='zeros',
                                    name='beta')

        if self.conditional:

            if self.hidden_units is not None:
                self.hidden_dense = Dense(
                    units=self.hidden_units,
                    activation=self.hidden_activation,
                    use_bias=False,
                    kernel_initializer=self.hidden_initializer)

            self.beta_dense = Dense(units=shape[0],
                                    use_bias=False,
                                    kernel_initializer='zeros')
            self.gamma_dense = Dense(units=shape[0],
                                     use_bias=False,
                                     kernel_initializer='zeros')

    def call(self, inputs):
        """如果是条件Layer Norm，则默认以list为输入，第二个是condition
        """
        if self.conditional:
            inputs, cond = inputs
            if self.hidden_units is not None:
                cond = self.hidden_dense(cond)
            for _ in range(K.ndim(inputs) - K.ndim(cond)):
                cond = K.expand_dims(cond, 1)
            beta = self.beta_dense(cond)
            gamma = self.gamma_dense(cond)
            beta, gamma = self.beta + beta, self.gamma + gamma
        else:
            beta, gamma = self.beta, self.gamma

        mean = K.mean(inputs, axis=-1, keepdims=True)
        variance = K.mean(K.square(inputs - mean), axis=-1, keepdims=True)
        std = K.sqrt(variance + self.epsilon)
        outputs = (inputs - mean) / std
        outputs = outputs * gamma + beta
        return outputs

    def get_config(self):
        config = {
            'conditional': self.conditional,
            'hidden_units': self.hidden_units,
            'hidden_activation': activations.serialize(self.hidden_activation),
            'hidden_initializer': initializers.serialize(self.hidden_initializer),
        }
        base_config = super(LayerNormalization, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class PositionEmbedding(Layer):
    """定义位置Embedding，这里的Embedding是可训练的。
    """
    def __init__(self,
                 input_dim,
                 output_dim,
                 merge_mode='add',
                 embeddings_initializer='zeros',
                 **kwargs):
        super(PositionEmbedding, self).__init__(**kwargs)
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.merge_mode = merge_mode
        self.embeddings_initializer = initializers.get(embeddings_initializer)

    def build(self, input_shape):
        super(PositionEmbedding, self).build(input_shape)
        self.embeddings = self.add_weight(
            name='embeddings',
            shape=(self.input_dim, self.output_dim),
            initializer=self.embeddings_initializer,
        )

    def call(self, inputs):
        input_shape = K.shape(inputs)
        batch_size, seq_len = input_shape[0], input_shape[1]
        pos_embeddings = self.embeddings[:seq_len]
        pos_embeddings = K.expand_dims(pos_embeddings, 0)

        if self.merge_mode == 'add':
            return inputs + pos_embeddings
        else:
            pos_embeddings = K.tile(pos_embeddings, [batch_size, 1, 1])
            return K.concatenate([inputs, pos_embeddings])

    def compute_output_shape(self, input_shape):
        if self.merge_mode == 'add':
            return input_shape
        else:
            return input_shape[:2] + (input_shape[2] + self.output_dim, )

    def get_config(self):
        config = {
            'input_dim': self.input_dim,
            'output_dim': self.output_dim,
            'merge_mode': self.merge_mode,
            'embeddings_initializer': initializers.serialize(self.embeddings_initializer),
        }
        base_config = super(PositionEmbedding, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class GroupDense(Layer):
    """分组全连接
    输入输出跟普通Dense一样，但参数更少，速度更快。
    """
    def __init__(self,
                 units,
                 groups=2,
                 activation='linear',
                 kernel_initializer='glorot_uniform',
                 **kwargs):
        super(GroupDense, self).__init__(**kwargs)
        self.units = units
        self.groups = groups
        self.activation = activations.get(activation)
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        super(GroupDense, self).build(input_shape)
        input_dim = input_shape[-1]
        if not isinstance(input_dim, int):
            input_dim = input_dim.value
        assert input_dim % self.groups == 0
        assert self.units % self.groups == 0
        self.kernel = self.add_weight(name='kernel',
                                      shape=(input_dim // self.groups,
                                             self.units // self.groups,
                                             self.groups),
                                      initializer=self.kernel_initializer)
        self.bias = self.add_weight(name='bias',
                                    shape=(self.units, ),
                                    initializer='zeros')

    def call(self, inputs):
        ndim, shape = K.ndim(inputs), K.shape(inputs)
        shape = [shape[i] for i in range(ndim)]
        inputs = K.reshape(inputs, shape[:-1] + [shape[-1] // self.groups, self.groups])
        outputs = tf.einsum('...ig,ijg->...gj', inputs, self.kernel)
        outputs = K.reshape(outputs, shape[:-1] + [self.units])
        outputs = outputs + self.bias
        outputs = self.activation(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        return input_shape[:-1] + (self.units, )

    def get_config(self):
        config = {
            'units': self.units,
            'groups': self.groups,
            'activation': activations.serialize(self.activation),
            'kernel_initializer': initializers.serialize(self.kernel_initializer),
        }
        base_config = super(GroupDense, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class FeedForward(Layer):
    """FeedForward层，其实就是两个Dense层的叠加
    """
    def __init__(self,
                 units,
                 groups=1,
                 activation='relu',
                 kernel_initializer='glorot_uniform',
                 pool_size=None,
                 **kwargs):
        super(FeedForward, self).__init__(**kwargs)
        self.units = units
        self.groups = groups
        self.activation = activations.get(activation)
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.pool_size = pool_size or 1

    def build(self, input_shape):
        super(FeedForward, self).build(input_shape)
        output_dim = input_shape[-1]
        if not isinstance(output_dim, int):
            output_dim = output_dim.value

        if self.groups is None or self.groups == 1:
            self.dense_1 = Dense(units=self.units,
                                 activation=self.activation,
                                 kernel_initializer=self.kernel_initializer)
            self.dense_2 = Dense(units=output_dim,
                                 kernel_initializer=self.kernel_initializer)
        else:
            self.dense_1 = GroupDense(units=self.units,
                                      groups=self.groups,
                                      activation=self.activation,
                                      kernel_initializer=self.kernel_initializer)
            self.dense_2 = GroupDense(units=output_dim,
                                      groups=self.groups,
                                      kernel_initializer=self.kernel_initializer)

    def call(self, inputs, mask=None):
        x = inputs
        # Pooling
        if self.pool_size > 1:
            if mask is not None:
                if not hasattr(self, 'mask_layer'):
                    self.mask_layer = search_layer(x, mask)
                mask = self.mask_layer.output_mask
            x_in_len = K.shape(x)[1]
            x = sequence_masking(x, mask, 0)
            x = divisible_temporal_padding(x, self.pool_size)
            x = pool1d(x, self.pool_size, self.pool_size, pool_mode='avg')
        # 执行FFN
        x = self.dense_1(x)
        x = self.dense_2(x)
        # 恢复长度
        if self.pool_size > 1:
            x = K.repeat_elements(x, self.pool_size, 1)[:, :x_in_len]
        # 返回结果
        return x

    def get_config(self):
        config = {
            'units': self.units,
            'activation': activations.serialize(self.activation),
            'kernel_initializer': initializers.serialize(self.kernel_initializer),
            'pool_size': self.pool_size,
        }
        base_config = super(FeedForward, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class EmbeddingDense(Layer):
    """运算跟Dense一致，但kernel用Embedding层的embeddings矩阵。
    根据Embedding层的名字来搜索定位Embedding层。
    """
    def __init__(self, embedding_name, activation='softmax', **kwargs):
        super(EmbeddingDense, self).__init__(**kwargs)
        self.embedding_name = embedding_name
        self.activation = activations.get(activation)

    def call(self, inputs):
        if not hasattr(self, 'kernel'):
            embedding_layer = search_layer(inputs, self.embedding_name)
            if embedding_layer is None:
                raise Exception('Embedding layer not found')

            self.kernel = K.transpose(embedding_layer.embeddings)
            self.units = K.int_shape(self.kernel)[1]
            self.bias = self.add_weight(name='bias',
                                        shape=(self.units, ),
                                        initializer='zeros')

        outputs = K.dot(inputs, self.kernel)
        outputs = K.bias_add(outputs, self.bias)
        outputs = self.activation(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        return input_shape[:-1] + (self.units, )

    def get_config(self):
        config = {
            'embedding_name': self.embedding_name,
            'activation': activations.serialize(self.activation),
        }
        base_config = super(EmbeddingDense, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class ConditionalRandomField(Layer):
    """纯Keras实现CRF层
    CRF层本质上是一个带训练参数的loss计算层。
    """
    def __init__(self, lr_multiplier=1, **kwargs):
        super(ConditionalRandomField, self).__init__(**kwargs)
        self.lr_multiplier = lr_multiplier  # 当前层学习率的放大倍数

    def build(self, input_shape):
        output_dim = input_shape[-1]
        if not isinstance(output_dim, int):
            output_dim = output_dim.value
        self.trans = self.add_weight(name='trans',
                                     shape=(output_dim, output_dim),
                                     initializer='glorot_uniform',
                                     trainable=True)
        if self.lr_multiplier != 1:
            K.set_value(self.trans, K.eval(self.trans) / self.lr_multiplier)
            self.trans = self.lr_multiplier * self.trans

    def target_score(self, y_true, y_pred, mask=None):
        """计算目标路径的相对概率（还没有归一化）
        要点：逐标签得分，加上转移概率得分。
        """
        y_true = sequence_masking(y_true, mask, 0)
        point_score = tf.einsum('bni,bni->b', y_true, y_pred)  # 逐标签得分
        trans_score = tf.einsum('bni,ij,bnj->b', y_true[:, :-1], self.trans,
                                y_true[:, 1:])  # 标签转移得分
        return point_score + trans_score

    def log_norm_step(self, inputs, states):
        """递归计算归一化因子
        要点：1、递归计算；2、用logsumexp避免溢出。
        """
        inputs, mask = inputs[:, :-1], inputs[:, -1:]
        states = K.expand_dims(states[0], 2)  # (batch_size, output_dim, 1)
        trans = K.expand_dims(self.trans, 0)  # (1, output_dim, output_dim)
        outputs = tf.reduce_logsumexp(states + trans, 1)  # (batch_size, output_dim)
        outputs = outputs + inputs
        outputs = mask * outputs + (1 - mask) * states[:, :, 0]
        return outputs, [outputs]

    def call(self, inputs, mask=None):
        """CRF本身不改变输出，它只是一个loss
        """
        if mask is not None:
            if not hasattr(self, 'mask_layer'):
                self.mask_layer = search_layer(inputs, mask)

        return inputs

    @property
    def output_mask(self):
        if hasattr(self, 'mask_layer'):
            return self.mask_layer.output_mask

    def dense_loss(self, y_true, y_pred):
        """y_true需要是one hot形式
        """
        mask = self.output_mask
        # 计算目标分数
        target_score = self.target_score(y_true, y_pred, mask)
        # 递归计算log Z
        init_states = [y_pred[:, 0]]
        if mask is None:
            mask = K.ones_like(y_pred[:, :, :1])
        else:
            mask = K.expand_dims(mask, 2)
        y_pred = K.concatenate([y_pred, mask])
        log_norm, _, _ = K.rnn(self.log_norm_step,
                               y_pred[:, 1:],
                               init_states)  # 最后一步的log Z向量
        log_norm = tf.reduce_logsumexp(log_norm, 1)  # logsumexp得标量
        # 计算损失 -log p
        return log_norm - target_score

    def sparse_loss(self, y_true, y_pred):
        """y_true需要是整数形式（非one hot）
        """
        # y_true需要重新明确一下dtype和shape
        y_true = K.cast(y_true, 'int32')
        y_true = K.reshape(y_true, [K.shape(y_true)[0], -1])
        # 转为one hot
        y_true = K.one_hot(y_true, K.shape(self.trans)[0])
        return self.dense_loss(y_true, y_pred)

    def dense_accuracy(self, y_true, y_pred):
        """训练过程中显示逐帧准确率的函数，排除了mask的影响
        此处y_true需要是one hot形式
        """
        y_true = K.argmax(y_true, 2)
        return self.sparse_accuracy(y_true, y_pred)

    def sparse_accuracy(self, y_true, y_pred):
        """训练过程中显示逐帧准确率的函数，排除了mask的影响
        此处y_true需要是整数形式（非one hot）
        """
        mask = self.output_mask
        # y_true需要重新明确一下dtype和shape
        y_true = K.cast(y_true, 'int32')
        y_true = K.reshape(y_true, [K.shape(y_true)[0], -1])
        # 逐标签取最大来粗略评测训练效果
        y_pred = K.cast(K.argmax(y_pred, 2), 'int32')
        isequal = K.cast(K.equal(y_true, y_pred), K.floatx())
        if mask is None:
            return K.mean(isequal)
        else:
            return K.sum(isequal * mask) / K.sum(mask)

    def get_config(self):
        config = {
            'lr_multiplier': self.lr_multiplier,
        }
        base_config = super(ConditionalRandomField, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

custom_objects = {
    'ZeroMasking': ZeroMasking,
    'MultiHeadAttention': MultiHeadAttention,
    'LayerNormalization': LayerNormalization,
    'PositionEmbedding': PositionEmbedding,
    'GroupDense': GroupDense,
    'FeedForward': FeedForward,
    'EmbeddingDense': EmbeddingDense,
    'ConditionalRandomField': ConditionalRandomField,
}

keras.utils.get_custom_objects().update(custom_objects)


import numpy as np
import torch
import torch.nn.functional as F

from classifier_control.classifier.utils.spatial_softmax import SpatialSoftmax
from classifier_control.classifier.utils.layers import Linear
from classifier_control.classifier.utils.subnetworks import ConvEncoder
from classifier_control.classifier.utils.general_utils import AttrDict
from classifier_control.classifier.models.base_model import BaseModel
from classifier_control.classifier.models.utils.utils import select_indices

class TempdistRegressor(BaseModel):
    def __init__(self, overrideparams, logger=None):
        super().__init__(logger)
        self._hp = self._default_hparams()
        self.overrideparams = overrideparams
        self.override_defaults(overrideparams)  # override defaults with config file
        self.postprocess_params()

        assert self._hp.batch_size != -1   # make sure that batch size was overridden

        self.tdist_classifiers = []
        self.build_network()
        self.proxy_ctrl_counter = 0

    def _default_hparams(self):
        default_dict = AttrDict({
            'tmax_label':10,  # the highest label for temporal distance, values are clamped after that
            'use_skips':False, #todo try resnet architecture!
            'ngf': 8,
            'nz_enc': 64,
            'classifier_restore_path':None, # not really needed here.
        })

        # add new params to parent params
        parent_params = super()._default_hparams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def build_network(self, build_encoder=True):
        self.encoder = ConvEncoder(self._hp)
        out_size = self.encoder.get_output_size()
        if self._hp.spatial_softmax:
            self.spatial_softmax = SpatialSoftmax(out_size[1], out_size[2], out_size[0])  # height, width, channel
            self.linear = Linear(in_dim=out_size[0]*2, out_dim=1, builder=self._hp.builder)
        else:
            self.linear = Linear(in_dim=256, out_dim=1, builder=self._hp.builder)
            self.fc_layers = []
            self.fc_layers.append(Linear(in_dim=out_size[0]*out_size[1]*out_size[2], out_dim=256, builder=self._hp.builder))
            for i in range(3):
                self.fc_layers.append(Linear(in_dim=256, out_dim=256, builder=self._hp.builder))

    def sample_image_pair(self, images):
        tlen = images.shape[1]

        # get positives:
        t0 = np.random.randint(0, tlen, self._hp.batch_size)
        t1 = np.array([np.random.randint(t0[b], tlen, 1) for b in range(images.shape[0])]).squeeze()

        t0, t1 = torch.from_numpy(t0), torch.from_numpy(t1)

        im_t0 = select_indices(images, t0)
        im_t1 = select_indices(images, t1)

        self.labels = torch.clamp_max(t1 - t0, self._hp.tmax_label-1).float()

        img_pair_stack = torch.stack([im_t0, im_t1], dim=1)
        return img_pair_stack

    def forward(self, inputs):
        """
        forward pass at training time
        :param
            images shape = batch x time x channel x height x width
        :return: model_output
        """
        image_pairs = self.sample_image_pair(inputs.demo_seq_images)
        self.img_pair = image_pairs
        model_output = self.make_prediction(image_pairs)
        return model_output

    def make_prediction(self, image_pairs_stacked):
        im_t0, im_t1 = image_pairs_stacked[:,0], image_pairs_stacked[:,1]
        embeddings = self.encoder(torch.cat([im_t0, im_t1], dim=1))
        if self._hp.spatial_softmax:
            embeddings = self.spatial_softmax(embeddings)
        else:
            embeddings = torch.flatten(embeddings, start_dim=1)
            for fc_layer in self.fc_layers:
                embeddings = F.relu(fc_layer(embeddings))
        self.tdist_estimates = self.linear(embeddings)
        model_output = AttrDict(tdist_estimates=self.tdist_estimates, img_pair=image_pairs_stacked)
        return model_output

    def _log_outputs(self, model_output, inputs, losses, step, log_images, phase):
        if log_images:
            self._logger.log_pair_predictions(self.img_pair, self.tdist_estimates, self.labels,'tdist_regression', step, phase)

    def loss(self, model_output):
        losses = AttrDict()
        setattr(losses, 'mse', torch.nn.MSELoss()(model_output.tdist_estimates.squeeze(), self.labels.to(self._hp.device)))

        # compute total loss
        losses.total_loss = torch.stack(list(losses.values())).sum()
        return losses

    def get_device(self):
        return self._hp.device

class TempdistRegressorTestTime(TempdistRegressor):
    def __init__(self, overrideparams, logger=None, restore_from_disk=True):
        super(TempdistRegressorTestTime, self).__init__(overrideparams, logger)
        if self._hp.classifier_restore_path is not None:
            checkpoint = torch.load(self._hp.classifier_restore_path, map_location=self._hp.device)
            self.load_state_dict(checkpoint['state_dict'])
        else:
            print('#########################')
            print("Warning Classifier weights not restored during init!!")
            print('#########################')

    def _default_hparams(self):
        parent_params = super()._default_hparams()
        parent_params.add_hparam('classifier_restore_path', None)
        return parent_params

    def forward(self, inputs):
        """
        forward pass at training time
        :param
            images shape = batch x time x channel x height x width
        :return: model_output
        """
        image_pairs = torch.stack([inputs['current_img'], inputs['goal_img']], dim=1)
        expected_distance = self.make_prediction(image_pairs).tdist_estimates.data.cpu().numpy().squeeze()
        return expected_distance

    def visualize_test_time(self, content_dict, visualize_indices, verbose_folder):
        pass

def ptrch2uint8(img):
    return ((img + 1)/2*255.).astype(np.uint8)


import numba
import numpy as np
import torch

from mmdet3d.ops.iou3d.iou3d_utils import nms_gpu, nms_normal_gpu, nms_weighted_gpu

def box3d_multiclass_nms(mlvl_bboxes,
                         mlvl_bboxes_for_nms,
                         mlvl_scores,
                         score_thr,
                         max_num,
                         cfg,
                         mlvl_dir_scores=None,
                         mlvl_attr_scores=None,
                         mlvl_bboxes2d=None):
    """Multi-class nms for 3D boxes.

    Args:
        mlvl_bboxes (torch.Tensor): Multi-level boxes with shape (N, M).
            M is the dimensions of boxes.
        mlvl_bboxes_for_nms (torch.Tensor): Multi-level boxes with shape
            (N, 5) ([x1, y1, x2, y2, ry]). N is the number of boxes.
        mlvl_scores (torch.Tensor): Multi-level boxes with shape
            (N, C + 1). N is the number of boxes. C is the number of classes.
        score_thr (float): Score thredhold to filter boxes with low
            confidence.
        max_num (int): Maximum number of boxes will be kept.
        cfg (dict): Configuration dict of NMS.
        mlvl_dir_scores (torch.Tensor, optional): Multi-level scores
            of direction classifier. Defaults to None.
        mlvl_attr_scores (torch.Tensor, optional): Multi-level scores
            of attribute classifier. Defaults to None.
        mlvl_bboxes2d (torch.Tensor, optional): Multi-level 2D bounding
            boxes. Defaults to None.

    Returns:
        tuple[torch.Tensor]: Return results after nms, including 3D \
            bounding boxes, scores, labels, direction scores, attribute \
            scores (optional) and 2D bounding boxes (optional).
    """
    # do multi class nms
    # the fg class id range: [0, num_classes-1]
    num_classes = mlvl_scores.shape[1] - 1
    bboxes = []
    scores = []
    labels = []
    dir_scores = []
    attr_scores = []
    bboxes2d = []
    for i in range(0, num_classes):
        # get bboxes and scores of this class
        cls_inds = mlvl_scores[:, i] > score_thr
        if not cls_inds.any():
            continue

        _scores = mlvl_scores[cls_inds, i]
        _bboxes_for_nms = mlvl_bboxes_for_nms[cls_inds, :]

        if cfg.use_rotate_nms:
            nms_func = nms_gpu
        else:
            nms_func = nms_normal_gpu

        selected = nms_func(_bboxes_for_nms, _scores, cfg.nms_thr)
        _mlvl_bboxes = mlvl_bboxes[cls_inds, :]
        bboxes.append(_mlvl_bboxes[selected])
        scores.append(_scores[selected])

        # _mlvl_bboxes = mlvl_bboxes[cls_inds, :]
        # selected, keep_boxes, keep_scores = nms_weighted_gpu(_bboxes_for_nms, _scores, cfg.nms_thr, _mlvl_bboxes)
        # bboxes.append(keep_boxes[selected])
        # scores.append(keep_scores[selected])
        # import pdb; pdb.set_trace()

        cls_label = mlvl_bboxes.new_full((len(selected), ),
                                         i,
                                         dtype=torch.long)
        labels.append(cls_label)

        if mlvl_dir_scores is not None:
            _mlvl_dir_scores = mlvl_dir_scores[cls_inds]
            dir_scores.append(_mlvl_dir_scores[selected])
        if mlvl_attr_scores is not None:
            _mlvl_attr_scores = mlvl_attr_scores[cls_inds]
            attr_scores.append(_mlvl_attr_scores[selected])
        if mlvl_bboxes2d is not None:
            _mlvl_bboxes2d = mlvl_bboxes2d[cls_inds]
            bboxes2d.append(_mlvl_bboxes2d[selected])

    if bboxes:
        bboxes = torch.cat(bboxes, dim=0)
        scores = torch.cat(scores, dim=0)
        labels = torch.cat(labels, dim=0)
        if mlvl_dir_scores is not None:
            dir_scores = torch.cat(dir_scores, dim=0)
        if mlvl_attr_scores is not None:
            attr_scores = torch.cat(attr_scores, dim=0)
        if mlvl_bboxes2d is not None:
            bboxes2d = torch.cat(bboxes2d, dim=0)
        if bboxes.shape[0] > max_num:
            _, inds = scores.sort(descending=True)
            inds = inds[:max_num]
            bboxes = bboxes[inds, :]
            labels = labels[inds]
            scores = scores[inds]
            if mlvl_dir_scores is not None:
                dir_scores = dir_scores[inds]
            if mlvl_attr_scores is not None:
                attr_scores = attr_scores[inds]
            if mlvl_bboxes2d is not None:
                bboxes2d = bboxes2d[inds]
    else:
        bboxes = mlvl_scores.new_zeros((0, mlvl_bboxes.size(-1)))
        scores = mlvl_scores.new_zeros((0, ))
        labels = mlvl_scores.new_zeros((0, ), dtype=torch.long)
        if mlvl_dir_scores is not None:
            dir_scores = mlvl_scores.new_zeros((0, ))
        if mlvl_attr_scores is not None:
            attr_scores = mlvl_scores.new_zeros((0, ))
        if mlvl_bboxes2d is not None:
            bboxes2d = mlvl_scores.new_zeros((0, 4))

    results = (bboxes, scores, labels)

    if mlvl_dir_scores is not None:
        results = results + (dir_scores, )
    if mlvl_attr_scores is not None:
        results = results + (attr_scores, )
    if mlvl_bboxes2d is not None:
        results = results + (bboxes2d, )

    return results

def aligned_3d_nms(boxes, scores, classes, thresh):
    """3d nms for aligned boxes.

    Args:
        boxes (torch.Tensor): Aligned box with shape [n, 6].
        scores (torch.Tensor): Scores of each box.
        classes (torch.Tensor): Class of each box.
        thresh (float): Iou threshold for nms.

    Returns:
        torch.Tensor: Indices of selected boxes.
    """
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    z1 = boxes[:, 2]
    x2 = boxes[:, 3]
    y2 = boxes[:, 4]
    z2 = boxes[:, 5]
    area = (x2 - x1) * (y2 - y1) * (z2 - z1)
    zero = boxes.new_zeros(1, )

    score_sorted = torch.argsort(scores)
    pick = []
    while (score_sorted.shape[0] != 0):
        last = score_sorted.shape[0]
        i = score_sorted[-1]
        pick.append(i)

        xx1 = torch.max(x1[i], x1[score_sorted[:last - 1]])
        yy1 = torch.max(y1[i], y1[score_sorted[:last - 1]])
        zz1 = torch.max(z1[i], z1[score_sorted[:last - 1]])
        xx2 = torch.min(x2[i], x2[score_sorted[:last - 1]])
        yy2 = torch.min(y2[i], y2[score_sorted[:last - 1]])
        zz2 = torch.min(z2[i], z2[score_sorted[:last - 1]])
        classes1 = classes[i]
        classes2 = classes[score_sorted[:last - 1]]
        inter_l = torch.max(zero, xx2 - xx1)
        inter_w = torch.max(zero, yy2 - yy1)
        inter_h = torch.max(zero, zz2 - zz1)

        inter = inter_l * inter_w * inter_h
        iou = inter / (area[i] + area[score_sorted[:last - 1]] - inter)
        iou = iou * (classes1 == classes2).float()
        score_sorted = score_sorted[torch.nonzero(
            iou <= thresh, as_tuple=False).flatten()]

    indices = boxes.new_tensor(pick, dtype=torch.long)
    return indices

@numba.jit(nopython=True)
def circle_nms(dets, thresh, post_max_size=83):
    """Circular NMS.

    An object is only counted as positive if no other center
    with a higher confidence exists within a radius r using a
    bird-eye view distance metric.

    Args:
        dets (torch.Tensor): Detection results with the shape of [N, 3].
        thresh (float): Value of threshold.
        post_max_size (int): Max number of prediction to be kept. Defaults
            to 83

    Returns:
        torch.Tensor: Indexes of the detections to be kept.
    """
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    scores = dets[:, 2]
    order = scores.argsort()[::-1].astype(np.int32)  # highest->lowest
    ndets = dets.shape[0]
    suppressed = np.zeros((ndets), dtype=np.int32)
    keep = []
    for _i in range(ndets):
        i = order[_i]  # start with highest score box
        if suppressed[
                i] == 1:  # if any box have enough iou with this, remove it
            continue
        keep.append(i)
        for _j in range(_i + 1, ndets):
            j = order[_j]
            if suppressed[j] == 1:
                continue
            # calculate center distance between i and j box
            dist = (x1[i] - x1[j])**2 + (y1[i] - y1[j])**2

            # ovr = inter / areas[j]
            if dist <= thresh:
                suppressed[j] = 1
    return keep[:post_max_size]


'''
Aligns all RIFDock outputs to the input PDB, then clusters them based on
structure similarity. Representatives of each cluster are picked based
on their RIFDock score.

Usage:
    cluster_sge.py <rif_workspace> [options]

Options:
    --task=NUM, -t  Provide a task number. Each task is a different
    scaffold.
'''
import gzip
import docopt
import glob
import sys, os
import numpy as np
from shutil import copyfile
from helix import workspace as ws

class Design(object):
    def __init__(self, path):
        self.path = path
        self.backbone_coords = []
        self.score = None
        self.get_score()

    def read_coords(self, chain=None):
        with gzip.open(self.path, mode='rt') as f:
            lines = f.readlines()

        coords = []
        atoms = 'N', 'CA', 'C'

        for line in lines:
            if not line.startswith('ATOM'): continue

            atom_name = line[13:16].strip().upper()
            residue_id = int(line[22:26])
            chain_id = line[21]
            atom_coord = np.array((
                float(line[30:38]),
                float(line[38:46]),
                float(line[46:54])))

            if atom_name in atoms:
                if chain:
                    if chain_id == chain:
                        coords.append(atom_coord)
                else:
                    coords.append(atom_coord)

        self.backbone_coords = np.array(coords)

    def get_score(self):
        # if len(glob.glob(os.path.dirname(self.path) + '/all.d*')) > 1:
            # print('WARNING: Multiple docking scorefiles detected.')
            # print('This may result in the wrong score being used.')
        scorefiles = sorted(glob.glob(os.path.join(
            os.path.dirname(self.path), 'all.dok*'
            )))
        # scorefile = os.path.join(os.path.dirname(self.path), 'all.dok')
        scorefile = scorefiles[-1]
        with open(scorefile, 'r') as f:
            for line in f:
                score = line[76:82]
                path = line[141:-1]

                if os.path.basename(path) == os.path.basename(self.path):
                    self.score = float(score)
                    return

class Cluster(object):
    def __init__(self, cluster):
        self.cluster = cluster
        self.designs = []
        return

    @property
    def rep(self):
        return self.find_rep()

    def add(self, design):
        self.designs.append(design)

    def find_rep(self):
        lowest_score = 9999
        for design in self.designs:
            if design.score < lowest_score:
                lowest_score = design.score
                rep = design

        return rep

class StructureCluster(object):
    def __init__(self, workspace, basename=None, threshold=None):
        self.threshold=threshold
        self.designs = []
        # Path sould be to parent rifdock directory.
        if basename:
            search = workspace.focus_dir + '/*/docked_full/{}*.pdb.gz'.format(basename)
        else:
            search = workspace.focus_dir + '/*/docked_full/*.pdb.gz'
        for path in glob.glob(
                search
                ):
            design = Design(path)
            design.read_coords(chain='A')
            self.designs.append(design)

        self.clusters = {}

    def cluster_coords(self, verbose=False):
        import scipy.spatial.distance as sp_dist
        import scipy.cluster.hierarchy as sp_clust
        from itertools import combinations

        num_designs = len(self.designs)
        if num_designs < 2: return

        # Calculate pariwise distance matrix

        dist_matrix = np.zeros((num_designs, num_designs))
        design_combos = combinations(enumerate(self.designs), 2)

        for (i, design_i), (j, design_j) in design_combos:
            dist_matrix[i,j] = self.calculate_rmsd(design_i,
                    design_j)
            dist_matrix[j,i] = dist_matrix[i,j]

        print(dist_matrix)

        # Cluster design such that no two designs in a cluster are
        # further apart than the threshold, in Angstroms.

        dist_vector = sp_dist.squareform(dist_matrix)
        mean_dist = np.mean(dist_vector)
        hierarchy = sp_clust.complete(dist_vector)
        print(self.threshold)
        print(mean_dist)
        print(self.threshold or mean_dist)
        clusters = sp_clust.fcluster(
                hierarchy, self.threshold or mean_dist, criterion='distance')
        print(clusters)
        # clusters = sp_clust.fcluster(
                # hierarchy, 20,
                # criterion='maxclust')

        for cluster, design in zip(clusters, self.designs):
            design.structure_cluster = cluster
            if cluster not in self.clusters:
                clst = Cluster(cluster)
                clst.add(design)
                self.clusters[cluster] = clst
            else:
                self.clusters[cluster].add(design)

        # Print some debugging information, if requested.

        if verbose == True:
            cluster_map = {}

            for cluster, design in zip(clusters, self.designs):
                cluster_map.setdefault(cluster, []).append(design)

            for cluster in sorted(set(clusters)):

                # Print out the designs contained in this cluster.

                print("Cluster {}:".format(cluster))
                for design in cluster_map[cluster]:
                    print(" ", design.path)
                print()

                # Print out pairwise distances for every cluster member.

                X = cluster_map[cluster]
                N = len(X)
                D = np.zeros((N, N))

                for i, design_i in enumerate(X):
                    for j, design_j in enumerate(X):
                        D[i,j] = self.calculate_rmsd(design_i, design_j)

                print(sp_dist.squareform(D))
                print()

                # Offer to display the cluster in pymol.

                command = ['pymol', '-qx', '-d', 'as cartoon']
                for design in cluster_map[cluster]:
                    command.append(design.path)

                # if raw_input("  View in pymol? [y/N] ") == 'y':
                    # import subprocess
                    # subprocess.check_output(command)

                print()

    def calculate_rmsd(self, design_1, design_2):
        # assert len(design_1.backbone_coords) and len(design_2.backbone_coords)
        difference = design_1.backbone_coords - design_2.backbone_coords
        num_atoms = design_1.backbone_coords.shape[0]

        return np.sqrt(np.sum(difference**2) / num_atoms)

if __name__=='__main__':
    # folders = sorted(glob.glob(sys.argv[1] + '/*_output'))
    args = docopt.docopt(__doc__)
    workspace = ws.workspace_from_dir(args['<rif_workspace>'])
    assert(type(workspace).__name__ == 'RIFWorkspace')

    if 'SGE_TASK_ID' in os.environ:
        task = int(os.environ['SGE_TASK_ID']) - 1
    elif args['--task']:
        task = int(args['--task']) - 1
    else:
        task = 0
    # targets = workspace.targets
    helices = workspace.scaffolds
    helix = helices[task]
    # folders = workspace.patches
    # workspace = ws.RIFWorkspace(workspace.root_dir, targets[task])

    # for helixlength in [3,4,6,8]:
    # for helix in workspace.helices:
    basename = workspace.basename(helix)
    clust = StructureCluster(workspace, basename=basename, threshold=10)
    print('Clustering...')
    clust.cluster_coords(verbose=True)
    clusters = clust.clusters

    outpath = os.path.join(workspace.focus_dir, 'cluster_representatives',
            basename)
    # if not os.path.exists(outpath):
        # print('PATH NO EXIST')
        # os.mkdir(outpath)
    os.makedirs(outpath, exist_ok=True)
    scores = open(os.path.join(
        outpath,
        '{}.scores'.format(basename)), 'w')

    overwrite = True
    # overwrite = False

    for clst in clusters:
        outfile = '{}_clst{}_rep.pdb.gz'.format(basename,
                clst)
        out = os.path.join(outpath, outfile)
        if overwrite:
            if os.path.exists(out):
                os.remove(out)
        print('CLUSTER {}'.format(clst))
        print('REP:')
        print(clusters[clst].rep.path)
        rep_dir = os.path.dirname(
                os.path.abspath(
                    clusters[clst].rep.path
                    )
                )
        relpath = os.path.relpath(
                rep_dir,
                outpath
                )
        # copyfile(clusters[clst].rep.path, os.path.join(outpath,
            # '{}turn_clst{}_rep.pdb.gz'.format(helixlength, clst)))
        os.symlink(
                os.path.join(relpath,
                    os.path.basename(clusters[clst].rep.path)),
                out
                )
        dokfile = sorted(glob.glob(rep_dir + '/*.dok*'))[-1]
        with open(dokfile, 'r') as f:
            for line in f:
                linesplit = line.split(' ')
                filename = os.path.basename(linesplit[-1]).strip('\n')
                if filename == os.path.basename(
                                clusters[clst].rep.path):
                    linesplit.insert(0, outfile)
                    # line.append(os.path.basename(out))
                    scores.write(' '.join(linesplit))

    scores.close()