is_debug = 0
is_down = 0
learning_rate = 0.01
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
from sklearn.model_selection import KFold

#========================================================================
# 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
"""
# Basic Args
seed = 1208
set_type = 'all'
fold_n = 5
key, target, ignore_list = MS_utils.get_basic_var()
ignore_list = [key, target, 'country_group', 'down_flg']
comment = sys.argv[1]

# Base
base_path = '../input/base_exclude*'
base_path = '../input/base_Av*'
base_path = '../input/base_group*'
base = utils.read_df_pkl(base_path)[[key, target, 'country_group']]

#========================================================================
# Make Validation
#  vi_col = 'f000_AvSigVersion'
#  ignore_list.append(vi_col)

#  cv = KFold(n_splits=fold_n, shuffle=False, random_state=seed)
#  if is_debug:
#      base_train = base[~base[target].isnull()].head(10000).sort_values(by=vi_col)
#  else:
#      base_train = base[~base[target].isnull()].sort_values(by=vi_col)

#  kfold = list(cv.split(base_train[[key, 'country_group']], base_train[target]))
#  del base_train
#  gc.collect()
#  base = base[[key, target, 'country_group']]

if is_down:
    kfold = MS_utils.get_kfold(base=base[base['down_flg']==0].reset_index(drop=True))
else:
    kfold = MS_utils.get_kfold(base=base)
kfold = zip(*kfold)

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

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

#========================================================================
# 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 = train[~train[col].isnull()][col].mode().values[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])
#========================================================================

#========================================================================
# Down
if is_down:
    train = train[train['down_flg']==0].reset_index(drop=True)
    down_cols = ['down_flg', 'country_group']
    train.drop(down_cols, axis=1, inplace=True)
    test.drop( down_cols, axis=1 , inplace=True)
#========================================================================

Y = train[target]

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

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

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

for model_type in model_type_list:

    if model_type=='lgb':
        params = params_lgb()
        #  params['num_leaves'] = 4
        #  params['num_threads'] = 36
        params['num_threads'] = 60
        params['learning_rate'] = learning_rate
    else:
        params = {}
    logger.info(f"{model_type} Train Start!!")
    for num_fold, (trn_idx, val_idx) in enumerate(kfold):

        #========================================================================
        # 複数スレッドでkfoldを取得する
        get_fold = np.int(sys.argv[3])
        if num_fold!=get_fold:
            continue
        #========================================================================

        x_train, y_train = train[use_cols].iloc[trn_idx, :], Y.iloc[trn_idx]
        x_val, y_val = train[use_cols].iloc[val_idx, :], Y.iloc[val_idx]

        logger.info(f"Fold{num_fold} | 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=x_test
            , params=params
            , seed=seed
            , get_score=metric
        )
        feim_list.append(feim.set_index('feature').rename(columns={'importance':f'imp_{num_fold}'}))

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

    pred_col = 'prediction'
    base[pred_col] = np.hstack((oof_pred, y_test))
    base = base[[key, pred_col]]
    #========================================================================
    # Saving
    utils.to_pkl_gzip(obj=base, path=f'../output/{start_time[4:12]}_stack_{model_type}_FOLD-{get_fold}_feat{len(x_train.columns)}_{comment}')
    #========================================================================
