import numpy as np
from texttable import Texttable
import os, sys, json, time, gc, math, datetime
from shutil import copyfile
from datageneratorNew import data_provide
from sklearn.metrics import classification_report, confusion_matrix
from keras.models import load_model
from keras.callbacks import Callback, ModelCheckpoint, LambdaCallback
from keras.utils.np_utils import categorical_probas_to_classes
np.random.seed(2004)

def loadmodelmodul(modelpath):
    from importlib.machinery import SourceFileLoader
    model = SourceFileLoader("model.py", modelpath).load_module()

    return model

def train(model, trainsettings, traindata_gen, num_trainsamples, validdata_gen, num_validsamples, start_from_epoch, batchSize):
    tvh = model.fit_generator(generator=traindata_gen,
                              steps_per_epoch=math.ceil(num_trainsamples/batchSize),
                              validation_data=validdata_gen,
                              validation_steps=math.ceil(num_validsamples/batchSize),
                              initial_epoch=start_from_epoch,
                              **trainsettings)
    return tvh


def validate(model, validdata_gen, no_samples, batchSize):
    vprobs, vclass, vtrues = _predict(model, validdata_gen, batchSize, nosamples=no_samples)
    return vprobs, vclass, vtrues


def test(model, testdata_gen, no_samples, batchSize):
    tprobs, tclass, ttrues = _predict(model, testdata_gen, batchSize, nosamples=no_samples)
    return tprobs, tclass, ttrues



def _predict(model, datagenerator, batchSize, nosamples):
    numIter = math.ceil(nosamples/batchSize)
    pprobs, ptrues = [], []

    for iter in range(0,int(numIter)):
        x, y = next(datagenerator)
        prob = model.predict(x, batch_size=batchSize, verbose=2)
        for x in range (0,batchSize):
            pprobs.append(prob[x])
            ptrues.append(y[x])
        print("iteration: " + str(iter))
    pclass = categorical_probas_to_classes(pprobs)
    ptrues = categorical_probas_to_classes(ptrues)

    return pprobs, pclass, ptrues

def create_model_folder(path_to_exp_folder, model_id):
    path_to_mod_folder = os.path.join(path_to_exp_folder, 'Model-{0}'.format(model_id))
    if not os.path.exists(path_to_mod_folder):
        os.makedirs(path_to_mod_folder)
    return path_to_mod_folder


def save_experiment_code(path_to_mod_folder, model_id, parameters):
    src = '{0}/experimentNew.py'.format(os.getcwd())
    dst = '{0}/{1}_exp_code.py'.format(path_to_mod_folder, model_id)
    copyfile(src, dst)
    src = '{0}/datageneratorNew.py'.format(os.getcwd())
    dst = '{0}/{1}_data_code.py'.format(path_to_mod_folder, model_id)
    copyfile(src, dst)
    src = '{0}/model.py'.format(os.getcwd())
    dst = '{0}/{1}_model_code.py'.format(path_to_mod_folder, model_id)
    copyfile(src, dst)

    with open(dst, "a+") as fl:
        fl.write('\n\n{0}\n###   Used parameters: {1}\n{0}'.format(('#' * 100), parameters))


def save_model_summary(model, path_to_mod_folder, model_id):
    summarypath = '{0}/{1}_mod_summ.txt'.format(path_to_mod_folder, model_id)
    originalstdout = sys.stdout
    sys.stdout = open(summarypath, 'w')
    print(model.summary())
    sys.stdout = originalstdout


def save_model(model, path_to_mod_folder, model_id, epoch=False, last_only=True):

    if last_only:
        delete_unnecessary_models(path_to_mod_folder)

    modelpath = '{0}/{1}-mod_{2}.h5'.format(path_to_mod_folder, model_id, epoch)
    model.save(modelpath, overwrite=True)


def load_saved_model(path_to_mod_folder, model_id, epoch=None, best=False):
    if epoch is None and not best: raise ValueError('Epoch parameter of load_saved_modul() must be specified unless best=True.')
    if best: epoch = 'best'
    else: epoch = epoch -1
    modelpath = '{0}/{1}-mod_{2}.h5'.format(path_to_mod_folder, model_id, epoch)
    return load_model(modelpath)


def set_saving_of_last_model(trainsettings, path_to_mod_folder, model_id, model):
    save_last_model_callback = LambdaCallback(
        on_epoch_end=lambda epoch, logs: save_model(model, path_to_mod_folder, model_id, epoch=epoch, last_only=True))
    trainsettings['callbacks'].append(save_last_model_callback)


def set_saving_of_best_model(trainsettings, path_to_mod_folder, model_id, monitor='val_acc'):
    modelpath = '{0}/{1}-mod_best.h5'.format(path_to_mod_folder, model_id)
    if not any([isinstance(x, ModelCheckpoint) for x in trainsettings['callbacks']]):
        model_checkpoint = ModelCheckpoint(
            filepath=modelpath,
            monitor=monitor,
            verbose=0,
            save_best_only=True,
            mode='auto')
        trainsettings['callbacks'].append(model_checkpoint)


class RunningHistory(Callback):
    def __init__(self, filepath, initial_run_hist=None):
        super(RunningHistory, self).__init__()
        self.filepath = filepath
        self.initial_run_hist = initial_run_hist

    def on_train_begin(self, logs={}):
        if self.initial_run_hist is None:
            self.run_hist = {'loss': [], 'acc': [], 'val_loss': [], 'val_acc': [], 'epoch': 0}
        else:
            self.run_hist = self.initial_run_hist

    def on_epoch_end(self, epoch, logs={}):
        self.run_hist['loss'].append(logs.get('loss'))
        self.run_hist['acc'].append(logs.get('acc'))
        self.run_hist['val_loss'].append(logs.get('val_loss'))
        self.run_hist['val_acc'].append(logs.get('val_acc'))
        self.run_hist['epoch'] = epoch
        self.save_run_hist()

    def save_run_hist(self):
        with open(self.filepath, "w+") as fl:
            rh = _convert_hist_to_float(self.run_hist)
            json.dump(rh, fl)


def set_returning_running_history(trainsettings, path_to_mod_folder, model_id, initial_run_hist=None):
    run_hist_path = '{0}/{1}-run_hist.json'.format(path_to_mod_folder, model_id)
    if not any([isinstance(x, RunningHistory) for x in trainsettings['callbacks']]):
        run_hist_callback = RunningHistory(
            filepath=run_hist_path,
            initial_run_hist=initial_run_hist)
        trainsettings['callbacks'].append(run_hist_callback)


def delete_unnecessary_models(path_to_mod_folder):
    previous_models = [x for x in os.listdir(path_to_mod_folder) if 'h5' in x and 'best' not in x]
    [os.remove(os.path.join(path_to_mod_folder, previous)) for previous in previous_models]



def delete_unnecessary_run_hists(path_to_mod_folder):
    run_hist_files = [x for x in os.listdir(path_to_mod_folder) if '-run_hist.json' in x]
    [os.remove(os.path.join(path_to_mod_folder, ruh_hist)) for ruh_hist in run_hist_files]


def _convert_hist_to_float(hist_dict):
    for key, value in hist_dict.items():
        if type(value) is list:
            hist_dict[key] = list(map(float, value))
    return hist_dict


def get_train_stats_dict(tvh, start_time, end_time, path_to_mod_folder, model_id, data_name):
    trainstats = {}

    trainstats['data_name'] = data_name
    trainstats['model_id'] = model_id
    trainstats['path_to_mod_folder'] = path_to_mod_folder

    trainstats['tvh_results'] = tvh

    total_epochs = len(tvh['loss'])
    trainstats['total_epochs'] = total_epochs

    best_epoch = tvh['val_loss'].index(min(tvh['val_loss'])) + 1
    trainstats['best_epoch'] = best_epoch

    elapsed_time = end_time - start_time
    trainstats['elapsed_time'] = elapsed_time

    one_epoch_duration = elapsed_time / (total_epochs - start_from_epoch)
    trainstats['one_epoch_duration'] = one_epoch_duration

    total_time_s = total_epochs * one_epoch_duration
    trainstats['total_time_s'] = total_time_s

    total_time_m = total_time_s / 60
    trainstats['total_time_m'] = total_time_m

    best_epoch_time_m = total_time_m / (total_epochs / best_epoch)
    trainstats['best_epoch_time_m'] = best_epoch_time_m

    return trainstats

def save_tvhistory_json(trainstats):
    results = trainstats['tvh_results']
    results = _convert_hist_to_float(results)
    histpath = '{0}/{1}-tvh_dict.json'.format(trainstats['path_to_mod_folder'], trainstats['model_id'])
    with open(histpath, "a+") as fl:
        json.dump(results, fl)


def load_run_hist_json(path_to_mod_folder, model_id):
    tvhistpath = '{0}/{1}-run_hist.json'.format(path_to_mod_folder, model_id)
    with open(tvhistpath, 'r') as fl:
        return json.load(fl)


def save_model_results(trainstats, list_of_class_names, vclass, vtrues, tclass=None, ttrues=None):

    tvlogpath = '{0}/{1}_tvh_log.txt'.format(trainstats['path_to_mod_folder'], trainstats['model_id'])

    originalstdout = sys.stdout
    sys.stdout = open(tvlogpath, 'a+')

    print_tvresults(trainstats)

    print_classification_report(vtrues, vclass, list_of_class_names, name='VALIDATION')
    print_confusion_matrix(vtrues, vclass, list_of_class_names, name='VALIDATION')

    if not any([tclass is None, ttrues is None, list_of_class_names is None]):
        print_classification_report(ttrues, tclass, list_of_class_names, name='TESTING')
        print_confusion_matrix(ttrues, tclass, list_of_class_names, name='TESTING')

    sys.stdout = originalstdout


def save_data_info(trainstats, num_train, num_valid, num_test):
    tvlogpath = '{0}/{1}_tvh_log.txt'.format(trainstats['path_to_mod_folder'], trainstats['model_id'])
    with open(tvlogpath, 'a+') as fl:
        print('DATA USED: {0}\n'.format(trainstats['data_name']), file=fl)
        print('\tNo.train: {0} No.valid: {1} No.test: {2}'.format(num_train, num_valid, num_test), file=fl)


def print_tvresults(trainstats):
    ts = trainstats
    results = trainstats['tvh_results']
    b_e = trainstats['best_epoch']

    print('\nMODEL {0} - RESULTS OF TRAINING AND VALIDATION'.format(ts['model_id']))
    print('\tTraining of {0} epochs took {1:.2f} minutes. Best epoch = {2}. Training took {3:.2f} minutes'
          .format(ts['total_epochs'], ts['total_time_m'], b_e, ts['best_epoch_time_m']))
    print('\tTrain acc: \t{0:.2f}%'.format(results['acc'][b_e - 1] * 100))
    print('\tTrain loss:\t{0:.2f}%'.format(results['loss'][b_e - 1] * 100))
    print('\tValid acc: \t{0:.2f}%'.format(results['val_acc'][b_e - 1] * 100))
    # Coloring to console '\033[92m' - green '\033[0m' - end of color
    print('\tValid loss:\t{0:.2f}%'.format(results['val_loss'][b_e - 1] * 100))
    print()


def save_train_epoch_log(trainstats):
    results = trainstats['tvh_results']
    tvlogpath = '{0}/{1}_tvh_log.txt'.format(trainstats['path_to_mod_folder'], trainstats['model_id'])
    digits = len(str(trainstats['total_epochs']))

    with open(tvlogpath, 'a+') as fl:
        print('\nTRAINING LOG:', file=fl)
        for epoch, value in enumerate(zip(results['loss'], results['acc'], results['val_loss'], results['val_acc'])):
            best = '***' if epoch + 1 == trainstats['best_epoch'] else ''  # Denotes best epoch in log with three stars
            print('{0:.1f}s |'.format(trainstats['one_epoch_duration']),
                  'Epoch {0:0>{digits}} /'.format(epoch + 1, digits=digits),
                  '{0:0>{digits}}: '.format(trainstats['total_epochs'], digits=digits),
                  'tl:{0:.4f}\tta:{1:.4f}\tvl:{2:.4f}\tva:{3:.4f} \t{4}'
                  .format(value[0], value[1], value[2], value[3], best), file=fl)


def save_file_result_log(trainstats, vprobs, vclass, vtrues, tprobs, tclass, ttrues, list_of_valid, list_of_test):
    filelogpath = '{0}/{1}_res_log.txt'.format(trainstats['path_to_mod_folder'], trainstats['model_id'])
    val = [list_of_valid, vtrues, vclass, vprobs, 'VALIDATION']
    test = [list_of_test, ttrues, tclass, tprobs, 'TESTING']

    with open(filelogpath, 'a+') as fl:

        def gen_file_results_log(lof, ftrue, fpred, fprob, name):
            if lof is not None:
                print('\n{0} FILE PREDICTIONS wth confidence:'.format(name), file=fl)
                [print('{0}\ttrue: {1}\tpred: {2}, \tprob: {3}, \t name: {4}'.format(true == pred, true, pred, prob, name), file=fl)
                 for true, pred, prob, name in zip(ftrue, fpred, fprob, lof)]

        for x in [val, test]:
            gen_file_results_log(lof=x[0], ftrue=x[1], fpred=x[2], fprob=x[3], name=x[4])


def print_classification_report(y_true, y_predicted, list_of_class_names, name='VALIDATION'):
    if any([y_true is None, y_predicted is None, list_of_class_names is None]):
        print('\nNo data was provided for creating {0} classification report.'.format(name))
    else:
        print('\n{0} | CLASSIFICATION REPORT:'.format(name))
        print(classification_report(y_true, y_predicted, target_names=list_of_class_names))


def print_confusion_matrix(y_true, y_predicted, list_of_class_names, name='VALIDATION'):
    if any([y_true is None, y_predicted is None, list_of_class_names is None]):
        print('\nNo data was provided for creating {0} confusion matrix.'.format(name))
    else:
        confmat = confusion_matrix(y_true, y_predicted)

        tab = Texttable()

        tab.header(['CONFUSION MATRIX'] +
                   ['true: {0}'.format(list_of_class_names[i]) for i in range(len(confmat[0]))] +
                   ['class recall'])

        lastrow = ['class precision:']
        diagonal = 0
        for i, valrow in enumerate(confmat):
            tab.add_row(['pred.: {0}'.format(list_of_class_names[i])] + list(valrow) + [
                '{0: .2f} %'.format((valrow[i] / sum(valrow)) * 100)])
            lastrow.append(
                '{0: .2f} %'.format((valrow[i] / sum([confmat[j][i] for j in range(len(confmat[0]))])) * 100))
            diagonal += valrow[i]
        accuracy = diagonal / sum([sum(confmat[i]) for i in range(len(confmat[0]))])
        lastrow.append('acc: {0: .2f} %'.format(accuracy * 100))
        tab.add_row(lastrow)
        tab.set_deco(tab.HEADER | tab.VLINES | tab.HLINES)
        confmatascii = tab.draw()

        print('\n{0} | CONFUSION MATRIX:'.format(name))
        print(confmatascii)
        print()


def rename_model_folder(trainstats):
    results = trainstats['tvh_results']
    model_folder_name = os.path.split(trainstats['path_to_mod_folder'])[-1]

    best_epoch = results['val_loss'].index(min(results['val_loss'])) + 1
    acc_t = results['acc'][best_epoch - 1] * 100
    acc_v = results['val_acc'][best_epoch - 1] * 100

    src = trainstats['path_to_mod_folder']
    dst = os.path.join(path_to_exp_folder, 'V{0:.2f}-T{1:.2f}-BE{2}_{3}'
                       .format(acc_v, acc_t, best_epoch, model_folder_name))
    os.rename(src, dst)


def run(path_to_exp_folder, exp_id, param_id, GPU_core, start_from_epoch):
    gc.disable()
    np.random.seed(2004)

    traindata_gen, validdata_gen, testdata_gen, \
    num_trainsamples, num_validsamples, num_testsamples, list_of_class_names, \
    list_of_train, list_of_valid, list_of_test, data_name, batchSize = data_provide()


    md = loadmodelmodul(modelpath=os.path.join(os.getcwd(), 'model.py'))

    parameters = md.tryparameters[param_id]
    model_id = param_id + 1000  # +1000 so the id is always at least 4 digits
    path_to_mod_folder = create_model_folder(path_to_exp_folder, model_id)

    if start_from_epoch == 0:
        np.random.seed(2004)
        model, trainsettings = md.create(compile=True, **parameters)
        initial_run_hist = None

        save_experiment_code(path_to_mod_folder, model_id, parameters)

        save_model_summary(model, path_to_mod_folder, model_id)

    else:
        model = load_saved_model(path_to_mod_folder, model_id, start_from_epoch, best=False)
        trainsettings = md.create(compile=False, **parameters)
        initial_run_hist = load_run_hist_json(path_to_mod_folder, model_id)

    set_saving_of_last_model(trainsettings, path_to_mod_folder, model_id, model)
    set_saving_of_best_model(trainsettings, path_to_mod_folder, model_id, monitor='val_loss')
    set_returning_running_history(trainsettings, path_to_mod_folder, model_id, initial_run_hist=initial_run_hist)

    if traindata_gen is not None:
        train_start_time = time.time()
        np.random.seed(2004)
        hist = train(model, trainsettings, traindata_gen, num_trainsamples, validdata_gen, num_validsamples,
                     start_from_epoch, batchSize)
        train_end_time = time.time()

        model = load_saved_model(path_to_mod_folder, model_id, epoch=None, best=True)
        vprobs, vclass, vtrues = validate(model, validdata_gen, num_validsamples, batchSize)
    else:
        raise ValueError(
            'The variable train_valid in datagenerator.py function data_provide() must be set to True for now.')

    if testdata_gen is not None:
        tprobs, tclass, ttrues = test(model, testdata_gen, num_testsamples, batchSize)
    else:
        tprobs = tclass = ttrues = None

    if start_from_epoch == 0:
         tvh = hist.history
    else:
         tvh = load_run_hist_json(path_to_mod_folder, model_id)

    trainstats = get_train_stats_dict(tvh, train_start_time, train_end_time, path_to_mod_folder, model_id, data_name)


    save_tvhistory_json(trainstats)

    save_data_info(trainstats, num_trainsamples, num_validsamples, num_testsamples)

    save_model_results(trainstats, list_of_class_names, vclass, vtrues, tclass, ttrues)

    save_train_epoch_log(trainstats)

    save_file_result_log(trainstats, vprobs, vclass, vtrues, tprobs, tclass, ttrues, list_of_valid, list_of_test)

    print('\n{0}{1}\nDONE using GPU core: {2} | Experiment {3} | Model {4}\nParameters: {5}\n{1}{6}'
          .format('\033[94m', '#' * 80, GPU_core, exp_id, model_id, parameters, '\033[0m'))

    print_tvresults(trainstats)

    rename_model_folder(trainstats)


if __name__ == '__main__':
    args = sys.argv
    path_to_exp_folder = args[1]
    exp_id = str(args[2])
    param_id = int(args[3])
    GPU_core = int(args[4])
    start_from_epoch = int(args[5])
    gc.disable()
    run(path_to_exp_folder, exp_id, param_id, GPU_core, start_from_epoch)
