import os
import logging
import tensorflow as tf
from medseg_dl import parameters
from medseg_dl.utils import utils_data, utils_misc
from medseg_dl.model import input_fn
from medseg_dl.model import model_fn
from medseg_dl.model import training
import sys


def main(dir_model, non_local='disable', attgate='disable', device=None, idx_dataset=0):
    # Since this is the training script all variables should be trainable
    b_training = True  # hardcoded since scripts are not identical

    # Load / generate parameters from model file, if available
    file_params = os.path.join(dir_model, 'params.yaml')
    if os.path.isfile(file_params):
        params = parameters.Params(path_yaml=file_params)
    else:
        params = parameters.Params(model_dir=dir_model, idx_dataset=idx_dataset)

    # Set logger
    utils_misc.set_logger(os.path.join(params.dict['dir_model'], 'train.log'), params.dict['log_level'])

    # Set random seed
    if params.dict['b_use_seed']:
        tf.set_random_seed(params.dict['random_seed'])

    # Set device for graph calc
    if device:
        os.environ['CUDA_VISIBLE_DEVICES'] = device
    else:
        os.environ['CUDA_VISIBLE_DEVICES'] = params.dict['device']

    """ Fetch data, generate pipeline and model """
    tf.reset_default_graph()

    # Fetch datasets, atm. saved as json
    logging.info('Fetching the datasets...')
    filenames_train, filenames_eval = utils_data.load_sets(params.dict['dir_data'],
                                                           params.dict['dir_model'],
                                                           path_parser_cfg=params.dict['path_parser_cfg'],
                                                           set_split=params.dict['set_split'],
                                                           b_recreate=True)

    # Create a tf.data pipeline
    logging.info('Creating the pipeline...')
    spec_pipeline = input_fn.gen_pipeline_train(filenames=filenames_train,
                                                shape_image=params.dict['shape_image'],
                                                shape_input=params.dict['shape_input'],
                                                shape_output=params.dict['shape_output'],
                                                channels_out=params.dict['channels_out'],
                                                size_batch=params.dict['size_batch'],
                                                size_buffer=params.dict['size_buffer'],
                                                num_parallel_calls=params.dict['num_parallel_calls'],
                                                repeat=params.dict['repeat'],
                                                b_shuffle=params.dict['b_shuffle'],
                                                patches_per_class=params.dict['patches_per_class'],
                                                sigma_offset=params.dict['sigma_offset'],
                                                sigma_noise=params.dict['sigma_noise'],
                                                sigma_pos=params.dict['sigma_pos'],
                                                b_mirror=params.dict['b_mirror'],
                                                b_rotate=params.dict['b_rotate'],
                                                b_scale=params.dict['b_scale'],
                                                b_warp=params.dict['b_warp'],
                                                b_permute_labels=params.dict['b_permute_labels'],
                                                angle_max=params.dict['angle_max'],
                                                scale_factor=params.dict['scale_factor'],
                                                delta_max=params.dict['delta_max'],
                                                b_verbose=False)



    # Create the model (incorporating losses, optimizer, metrics)
    logging.info('Creating the model...')
    spec_model = model_fn.model_fn(spec_pipeline,
                                   input_metrics=None,
                                   b_training=b_training,
                                   channels=params.dict['channels'],
                                   channels_out=params.dict['channels_out'],
                                   batch_size=params.dict['size_batch'],
                                   b_dynamic_pos_mid=params.dict['b_dynamic_pos_mid'],
                                   b_dynamic_pos_end=params.dict['b_dynamic_pos_end'],
                                   non_local=non_local,
                                   non_local_num=1,
                                   attgate=attgate,
                                   filters=params.dict['filters'],
                                   dense_layers=params.dict['dense_layers'],
                                   alpha=params.dict['alpha'],
                                   dropout_rate=params.dict['rate_dropout'],
                                   rate_learning=params.dict['rate_learning'],
                                   beta1=params.dict['beta1'],
                                   beta2=params.dict['beta2'],
                                   epsilon=params.dict['epsilon'])


    # Train the actual model
    logging.info('Starting training for %i epoch(s)', params.dict['num_epochs'])
    training.sess_train(spec_pipeline, spec_model, params)


if __name__ == '__main__':
    # model
    model_dir = '/home/stage13_realshuffle_attgate_batch24_l2_repeat10_shuffle_eval_reminder_timout6000_K'

    # crucial parameters
    logging.info('Received the following input(s): %s', str(sys.argv))

    non_local = 'disable'
    if len(sys.argv) > 1:
        non_local = str(sys.argv[1])
    logging.info('Adding nonlocal block after %s', non_local)

    attgate = 'active'
    if len(sys.argv) > 2:
        attgate = istr(sys.argv[2])
    logging.info('Adding attention gate after %s', attgate)

    device = '3'
    if len(sys.argv) > 3:
        device = str(sys.argv[3])
    logging.info('Calculating on device %s', device)

    idx_dataset = 0
    if len(sys.argv) > 4:
        idx_dataset = int(sys.argv[4])
    logging.info('Using dataset %i', idx_dataset)

    main(model_dir, non_local=non_local, attgate=attgate, device=device, idx_dataset=idx_dataset)