import argparse
import os
import numpy as np
from sklearn.preprocessing import MultiLabelBinarizer

from functions import model_functions as model_func
from functions import general_functions as general
from functions import featurizing_functions as ft

parser = argparse.ArgumentParser(description='Training using compound properties features.')

parser.add_argument('featurizer', type=str, help='type of descriptor (choice: daylight, ecfp)')
parser.add_argument('dataset', type=str,
                    help='dataset to be used (choice: dataset1.pkl, dataset2.pkl)')
parser.add_argument('method', type=str,
                    help='ML methods to use for training (choice: [simplenn, randomforest, '
                         'knearest, gradientboosting, svm, convnn])')
parser.add_argument('--num_split', type=int, default=10,
                    help='Number of splits for averaging of multiclass classification performance.')
parser.add_argument('--seed', type=int, default=7,
                    help='seed number for random shuffling.')
parser.add_argument('--save_model', type=bool, default=True,
                    help='Whether trained model is saved.')
parser.add_argument('--model_param', nargs='*', default=0,
                    help='Input parameter for model used [IMPORTANT: ORDERING OF INPUT PARAMETER '
                         'MATTERS](see function define_model for better understanding).')
parser.add_argument('--filename_append', type=str, default='',
                    help='Filename extension to distinct runs from each other.')
parser.add_argument('--loss', type=str, default='binary_crossentropy')
parser.add_argument('--optimizer', type=str, default='adam')
parser.add_argument('--epochs', type=int, default=20)
args = parser.parse_args()

# Check if inputs are available
if args.method not in ['simplenn', 'gradientboosting', 'randomforest', 'knearest']:
    raise Exception("%s not available. Choose from simplenn, gradientboosting, randomforest or knearest." %args.method)
if args.featurizer not in ['daylight', 'ecfp']:
    raise Exception("Descriptor %s not available. Choose from daylight or ecfp." %args.featurizer)

filename = "%s/%s_%s_%s" %(args.featurizer, args.method, args.dataset[:args.dataset.rfind('.')], args.filename_append)
logfile = "./logfile/" + filename + ".log"

general.check_path_exists(logfile)
sys.stdout = open(logfile, 'wt')

# Import data
temp_path = general.file_pathway(args.dataset)
if os.path.exists(temp_path):
    data = general.import_pandas_dataframe(temp_path)
    print("Shape of data:", data.shape)
else:
    raise Exception("%s does not exist." % args.dataset)

if not 'fingerprint' in data.columns:
    if args.featurizer=='daylight':
        print ("Daylight Fingerprinting...")
        data['fingerprint'] = data['mol'].apply(ft.get_rdk)
        print ("Daylight Fingerprinting done.")
    elif args.featurizer=='ecfp':
        print ("Calculating ECFP...")
        data['fingerprint'] = data['mol'].apply(ft.get_ecfp)
        print ("ECFP calculation done.")

# Input (x) and class (y)
X = np.array(np.stack(data['fingerprint']), dtype=int)

mlb = MultiLabelBinarizer().fit(data['agrochemical'])
Y = mlb.transform(data['agrochemical'])

# Machine learning algorithms to use
method = args.method
model_param = args.model_param

num_split = args.num_split
seed = args.seed

if method in ['simplenn']:
    layers_dim = [X.shape[1], 512, 128, 32, 8, Y.shape[1]]
    activation = ['relu', 'relu', 'relu', 'relu', 'sigmoid']

    epochs = args.epochs
    print("Number of epochs:", epochs)

    model_func.define_model(method)
    model = model_func.build_simplenn_model(layers_dim=layers_dim, activation=activation,
                                            loss=args.loss, optimizer=args.optimizer)

else:
    epochs = 0
    model = model_func.define_model(method, model_param)

model = model_func.train_model(model, num_split, seed, X, Y,
                               neural_network_epochs=epochs)

if args.save_model:
    save_filepath = './saved_model/' + filename + '.h'
    general.check_path_exists(save_filepath)
    model_func.save_model(model, save_filepath, epochs)

print("\nExiting program.")
sys.stdout.close()
