import numpy as np

import tensorflow as tf

from tensorflow.python.platform import flags

from data_generator import ImageDataGenerator

import logging

from utils import _eval_dice, _connectivity_region_analysis, parse_fn, _crop_object_region, _get_coutour_sample, parse_fn_haus,_eval_haus

import time 

import os

import SimpleITK as sitk

def train(model, saver, sess, train_file_list, test_file, args, resume_itr=0):

    if args.log:

        train_writer = tf.summary.FileWriter(args.log_dir + '/' + args.phase + '/', sess.graph)

    # Data loaders

    with tf.device('/cpu:0'):

        tr_data_list, train_iterator_list, train_next_list = [],[],[]

        for i in range(len(train_file_list)):

            tr_data = ImageDataGenerator(train_file_list[i], mode='training', \

                                         batch_size=args.meta_batch_size, num_classes=args.n_class, shuffle=True)

            tr_data_list.append(tr_data)

            train_iterator_list.append(tf.data.Iterator.from_structure(tr_data.data.output_types,tr_data.data.output_shapes))

            train_next_list.append(train_iterator_list[i].get_next())

    # Ops for initializing different iterators

    training_init_op = []

    train_batches_per_epoch = []

    for i in range(len(train_file_list)):

        training_init_op.append(train_iterator_list[i].make_initializer(tr_data_list[i].data))

        sess.run(training_init_op[i])  # initialize training sample generator at itr=0

    # Training begins

    best_test_dice = 0

    best_test_haus = 0

    for epoch in xrange(0, args.epoch):

        for itr in range(resume_itr, args.train_iterations):

            start = time.time()

            # Sampling training and test tasks

            num_training_tasks = len(train_file_list)

            num_meta_train = 2#num_training_tasks-1

            num_meta_test = 1#num_training_tasks-num_meta_train  # as setting num_meta_test = 1

            # Randomly choosing meta train and meta test domains

            task_list = np.random.permutation(num_training_tasks)

            meta_train_index_list = task_list[:2]

            meta_test_index_list = task_list[-1:]

            # Sampling meta-train, meta-test data

            for i in range(num_meta_train):

                task_ind = meta_train_index_list[i]

                if i == 0:

                    inputa, labela = sess.run(train_next_list[task_ind])

                elif i == 1:

                    inputa1, labela1 = sess.run(train_next_list[task_ind])

                else:

                    raise RuntimeError('check number of meta-train domains.')

            for i in range(num_meta_test):

                task_ind = meta_test_index_list[i]

                if i == 0:

                    inputb, labelb = sess.run(train_next_list[task_ind])

                else:

                    raise RuntimeError('check number of meta-test domains.')

            input_group = np.concatenate((inputa[:2],inputa1[:1],inputb[:2]), axis=0)

            label_group = np.concatenate((labela[:2],labela1[:1],labelb[:2]), axis=0)

            contour_group, metric_label_group = _get_coutour_sample(label_group)

            feed_dict = {model.inputa: inputa, model.labela: labela, \

                         model.inputa1: inputa1, model.labela1: labela1, \

                         model.inputb: inputb, model.labelb: labelb, \

                         model.input_group:input_group, \

                         model.label_group:label_group, \

                         model.contour_group:contour_group, \

                         model.metric_label_group:metric_label_group, \

                         model.KEEP_PROB: 1.0}

            output_tensors = [model.task_train_op, model.meta_train_op, model.metric_train_op]

            output_tensors.extend([model.summ_op, model.seg_loss_b, model.compactness_loss_b, model.smoothness_loss_b, model.target_loss, model.source_loss])

            _, _, _, summ_writer, seg_loss_b, compactness_loss_b, smoothness_loss_b, target_loss, source_loss = sess.run(output_tensors, feed_dict)

            # output_tensors = [model.task_train_op]

            # output_tensors.extend([model.source_loss])

            # _, source_loss = sess.run(output_tensors, feed_dict)

            if itr % args.print_interval == 0:

                logging.info("Epoch: [%2d] [%6d/%6d] time: %4.4f inner lr:%.8f outer lr:%.8f" % (epoch, itr, args.train_iterations, (time.time()-start), model.inner_lr.eval(), model.outer_lr.eval()))

                logging.info('sou_loss: %.7f, tar_loss: %.7f, tar_seg_loss: %.7f, tar_compactness_loss: %.7f, tar_smoothness_loss: %.7f' % (source_loss, target_loss, seg_loss_b, compactness_loss_b, smoothness_loss_b))

            if itr % args.summary_interval == 0:

                train_writer.add_summary(summ_writer, itr)

                train_writer.flush()

            if (itr!=0) and itr % args.save_freq == 0:

                saver.save(sess, args.checkpoint_dir + '/epoch_' + str(epoch) + '_itr_'+str(itr) + ".model.cpkt")

            # Testing periodically

            if (itr!=0) and itr % args.test_freq == 0:

                test_dice, test_dice_arr, test_haus, test_haus_arr = test(sess, test_file, model, args)

                if test_dice > best_test_dice:

                    best_test_dice = test_dice

                with open((os.path.join(args.log_dir,'eva.txt')), 'a') as f:

                    print >> f, 'Iteration %d :' % (itr)

                    print >> f, '   Unseen domain testing results: Dice: %f' %(test_dice), test_dice_arr

                    print >> f, '   Current best accuracy %f' %(best_test_dice)

                    print >> f, '   Unseen domain testing results: Haus: %f' %(test_haus), test_haus_arr

                    print >> f, '   Current best accuracy %f' %(best_test_haus)

                # Save model

def test(sess, test_list, model, args):

    dice = []

    haus = []

    start = time.time()

    with open(test_list, 'r') as fp:

        rows = fp.readlines()

    test_list  = [row[:-1] if row[-1] == '\n' else row for row in rows]

    for fid, filename in enumerate(test_list):

        image, mask, spacing = parse_fn_haus(filename)

        pred_y = np.zeros(mask.shape)

        frame_list = [kk for kk in range(1, image.shape[2] - 1)]

        for ii in xrange(int(np.floor(image.shape[2] // model.test_batch_size))):

            vol = np.zeros([model.test_batch_size, model.volume_size[0], model.volume_size[1], model.volume_size[2]])

            for idx, jj in enumerate(frame_list[ii * model.test_batch_size: (ii + 1) * model.test_batch_size]):

                vol[idx, ...] = image[..., jj - 1: jj + 2].copy()

            pred_student = sess.run((model.outputs), feed_dict={model.test_input: vol, \

                                                                    model.KEEP_PROB: 1.0,\

                                                                    model.training_mode: True})

            for idx, jj in enumerate(frame_list[ii * model.test_batch_size: (ii + 1) * model.test_batch_size]):

                pred_y[..., jj] = pred_student[idx, ...].copy()

        processed_pred_y = _connectivity_region_analysis(pred_y)

        dice_subject = _eval_dice(mask, processed_pred_y)

        # print spacing

        dice.append(dice_subject)

        # haus.append(haus_subject)

        # _save_nii_prediction(mask, processed_pred_y, pred_y, args.result_dir, '_' + filename[-26:-20])

    dice_avg = np.mean(dice, axis=0).tolist()[0]

    # haus_avg = np.mean(haus, axis=0).tolist()[0]

    logging.info("dice_avg %.4f" % (dice_avg))

    # logging.info("haus_avg %.4f" % (haus_avg))

    return dice_avg, dice, 0, 0

    # return dice_avg, dice, haus_avg, haus

def _save_nii_prediction(gth, comp_pred, pre_pred, out_folder, out_bname):

    sitk.WriteImage(sitk.GetImageFromArray(gth), out_folder + out_bname + 'gth.nii.gz') 

    sitk.WriteImage(sitk.GetImageFromArray(pre_pred), out_folder + out_bname + 'premask.nii.gz') 

    sitk.WriteImage(sitk.GetImageFromArray(comp_pred), out_folder + out_bname + 'mask.nii.gz')