import gc
import time
import warnings
import yaml
from argparse import ArgumentParser
from pathlib import Path

import joblib
import logzero
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from joblib import Parallel, delayed
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import StratifiedKFold

from _model_map import models
from _feature_map import features
from utils import fast_auc, load_train_dataset, load_test_dataset, load_dump, timer


def get_dataset_filename(config, dataset_type):
    assert (type(dataset_type) == str and
            dataset_type in ('train', 'target', 'test', 'sample_submit'))

    dataset_dir = Path(config['dataset']['input_directory'])
    filename = Path(config['dataset']['files'][dataset_type])
    return dataset_dir / filename


def setup_logger(output_dir, logname):
    return logzero.setup_logger(
        logfile=f'{output_dir}/{logname}.log',
        level=10,
        formatter=None,
    )


def create_features(config, train_df, test_df, kfold):
    def _process(feature):
        feat = get_feature(feature)

        encode_type = None
        if feature in config['features']['count_encode']:
            encode_type = 'count'
        elif feature in config['features']['category_encode']:
            encode_type = 'category'
        elif feature in config['features']['rank_encode']:
            encode_type = 'rank'
        elif feature in config['features']['target_encode']:
            encode_type = 'target'
            feat.prepare_target(config['target'])

        if feat.create(train_df, test_df, kfold, encode_type):
            feat.save()

    with timer('Create features...'):
        Parallel(n_jobs=1)([delayed(_process)(f) for f in config['features']['all']])


def get_feature(feature):
    return features[feature]()


def get_kfold(config):
    cv = config['cv_strategy']
    if cv['method'] == 'StratifiedKFold':
        kfold = StratifiedKFold(n_splits=cv['n_splits'],
                                shuffle=True,
                                random_state=cv['seed'])
    else:
        raise NotImplementedError

    return kfold


def extract_use_features(config):
    use_features = []

    for f in config['features']['all']:
        if f in config['features']['count_encode']:
            use_features.append(f'cnt_{f}')
        elif f in config['features']['category_encode']:
            use_features.append(f'lbl_{f}')
        elif f in config['features']['rank_encode']:
            use_features.append(f'rank_{f}')
        elif f in config['features']['target_encode']:
            cv_n_splits = config['cv_strategy']['n_splits']
            cv_seed = config['cv_strategy']['seed']
            use_features.append(f'te_{f}_split{cv_n_splits}_seed{cv_seed}')
        else:
            use_features.append(f)

    return use_features


def train_model(x_train, y_train, kfold, config, logger):
    feature_importances = pd.DataFrame()
    scores = []
    clfs = []

    for i, (train_index, valid_index) in enumerate(kfold.split(x_train, y_train)):
        start_time = time.time()
        logger.info(f'Fold {i+1}')

        x_train_fold = x_train.iloc[train_index]
        y_train_fold = y_train.iloc[train_index]

        x_valid_fold = x_train.iloc[valid_index]
        y_valid_fold = y_train.iloc[valid_index]

        clf = models[config['model']['name']]()
        y_pred_valid = clf.train_and_validate(x_train_fold,
                                              x_valid_fold,
                                              y_train_fold,
                                              y_valid_fold,
                                              params=config['model'],
                                              logger=logger)

        scores.append(fast_auc(y_valid_fold, y_pred_valid))
        logger.info(f'Fold roc_auc: {roc_auc_score(y_valid_fold.values, y_pred_valid)}')

        fold_importance = pd.DataFrame()
        fold_importance['feature'] = x_train_fold.columns
        fold_importance['gain'] = clf.model.feature_importance(importance_type='gain')
        fold_importance['fold'] = i + 1
        feature_importances = pd.concat([feature_importances, fold_importance],
                                        axis=0, sort=False)

        clfs.append(clf)

        # resume training when error occurs
        joblib.dump(clf, f'tmp/lgb_fold{i+1}.pkl')
        joblib.dump(feature_importances, f'tmp/importances_fold{i+1}.pkl')

        elapsed = time.time() - start_time
        logger.info(f'Fold {i+1} done in {elapsed:.0f} s\n')

    feature_importances['gain'] /= kfold.n_splits
    logger.info(f'CV mean score: {np.mean(scores):.4f}')

    return clfs, feature_importances


def predict_model(clfs, use_features, config):
    n_splits = len(clfs)
    predictions = []

    for i, clf in enumerate(clfs):
        x_test = load_test_dataset(use_features)

        for f in config['features']['target_encode']:
            if f is None:
                pass

            col_name = get_feature(f).name
            drop_cols = [f'te_{col_name}_fold{k+1}' for k in range(n_splits) if i != k]
            x_test.drop(drop_cols, axis=1, inplace=True)
            x_test.rename(columns={f'te_{col_name}_fold{i+1}': f'te_{col_name}'},
                          inplace=True)

        print(f'predict: fold {i+1}/{n_splits}')
        y_pred = clf.predict(x_test)
        predictions.append(y_pred)

    return np.mean(predictions, axis=0)


def save_feature_importances(importances, basename, output_dir):
    mean_gain = importances[['feature', 'gain']].groupby('feature').mean()
    importances['mean_gain'] = importances['feature'].map(mean_gain['gain'])

    with warnings.catch_warnings():
        warnings.simplefilter('ignore')

        plt.figure(figsize=(20, 30))
        sns.barplot(x='gain', y='feature',
                    data=importances.sort_values('mean_gain', ascending=False))
        plt.title('LightGBM feature importances')
        plt.tight_layout()
        plt.savefig(f'{output_dir}/importances_{basename}.png')


def main():
    parser = ArgumentParser()
    parser.add_argument('--config', default='./configs/lgb_template.yaml')
    parser.add_argument('--create-features', action='store_true')
    options = parser.parse_args()
    config = yaml.safe_load(open(options.config))

    kfold = get_kfold(config)

    if options.create_features:
        train_path = get_dataset_filename(config, 'train')
        test_path = get_dataset_filename(config, 'test')

        with timer('Load train/test dump files'):
            train_df = load_dump(train_path)
            test_df = load_dump(test_path)

        create_features(config, train_df, test_df, kfold)

        del train_df, test_df
        gc.collect()

    target_col = config['target']
    target_path = get_dataset_filename(config, 'target')
    use_features = extract_use_features(config)
    x_train = load_train_dataset(use_features)
    y_train = load_dump(target_path)[target_col]

    output_dir = config['dataset']['output_directory']
    basename = Path(options.config).stem
    logger = setup_logger(output_dir, basename)

    clfs, importances = train_model(x_train, y_train, kfold, config, logger)

    save_feature_importances(importances, basename, output_dir)

    del x_train, y_train, importances
    gc.collect()

    pred = predict_model(clfs, use_features, config)

    print('Creating a submission csv file...')
    submission_path = get_dataset_filename(config, 'sample_submit')
    submission = pd.read_csv(submission_path)
    submission[target_col] = pred
    submission.to_csv(f'{output_dir}/submit_{basename}.csv.gz', index=False)
    print('Done.')

if __name__ == '__main__':
    main()

