is_debug = 0
import os
import re
import gc
import sys
import glob
import pandas as pd
import numpy as np
import datetime
from sklearn.preprocessing import LabelEncoder

#========================================================================
# original library 
HOME = os.path.expanduser('~')
sys.path.append(f"{HOME}/kaggle/data_analysis/library/")
sys.path.append(f"{HOME}/kaggle/data_analysis/model/")
import MS_utils
from params_MS import params_lgb
import utils, ml_utils
from utils import logger_func
try:
    if not logger:
        logger=logger_func()
except NameError:
    logger=logger_func()
#========================================================================

"""
argv[1]: comment
argv[2]: feature_key
argv[3]: group
"""

# Columns
base_path = '../input/base_exclude*'
base_path = '../input/base_Av*'
key, target, ignore_list = MS_utils.get_basic_var()
ignore_list = [key, target, 'country_group', 'down_flg']
base = utils.read_df_pkl(base_path)[[key, target, 'country_group']]

# Basic Args
seed = 1208
set_type = 'all'

comment = sys.argv[1]

if sys.argv[2].count('f'):
    train, test = MS_utils.get_feature_set(feat_key=sys.argv[2], base_path=base_path)
else:
    train, test = MS_utils.get_dataset(base=base)
print(train.shape, test.shape)

if is_debug:
    train = train.head(10000)
    test = test.head(5000)

#========================================================================
# Categorical Encode
cat_cols = utils.get_categorical_features(df=train, ignore_list=ignore_list)
print(f"Categorical: {cat_cols}")

#Fit LabelEncoder
for col in cat_cols:
    # 最も頻度の多いカテゴリでimpute
    max_freq = list(train[col].value_counts().index)[0]
    train[col].fillna(max_freq, inplace=True)
    test[col].fillna(max_freq, inplace=True)
    le = LabelEncoder().fit(pd.concat([train[col], test[col]], axis=0).value_counts().index.tolist())
    train[col] = le.transform(train[col])
    test[col]  = le.transform(test[col])
#========================================================================

model_type_list = ['lgb', 'rmf', 'lgr']
model_type_list = ['lgb']
model_type = 'lgb'
metric = 'auc'

feim_list = []
score_list = []
stack_list = []
y_test = np.zeros(len(test))

use_cols = [col for col in train.columns if col not in ignore_list]
x_test = test[use_cols]
start_time = "{0:%Y%m%d_%H%M%S}".format(datetime.datetime.now())

base_sum = 0
group_list = sorted(train['country_group'].unique().tolist())

#  for model_type in model_type_list:
for group in group_list:

    set_type = f"group{group}"
    tmp_train = train[train['country_group']==group].sample(frac=0.4)
    tmp_test = test[test['country_group']==group].head(100)
    Y = tmp_train[target]
    oof_pred = np.zeros(len(tmp_train))
    y_test = np.zeros(len(tmp_test))

    kfold = MS_utils.get_kfold(base=tmp_train[[key, target]], is_group=False, fold_n=3 )

    if model_type=='lgb':
        params = params_lgb()
        try:
            #  params['learning_rate'] = float(sys.argv[1])
            params['learning_rate'] = 0.05
            params['learning_rate'] = 0.1
        except IndexError:
            pass
        params['num_leaves'] = 2**6-1
        params['num_threads'] = 36
        params['num_threads'] = -1
    else:
        params = {}
    logger.info(f"{model_type} Train Start!!")
    for fold_n, (trn_idx, val_idx) in enumerate(kfold):
        x_train, y_train = tmp_train[use_cols].iloc[trn_idx, :], Y.iloc[trn_idx]
        x_val, y_val = tmp_train[use_cols].iloc[val_idx, :], Y.iloc[val_idx]

        logger.info(f"Fold{fold_n} | Train:{x_train.shape} | Valid:{x_val.shape}")

        score, tmp_oof, tmp_pred, feim, _ = ml_utils.Classifier(
            model_type=model_type
            , x_train=x_train
            , y_train=y_train
            , x_val=x_val
            , y_val=y_val
            , x_test=tmp_test
            , params=params
            , seed=seed
            , get_score=metric
        )
        feim_list.append(feim.set_index('feature').rename(columns={'importance':f'imp_{fold_n}'}))

        logger.info(f"Fold{fold_n} CV: {score}")
        score_list.append(score)
        oof_pred[val_idx] = tmp_oof
        y_test += tmp_pred

    feim = pd.concat(feim_list, axis=1)
    feim_cols = [col for col in feim.columns if col.count('imp_')]
    feim['importance'] = feim[feim_cols].mean(axis=1)
    feim.drop(feim_cols, axis=1, inplace=True)
    feim.sort_values(by='importance', ascending=False, inplace=True)
    feim['rank'] = np.arange(len(feim)) + 1

    cv_score = np.mean(score_list)
    logger.info(f'''
    #========================================================================
    # CV: {cv_score} | Group: {group} | Model: {model_type}
    #========================================================================''')

    y_test /= (fold_n+1)

    pred_col = 'prediction'
    tmp_train[pred_col] = oof_pred
    tmp_test[pred_col] = y_test
    stack_cols = [key, target, pred_col]
    tmp_stack = pd.concat([tmp_train, tmp_test], axis=0)[stack_cols]

    if base_sum==0:
        base_sum = tmp_stack[pred_col].sum()
    else:
        tmp_sum = tmp_stack[pred_col].sum()
        ratio = base_sum/tmp_sum
        tmp_stack[pred_col] = tmp_stack[pred_col] * ratio

    stack_list.append(tmp_stack)
    feim.to_csv(f'../valid/{start_time[4:12]}_{model_type}_SET-{set_type}_feat{len(x_train.columns)}_{comment}_CV{str(cv_score)[:7]}_LB.csv', index=True)

del train, test
gc.collect()

df_stack = pd.concat(stack_list, axis=0)
base = utils.read_df_pkl(base_path)[[key, target]]
base_train = base[~base[target].isnull()]
base_test = base[base[target].isnull()]
print(f"Before Stack Shape: {df_stack.shape}")
train = base_train.drop(target, axis=1).merge(df_stack[~df_stack[target].isnull()], on=key, how='inner')
test = base_test.drop(target, axis=1).merge(df_stack[df_stack[target].isnull()], on=key, how='inner')
df_stack = pd.concat([train, test], ignore_index=True, axis=0)
print(f"After Stack Shape: {df_stack.shape}")

y_train = train[target].values
y_pred = train[pred_col].values
from sklearn.metrics import roc_auc_score
cv_score = roc_auc_score(y_train, y_pred)
logger.info(f'''
#========================================================================
# CV: {cv_score}
#========================================================================''')

#========================================================================
# Saving
feim.to_csv(f'../valid/{start_time[4:12]}_{model_type}_SET-{set_type}_feat{len(x_train.columns)}_{comment}_CV{str(cv_score)[:7]}_LB.csv', index=True)
utils.to_pkl_gzip(obj=df_stack, path=f'../stack/{start_time[4:12]}_{model_type}_SET-{set_type}_feat{len(x_train.columns)}_{comment}_CV{str(cv_score)[:7]}_LB')

submit = pd.read_csv('../input/sample_submission.csv').set_index(key)
submit[target] = test[pred_col].values
submit.to_csv(f'../submit/{start_time[4:12]}_{model_type}_SET-{set_type}_feat{len(x_train.columns)}_{comment}_CV{str(cv_score)[:7]}_LB.csv', index=True)
#========================================================================
