import sys

import os

import numpy as np

import tensorflow as tf

from tensorflow.python.platform import flags

from data_generator import ImageDataGenerator

from saml_func import SAML

from train import train

from train import test

import datetime

import argparse

from utils import check_folder, show_all_variables

import logging

currtime = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')

tf.set_random_seed(2)

def parse_args(train_date):

    desc = "Tensorflow implementation of DenseUNet for prostate segmentation"

    parser = argparse.ArgumentParser(description=desc)

    parser.add_argument('--gpu', type=str, default='0', help='train or test or guide')

    parser.add_argument('--phase', type=str, default='train', help='train or test or guide')

    parser.add_argument('--n_class', type=int, default=2, help='The size of class')

    ## Training operations

    parser.add_argument('--target_domain', type=str, default='ISBI', help='dataset_name')

    parser.add_argument('--volume_size', type=list, default=[384, 384, 3], help='The size of input data')

    parser.add_argument('--label_size', type=list, default=[384, 384, 1], help='The size of label')

    parser.add_argument('--epoch', type=int, default=1, help='The number of epochs to run')

    parser.add_argument('--train_iterations', type=int, default=10000, help='The number of training iterations')

    parser.add_argument('--meta_batch_size', type=int, default=5, help='number of images sampled per source domain')

    parser.add_argument('--test_batch_size', type=int, default=1, help='number of images sampled per source domain')

    parser.add_argument('--inner_lr', type=float, default=1e-4, help='The learning rate')

    parser.add_argument('--outer_lr', type=float, default=1e-3, help='The learning rate')

    parser.add_argument('--metric_lr', type=float, default=1e-3, help='The learning rate')

    parser.add_argument('--margin', type=float, default=10.0, help='The learning rate')

    parser.add_argument('--compactness_loss_weight', type=float, default=1.0, help='The learning rate')

    parser.add_argument('--smoothness_loss_weight', type=float, default=0.005, help='The learning rate')

    parser.add_argument('--clipNorm', type=int, default=True, help='number of images sampled per source domain')

    parser.add_argument('--gradients_clip_value', type=float, default=10.0, help='The learning rate')

    # Logging, saving, and testing options

    parser.add_argument('--resume', type=int, default=False, help='number of images sampled per source domain')

    parser.add_argument('--log', type=int, default=True, help='write tensorboard')

    parser.add_argument('--decay_step', type=float, default=500, help='The learning rate')

    parser.add_argument('--decay_rate', type=float, default=0.95, help='The learning rate')

    parser.add_argument('--test_freq', type=int, default=200, help='The number of ckpt_save_freq')

    parser.add_argument('--save_freq', type=int, default=200, help='The number of ckpt_save_freq')

    parser.add_argument('--print_interval', type=int, default=5, help='The frequency to write tensorboard')

    parser.add_argument('--summary_interval', type=int, default=20, help='The frequency to write tensorboard')

    parser.add_argument('--restored_model', type=str, default=None, help='Model to restore')

    parser.add_argument('--test_model', type=str, default=None, help='Model to restore')

    # parser.add_argument('--dropout', type=str, default=1, help='dropout rate')

    # parser.add_argument('--cost_kwargs', type=str, default=1, help='cost_kwargs')

    # parser.add_argument('--opt_kwargs', type=str, default=1, help='opt_kwargs')

    parser.add_argument('--checkpoint_dir', type=str, default='../output/' + train_date + '/checkpoints/' ,

                        help='Directory name to save the checkpoints')

    parser.add_argument('--result_dir', type=str, default='../output/' + train_date + '/results/',

                        help='Directory name to save the generated images')

    parser.add_argument('--log_dir', type=str, default='../output/' + train_date + '/logs/',

                        help='Directory name to save training logs')

    parser.add_argument('--sample_dir', type=str, default='../output/' + train_date + '/samples/',

                        help='Directory name to save the samples on training')

    return check_args(parser.parse_args())

"""checking arguments"""

def check_args(args):

    # --checkpoint_dir

    check_folder(args.checkpoint_dir)

    # --result_dir

    check_folder(args.result_dir)

    # --result_dir

    check_folder(args.log_dir)

    # --sample_dir

    check_folder(args.sample_dir)

    return args

def main():

    train_date = 'xxx'

    args = parse_args(train_date)

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    # define logger

    logging.basicConfig(filename=args.log_dir+"/"+args.phase+'_log.txt', level=logging.DEBUG, format='%(asctime)s %(message)s')

    logging.getLogger().addHandler(logging.StreamHandler())

    # print all parameters

    logging.info("Usage:")

    logging.info("    {0}".format(" ".join([x for x in sys.argv]))) 

    logging.debug("All settings used:")

    os.system('cp main.py %s' % (args.log_dir)) # bkp of train procedure

    os.system('cp saml_func.py %s' % (args.log_dir)) # bkp of train procedure

    os.system('cp train.py %s' % (args.log_dir)) # bkp of train procedure

    os.system('cp utils.py %s' % (args.log_dir)) # bkp of train procedure

    os.system('cp data_generator.py %s' % (args.log_dir))

    filelist_root = '../dataset'

    source_list = ['HK', 'ISBI', 'ISBI_1.5', 'I2CVB','UCL', 'BIDMC']#'ISBI_1.5', 'I2CVB', 'UCL','BIDMC']#, 'I2CVB', 'ISBI_1.5', 'UCL', 'BIDMC']#'I2CVB', 'UCL', 'BIDMC', 'HK']

    source_list.remove(args.target_domain)

    # Constructing model

    model = SAML(args)

    model.construct_model_train()

    model.construct_model_test()

    model.summ_op = tf.summary.merge_all()

    saver = tf.train.Saver(tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES))

    sess = tf.InteractiveSession()

    tf.global_variables_initializer().run()

    show_all_variables()

    # restore model ----

    resume_itr = 0

    model_file = None

    if args.resume:

        model_file = tf.train.latest_checkpoint(args.checkpoint_dir)

        if model_file:

            ind1 = model_file.index('model')

            resume_itr = int(model_file[ind1+5:])

            print("Restoring model weights from " + model_file)

            saver.restore(sess, model_file)

    train_file_list = [os.path.join(filelist_root, source_domain+'_train_list') for source_domain in source_list]

    test_file_list = [os.path.join(filelist_root, args.target_domain+'_train_list')]

    # start training ----

    if args.phase == 'train':

        train(model, saver, sess, train_file_list, test_file_list[0], args, resume_itr)

    else:

        args.test_model = 'xxx'

        saver.restore(sess, args.test_model)

        logging.info("testing model restored %s" % args.test_model)

        test_dice, test_dice_arr, test_haus, test_haus_arr = test(sess, test_file_list[0], model, args)

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

            print >> f, 'testing model %s :' % (args.test_model)

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

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

if __name__ == "__main__":

    main()