import time

import argparse

import torch

import os

import models

from util import util

class BasicParams:

    """

    This class define the console parameters

    """

    def __init__(self):

        """

        Reset the class. Indicates the class hasn't been initialized

        """

        self.initialized = False

        self.isTrain = True

        self.isTest = True

    def initialize(self, parser):

        """

        Define the common console parameters

        """

        parser.add_argument('--gpu_ids', type=str, default='0',

                            help='which GPU would like to use: e.g. 0 or 0,1, -1 for CPU')

        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints',

                            help='models, settings and intermediate results are saved in folder in this directory')

        parser.add_argument('--experiment_name', type=str, default='test',

                            help='name of the folder in the checkpoint directory')

        # Dataset parameters

        parser.add_argument('--omics_mode', type=str, default='a',

                            help='omics types would like to use in the model, options: [abc | ab | a | b | c]')

        parser.add_argument('--data_root', type=str, default='./data',

                            help='path to input data')

        parser.add_argument('--batch_size', type=int, default=32,

                            help='input data batch size')

        parser.add_argument('--num_threads', default=0, type=int,

                            help='number of threads for loading data')

        parser.add_argument('--set_pin_memory', action='store_true',

                            help='set pin_memory in the dataloader to increase data loading performance')

        parser.add_argument('--not_stratified', action='store_true',

                            help='do not apply the stratified mode in train/test split if set true')

        parser.add_argument('--use_sample_list', action='store_true',

                            help='provide a subset sample list of the dataset, store in the path data_root/sample_list.tsv, if False use all the samples')

        parser.add_argument('--use_feature_lists', action='store_true',

                            help='provide feature lists of the input omics data, e.g. data_root/feature_list_A.tsv, if False use all the features')

        parser.add_argument('--detect_na', action='store_true',

                            help='detect missing value markers during data loading, stay False can improve the loading performance')

        parser.add_argument('--file_format', type=str, default='tsv',

                            help='file format of the omics data, options: [tsv | csv | hdf]')

        # Model parameters

        parser.add_argument('--model', type=str, default='vae_classifier',

                            help='chooses which model want to use, options: [vae_classifier | vae_regression | vae_survival | vae_multitask]')

        parser.add_argument('--net_VAE', type=str, default='fc_sep',

                            help='specify the backbone of the VAE, default is the one dimensional CNN, options: [conv_1d | fc_sep | fc]')

        parser.add_argument('--net_down', type=str, default='multi_FC_classifier',

                            help='specify the backbone of the downstream task network, default is the multi-layer FC classifier, options: [multi_FC_classifier | multi_FC_regression | multi_FC_survival | multi_FC_multitask]')

        parser.add_argument('--norm_type', type=str, default='batch',

                            help='the type of normalization applied to the model, default to use batch normalization, options: [batch | instance | none ]')

        parser.add_argument('--filter_num', type=int, default=8,

                            help='number of filters in the last convolution layer in the generator')

        parser.add_argument('--conv_k_size', type=int, default=9,

                            help='the kernel size of convolution layer, default kernel size is 9, the kernel is one dimensional.')

        parser.add_argument('--dropout_p', type=float, default=0.2,

                            help='probability of an element to be zeroed in a dropout layer, default is 0 which means no dropout.')

        parser.add_argument('--leaky_slope', type=float, default=0.2,

                            help='the negative slope of the Leaky ReLU activation function')

        parser.add_argument('--latent_space_dim', type=int, default=128,

                            help='the dimensionality of the latent space')

        parser.add_argument('--seed', type=int, default=42,

                            help='random seed')

        parser.add_argument('--init_type', type=str, default='normal',

                            help='choose the method of network initialization, options: [normal | xavier_normal | xavier_uniform | kaiming_normal | kaiming_uniform | orthogonal]')

        parser.add_argument('--init_gain', type=float, default=0.02,

                            help='scaling factor for normal, xavier and orthogonal initialization methods')

        # Loss parameters

        parser.add_argument('--recon_loss', type=str, default='BCE',

                            help='chooses the reconstruction loss function, options: [BCE | MSE | L1]')

        parser.add_argument('--reduction', type=str, default='mean',

                            help='chooses the reduction to apply to the loss function, options: [sum | mean]')

        parser.add_argument('--k_kl', type=float, default=0.01,

                            help='weight for the kl loss')

        parser.add_argument('--k_embed', type=float, default=0.001,

                            help='weight for the embedding loss')

        # Other parameters

        parser.add_argument('--deterministic', action='store_true',

                            help='make the model deterministic for reproduction if set true')

        parser.add_argument('--detail', action='store_true',

                            help='print more detailed information if set true')

        parser.add_argument('--epoch_to_load', type=str, default='latest',

                            help='the epoch number to load, set latest to load latest cached model')

        parser.add_argument('--experiment_to_load', type=str, default='test',

                            help='the experiment to load')

        self.initialized = True  # set the initialized to True after we define the parameters of the project

        return parser

    def get_params(self):

        """

        Initialize our parser with basic parameters once.

        Add additional model-specific parameters.

        """

        if not self.initialized:  # check if this object has been initialized

            # if not create a new parser object

            parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

            #  use our method to initialize the parser with the predefined arguments

            parser = self.initialize(parser)

        # get the basic parameters

        param, _ = parser.parse_known_args()

        # modify model-related parser options

        model_name = param.model

        model_param_setter = models.get_param_setter(model_name)

        parser = model_param_setter(parser, self.isTrain)

        # save and return the parser

        self.parser = parser

        return parser.parse_args()

    def print_params(self, param):

        """

        Print welcome words and command line parameters.

        Save the command line parameters in a txt file to the disk

        """

        message = ''

        message += '\nWelcome to OmiEmbed\nby Xiaoyu Zhang x.zhang18@imperial.ac.uk\n\n'

        message += '-----------------------Running Parameters-----------------------\n'

        for key, value in sorted(vars(param).items()):

            comment = ''

            default = self.parser.get_default(key)

            if value != default:

                comment = '\t[default: %s]' % str(default)

            message += '{:>18}: {:<15}{}\n'.format(str(key), str(value), comment)

        message += '----------------------------------------------------------------\n'

        print(message)

        # Save the running parameters setting in the disk

        experiment_dir = os.path.join(param.checkpoints_dir, param.experiment_name)

        util.mkdir(experiment_dir)

        file_name = os.path.join(experiment_dir, 'cmd_parameters.txt')

        with open(file_name, 'w') as param_file:

            now = time.strftime('%c')

            param_file.write('{:s}\n'.format(now))

            param_file.write(message)

            param_file.write('\n')

    def parse(self):

        """

        Parse the parameters of our project. Set up GPU device. Print the welcome words and list parameters in the console.

        """

        param = self.get_params()  # get the parameters to the object param

        param.isTrain = self.isTrain

        param.isTest = self.isTest

        # Print welcome words and command line parameters

        self.print_params(param)

        # Set the internal parameters

        # epoch_num: the total epoch number

        if self.isTrain:

            param.epoch_num = param.epoch_num_p1 + param.epoch_num_p2 + param.epoch_num_p3

        # downstream_task: for the classification task a labels.tsv file is needed, for the regression task a values.tsv file is needed

        if param.model == 'vae_classifier':

            param.downstream_task = 'classification'

        elif param.model == 'vae_regression':

            param.downstream_task = 'regression'

        elif param.model == 'vae_survival':

            param.downstream_task = 'survival'

        elif param.model == 'vae_multitask' or param.model == 'vae_multitask_gn':

            param.downstream_task = 'multitask'

        elif param.model == 'vae_alltask' or param.model == 'vae_alltask_gn':

            param.downstream_task = 'alltask'

        else:

            raise NotImplementedError('Model name [%s] is not recognized' % param.model)

        # add_channel: add one extra dimension of channel for the input data, used for convolution layer

        # ch_separate: separate the DNA methylation matrix base on the chromosome

        if param.net_VAE == 'conv_1d':

            param.add_channel = True

            param.ch_separate = False

        elif param.net_VAE == 'fc_sep':

            param.add_channel = False

            param.ch_separate = True

        elif param.net_VAE == 'fc':

            param.add_channel = False

            param.ch_separate = False

        else:

            raise NotImplementedError('VAE model name [%s] is not recognized' % param.net_VAE)

        # omics_num: the number of omics types

        param.omics_num = len(param.omics_mode)

        # Set up GPU

        str_gpu_ids = param.gpu_ids.split(',')

        param.gpu_ids = []

        for str_gpu_id in str_gpu_ids:

            int_gpu_id = int(str_gpu_id)

            if int_gpu_id >= 0:

                param.gpu_ids.append(int_gpu_id)

        if len(param.gpu_ids) > 0:

            torch.cuda.set_device(param.gpu_ids[0])

        self.param = param

        return self.param