import gc
import os
import random
import yaml
from argparse import ArgumentParser
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow.keras.backend as K 
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import StratifiedKFold
from sklearn.preprocessing import LabelEncoder, StandardScaler

from models import NNBuilder
from models import FitDataSequence, PredictDataSequence


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 create_dataset(train_df, test_df, config):
    x_train_cat, x_train_num = pd.DataFrame(), pd.DataFrame()
    x_test_cat, x_test_num = pd.DataFrame(), pd.DataFrame()

    for col in config['use_cols']['numerical']:
        print('numerical:', col)

        mean = train_df[col].mean()
        x_train_num[col] = train_df[col].fillna(mean).apply(np.log1p)
        x_test_num[col] = test_df[col].fillna(mean).apply(np.log1p)

    for col in config['use_cols']['count']:
        print('count:', col)

        x_train_num[col] = train_df.groupby(col)[col].transform('count').apply(np.log1p)
        x_test_num[col] = test_df.groupby(col)[col].transform('count').apply(np.log1p)

        mean = x_train_num[col].mean()
        x_train_num[col].fillna(mean, inplace=True)
        x_test_num[col].fillna(mean, inplace=True)

    for cols in config['use_cols']['pair_count']:
        col1, col2 = cols['col1'], cols['col2']
        col_name = '{}_{}'.format(col1, col2)
        print('pair count:', col_name)

        x_train_num[col_name] = train_df[[col1, col2]].groupby([col1, col2])[col2].transform('count').apply(np.log1p)
        x_test_num[col_name] = test_df[[col1, col2]].groupby([col1, col2])[col2].transform('count').apply(np.log1p)

        x_train_num[col_name].fillna(0, inplace=True)
        x_test_num[col_name].fillna(0, inplace=True)

    for col in config['use_cols']['categorical']:
        print('categorical:', col)

        train_df[col] = train_df[col].astype('str')
        test_df[col] = test_df[col].astype('str')
    
        le = LabelEncoder()
        le.fit(
            np.unique(train_df[col].unique().tolist() +
                      test_df[col].unique().tolist())
        )

        train_df[col] = le.transform(train_df[col]) + 1
        test_df[col]  = le.transform(test_df[col]) + 1

        agg_tr = (train_df.groupby([col])
                    .aggregate({'MachineIdentifier': 'count'})
                    .reset_index()
                    .rename({'MachineIdentifier': 'Train'}, axis=1))
        agg_te = (test_df.groupby([col])
                    .aggregate({'MachineIdentifier': 'count'})
                    .reset_index()
                    .rename({'MachineIdentifier': 'Test'}, axis=1))

        agg = pd.merge(agg_tr, agg_te, on=col, how='outer').replace(np.nan, 0)
        agg = agg[(agg['Train'] > 1000)].reset_index(drop=True)
        agg['Total'] = agg['Train'] + agg['Test']
        agg = agg[(agg['Train'] / agg['Total'] > 0.2) & (agg['Train'] / agg['Total'] < 0.8)]
        agg[col + 'Copy'] = agg[col]

        x_train_cat[col] = pd.merge(
                               train_df[[col]], 
                               agg[[col, col + 'Copy']], 
                               on=col,
                               how='left'
                           )[col + 'Copy'].replace(np.nan, 0)

        x_test_cat[col] = pd.merge(
                               test_df[[col]], 
                               agg[[col, col + 'Copy']], 
                               on=col,
                               how='left'
                          )[col + 'Copy'].replace(np.nan, 0)

        _, indexer = pd.factorize(x_train_cat[col], sort=True)
        x_train_cat[col] = indexer.get_indexer(x_train_cat[col])
        x_test_cat[col] = indexer.get_indexer(x_test_cat[col])

    del agg, agg_tr, agg_te, le, indexer
    gc.collect()

    scaler = StandardScaler()
    scaler.fit(np.vstack([x_train_num.values, x_test_num.values]))
    x_train_num = scaler.transform(x_train_num)
    x_test_num = scaler.transform(x_test_num)

    return (x_train_cat, x_train_num), (x_test_cat, x_test_num)


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 train_model(x_train_cat, x_train_num, y_train, config):
    model_params = config['model']['model_params']
    train_params = config['model']['train_params']

    kfold = get_kfold(config)

    embed_dims = model_params['embed_dims']
    numerical_size = x_train_num.shape[1]
    epochs = train_params['epochs']
    lr = train_params['lr']
    batch_size = train_params['batch_size']

    result = []

    for i, (train_idx, valid_idx) in enumerate(kfold.split(y_train, y_train)):
        print('Fold', i+1)

        x_train_cat_fold = x_train_cat.iloc[train_idx]
        x_train_num_fold = x_train_num[train_idx]
        x_val_cat_fold = x_train_cat.iloc[valid_idx]
        x_val_num_fold = x_train_num[valid_idx]

        y_train_fold = y_train.iloc[train_idx]
        y_val_fold = y_train.iloc[valid_idx]

        model = NNBuilder(embed_dims, numerical_size, lr).build_model()

        train_data_gen = FitDataSequence(x_train_cat_fold, x_train_num_fold, y_train_fold, batch_size)
        valid_data_gen = FitDataSequence(x_val_cat_fold, x_val_num_fold, y_val_fold, batch_size)

        best_auc = 0.
        patience = 0

        # NOTE: AUC callback maybe needs more RAM
        for epoch in range(epochs):
            print('epoch', epoch + 1)

            fit = model.fit_generator(
                generator=train_data_gen,
                steps_per_epoch=len(train_data_gen),
                epochs=1,
                validation_data=valid_data_gen,
                validation_steps=len(valid_data_gen),
                shuffle=True,
                verbose=1,
            )

            pred = model.predict_generator(
                generator=PredictDataSequence(x_train_cat_fold, x_train_num_fold, batch_size),
                verbose=0,
            )
            train_auc = roc_auc_score(y_train_fold, pred)
            print('Train AUC: {}'.format(train_auc))

            pred = model.predict_generator(
                generator=PredictDataSequence(x_val_cat_fold, x_val_num_fold, batch_size),
                verbose=0,
            )
            val_auc = roc_auc_score(y_val_fold, pred)
            print('Valid AUC: {}'.format(val_auc))

            if val_auc > best_auc:
                best_auc = val_auc
                patience = 0
                model.save_weights('./tmp/model_fold{}_best.h5'.format(i+1))
            else:
                patience += 1

            if patience == 1:
                current_lr = K.get_value(model.optimizer.lr)
                new_lr = max(current_lr * 0.75, 1e-5)
                K.set_value(model.optimizer.lr, new_lr)

            if patience >= 3:
                result.append(best_auc)
                break

    print('CV mean score:', np.mean(result))


def predict_model(x_test_cat, x_test_num, config):
    n_splits = config['cv_strategy']['n_splits']
    embed_dims = config['model']['model_params']['embed_dims']
    numerical_size = x_test_num.shape[1]
    lr = 0

    preds = []

    for i in range(n_splits):
        model = NNBuilder(embed_dims, numerical_size, lr).build_model()
        model.load_weights('./tmp/model_fold{}_best.h5'.format(i+1))

        pred = model.predict_generator(
            generator=PredictDataSequence(x_test_cat_fold, x_test_num_fold, 4096),
            verbose=0,
        )
        preds.append(pred)

    return np.mean(preds, dim=1)


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

    # TODO: ensure reproducibility
    seed = config['model']['train_params']['seed']
    os.environ['PYTHONHASHSHEED'] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    session_conf = tf.ConfigProto(
        intra_op_parallelism_threads=1,
        inter_op_parallelism_threads=1
    )
    tf.set_random_seed(seed)
    sess = tf.Session(graph=tf.get_default_graph(), config=session_conf)
    K.set_session(sess)

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

        print('Load train/test dump files.')
        train_df = joblib.load(train_path)
        test_df = joblib.load(test_path)

        print('Create dataset.')
        x_train, x_test = create_dataset(train_df, test_df, config)
        joblib.dump(x_train, './.cache/train_nn_dataset.pkl')
        joblib.dump(x_test, './.cache/test_nn_dataset.pkl')

        del x_train, x_test
        gc.collect()

    x_train_cat, x_train_num = joblib.load('./.cache/train_nn_dataset.pkl')
    x_test_cat, x_test_num = joblib.load('./.cache/test_nn_dataset.pkl')

    target_col = config['target']
    target_path = get_dataset_filename(config, 'target')
    y_train = joblib.load(target_path)[target_col]

    train_model(x_train_cat, x_train_num, y_train, config)
    preds = predict_model(x_test_cat, x_test_num, config)

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

    submission_path = get_dataset_filename(config, 'sample_submit')
    sub = pd.read_csv(submission_path)
    sub[target_col] = preds
    sub.to_csv('{}/submit_{}.csv.gz'.format(output_dir, basename), index=False)


if __name__ == '__main__':
    main()

