import os
import time
import sys
import argparse
import yaml
import shutil
import threading
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import KFold
from joblib import Parallel, delayed

import ATE
from ..data_utils import load_batches, encode_data_frame, x_y_split
from ..plot_utils import set_plotting_style
from ..plot_reg_performance import plot_reg_performance
from ..model_loader import get_model_factory, load_model_from_file
from ..hyperopt.grid_search import grid_search
from ..hyperopt.bayesian_optimization import bayesian_optimization
from ..metric_loader import get_metric_factory


def main():
    '''Main command line entry point. Trains model with given parameters.'''

    random_state = 1

    # parse args
    parser = argparse.ArgumentParser(
        description='Search TBR model hyperparameter space')
    parser.add_argument('in_dir', type=str,
                        help='directory containing input batches')
    parser.add_argument('batch_low', type=int,
                        help='start batch index (inclusive)')
    parser.add_argument('batch_high', type=int,
                        help='end batch index (exclusive)')
    parser.add_argument('in_search_file', type=str,
                        help='input search.csv file generated by tbr_search')
    parser.add_argument('out_dir', type=str,
                        help='output directory')
    parser.add_argument('model_type', type=str,
                        help='model type (needed for initialization)')

    parser.add_argument('--feature-def', type=str,
                        help='path to file with line-separated (encoded) whitelisted feature names')
    parser.add_argument('--k-folds', type=int, default=5,
                        help='k for k-fold cross-validation')
    parser.add_argument('--score', type=str, default='r2',
                        help='metric for model quality evaluation, supported values: "r2" (default), "mae", "adjusted_r2", "std_error"')
    parser.add_argument('--score-ascending', default=False, action='store_true',
                        help='True if metric should be sorted in ascending (from the smallest to the largest)')
    parser.add_argument('--n-best-models', type=int, default=10,
                        help='how many best models to use')
    parser.add_argument('--save-trained-models', default=False, action='store_true',
                        help='save model files for each checkpoint')
    parser.add_argument('--save-plots', default=False, action='store_true',
                        help='produce performance plots for each fold')
    parser.add_argument('--n-jobs', type=int, default=1,
                        help='number of parallel jobs')
    parser.add_argument('--set-size', type=int, action='append')
    args = parser.parse_args()

    set_plotting_style()

    # load data
    df = load_batches(args.in_dir, (args.batch_low, args.batch_high))
    df_enc = encode_data_frame(df, ATE.Domain())
    X, y_multiple = x_y_split(df_enc)
    y = y_multiple['tbr']

    if args.feature_def is not None:
        with open(args.feature_def, 'r') as f:
            included_features = [line.strip() for line in f.readlines()
                                 if len(line) > 0]
            X = X[included_features].copy()

    X = X.sort_index(axis=1)
    print(f'Features in order are: {list(X.columns)}')

    all_metrics = [init_metric()
                   for init_metric in get_metric_factory().values()]

    search_results = pd.read_csv(args.in_search_file)
    search_results.sort_values(
        by=['mean_metric_%s' % args.score], inplace=True, ascending=args.score_ascending)
    search_results = search_results.iloc[0:args.n_best_models].copy()

    model_creator = get_model_factory()[args.model_type]

    # prepare output
    out_file = os.path.join(args.out_dir, 'benchmark.csv')

    extra_columns = ['train_size', 'test_size', 'actual_idx']
    for some_metric in all_metrics:
        extra_columns += ['metric_%s%d' %
                          (some_metric.id, i) for i in range(args.k_folds)]
        extra_columns.append('mean_metric_%s' % some_metric.id)
    extra_columns += ['mean_time_train', 'mean_time_pred']
    extra_columns += ['time_train%d' % i for i in range(args.k_folds)]
    extra_columns += ['time_pred%d' % i for i in range(args.k_folds)]

    data = {column: [] for column in extra_columns}

    rows = list(search_results.to_dict(orient='index').items())

    kfold = KFold(n_splits=args.k_folds, shuffle=True,
                  random_state=random_state)

    for total_size in args.set_size:
        print(f'Retraining models on set of size {total_size}')

        # prepare train, test folds
        X_restricted = X.iloc[0:total_size]
        y_restricted = y.iloc[0:total_size]
        kfold_indices = [indices for indices
                         in kfold.split(X_restricted, y_restricted)]

        data = benchmark_models(data, X_restricted, y_restricted, model_creator,
                                kfold_indices, rows, extra_columns, all_metrics, args.n_best_models,
                                args.save_trained_models, args.save_plots, args.out_dir, args.n_jobs)

        benchmark = pd.DataFrame(data=data)
        print(f'Writing {benchmark.shape[0]} results to {out_file}')
        benchmark.to_csv(out_file)

    print('Done.')


def benchmark_models(data, X, y, model_creator, kfold_indices, rows, extra_columns, all_metrics,
                     n_best_models, keep_trained_models, save_plots, out_dir, n_parallel):
    train_size = int(np.median([len(train_index)
                                for train_index, test_index in kfold_indices]))
    test_size = int(np.median([len(test_index)
                               for train_index, test_index in kfold_indices]))

    # make sure all model arguments are recorded
    first_model_args = drop_non_model_args(rows[0][1])
    for arg_name in first_model_args:
        if arg_name not in data.keys() and arg_name not in ['out']:
            data[arg_name] = []

    results = Parallel(n_jobs=n_parallel)(delayed(single_model_job)(X, y, model_creator, kfold_indices, rows, model_idx, extra_columns,
                                                                    all_metrics, keep_trained_models, save_plots, out_dir, train_size, test_size)
                                          for model_idx in range(n_best_models))

    for model_args, extra_values in results:
        data['train_size'].append(train_size)
        data['test_size'].append(test_size)

        for arg_name in data.keys():
            if arg_name in model_args:
                data[arg_name].append(model_args[arg_name])

        for column, value in extra_values.items():
            if column not in ['train_size', 'test_size']:
                data[column].append(value)

    return data


def drop_non_model_args(model_args):
    all_keys = list(model_args.keys())
    for key in all_keys:
        if key.startswith('Unnamed') or key.startswith('mean_') or key.startswith('time_') \
                or key.startswith('metric_') or key.startswith('score'):
            del model_args[key]

    return model_args


def single_model_job(X, y, model_creator, kfold_indices, rows, model_idx, extra_columns,
                     all_metrics, keep_trained_models, save_plots, out_dir, train_size, test_size):
    actual_model_idx, model_args = rows[model_idx]

    model_args = drop_non_model_args(model_args)
    model_dir = os.path.join(
        out_dir, f'model{actual_model_idx}_train{train_size}_test{test_size}')

    # set model to produce output
    if keep_trained_models:
        model_args['out'] = os.path.join(model_dir, 'fold%d')

    # prepare output directory
    if os.path.exists(model_dir):
        print('WARNING: path "%s" exists, deleting previous contents' % model_dir)
        shutil.rmtree(model_dir)
    os.makedirs(model_dir)

    return benchmark_single_model(model_dir, model_creator, model_args, X, y,
                                  kfold_indices, all_metrics, extra_columns, actual_model_idx, save_plots)


def benchmark_single_model(model_dir, model_creator, model_args, X, y, kfold_indices, all_metrics, extra_columns, actual_idx, save_plots):
    extra_values = {column: [] for column in extra_columns}
    fold_idx = 0
    k_folds = len(kfold_indices)
    all_failed = True
    for train_index, test_index in kfold_indices:
        print(f'Starting fold {fold_idx+1} of {k_folds}.')
        X_train, X_test = X.iloc[train_index], X.iloc[test_index]
        y_train, y_test = y.iloc[train_index], y.iloc[test_index]

        fold_args = dict(model_args)
        if 'out' in fold_args:
            fold_args['out'] = fold_args['out'] % fold_idx

        try:
            model = model_creator(arg_dict=fold_args)
            train_time = train(model, X_train, y_train)
            evaluations, pred_time = test(
                model, X_test, y_test, all_metrics)

            for metric_id, value in evaluations.items():
                extra_values['metric_%s%d' % (metric_id, fold_idx)] = value

            extra_values['time_train%d' % fold_idx] = train_time
            extra_values['time_pred%d' % fold_idx] = pred_time

            if save_plots:
                plot_perf_path = os.path.join(model_dir, 'fold%d' % fold_idx)
                plot(plot_perf_path, model, X_test, y_test)

            all_failed = False
        except Exception as e:
            print(
                f'WARNING: Fold {fold_idx+1} of {k_folds} failed with error: {e}')

        print(f'Fold {fold_idx+1} of {k_folds} done.')
        fold_idx += 1

    if all_failed:
        return np.nan * np.ones((k_folds, 1)), np.nan, extra_values

    for some_metric in all_metrics:
        extra_values['mean_metric_%s' % some_metric.id] = \
            np.mean(np.array([extra_values['metric_%s%d' % (some_metric.id, i)]
                              for i in range(k_folds)]))
    extra_values['mean_time_train'] = \
        np.mean(np.array([extra_values['time_train%d' % i]
                          for i in range(k_folds)]))
    extra_values['mean_time_pred'] = \
        np.mean(np.array([extra_values['time_pred%d' % i]
                          for i in range(k_folds)]))
    extra_values['actual_idx'] = actual_idx

    return model_args, extra_values


def train(model, X_train, y_train):
    print(f'Training regressor on set of size {X_train.shape[0]}')
    tic = time.time()
    model.train(X_train.to_numpy(), y_train.to_numpy())
    toc = time.time()
    return (toc - tic) / X_train.shape[0]


def test(model, X_test, y_test, metrics):
    print(f'Testing regressor on set of size {X_test.shape[0]}')
    tic = time.time()
    y_pred = model.predict(X_test.to_numpy())
    toc = time.time()

    y_test = y_test.to_numpy()

    evaluations = {}
    for metric in metrics:
        evaluation = metric.evaluate(X_test, y_test, y_pred)
        print(
            f'Evaluation on test set of size {X_test.shape[0]} gives {metric.name} result: {evaluation}')
        evaluations[metric.id] = evaluation
    return evaluations, (toc - tic) / X_test.shape[0]


def plot(save_plot_path, model, X_test, y_test):
    df = X_test.copy()
    df.insert(0, 'tbr', -1.)
    df.insert(0, 'tbr_pred', -1.)
    df['tbr'] = y_test
    df['tbr_pred'] = model.predict(X_test)

    fig, ax = plot_reg_performance(df, density_bins=80)
    plt.tight_layout()

    plt.savefig('%s.png' % save_plot_path)
    plt.savefig('%s.pdf' % save_plot_path)
    plt.close()


if __name__ == '__main__':
    main()
