is_debug = 0
get_model = 1
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 = 4
key, raw_target, ignore_list = MS_utils.get_basic_var()
ignore_list = [key, raw_target]
comment = sys.argv[1]

# Base
vi_col = 'f000_AvSigVersion'
base_path = '../input/base_exclude*'
base_path = '../input/base_Av*'
#  base_path = '../input/base_group*'
base = utils.read_df_pkl(base_path)

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

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

#  train, test = MS_utils.get_dataset(base=base)
test, train = MS_utils.get_dataset(base=base)

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

#========================================================================
# 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])
#========================================================================

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

target_list = [
    #  'f000_SmartScreen'
    #  'f000_AvSigVersion'
    #  ,'f000_EngineVersion'
    #  ,'f000_Census_OSVersion'
    'f000_AVProductStatesIdentifier'
    #  ,'f000_CountryIdentifier'
    #  ,'f000_AVProductsInstalled'
    ]
#  [::-1]

for target in target_list:

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

    target_mode = train[~train[target].isnull()][target].mode()[0]
    Y = train[target].fillna(target_mode)

    ignore_list = [key, raw_target, target, 'country_group', pred_col]
    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())

    if model_type=='lgb':
        params = params_lgb()
        #  params['num_leaves'] = 4
        params['num_threads'] = 35
        try:
            #  params['learning_rate'] = float(sys.argv[1])
            params['learning_rate'] = 0.05
        except IndexError:
            pass
    else:
        params = {}

    logger.info(f"{model_type} Train Start!! Target: {target}")

    for num_fold, (trn_idx, val_idx) in enumerate(kfold):
        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(
        score, tmp_oof, tmp_pred, feim, model = ml_utils.Regressor(
            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
            , get_model=get_model
        )
        if get_model:
            utils.to_pkl_gzip(obj=model, path=f'../model/{start_time[4:11]}_{comment}_{target}_{model_type}_fold{num_fold}_feat{len(use_cols)}_{metric}-{score}')
            del model
            gc.collect()

        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

    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)

    cv_score = np.mean(score_list)
    logger.info(f'''
    #========================================================================
    # Model : {model_type}
    # Target: {target}
    # CV    : {cv_score}
    #========================================================================''')

    y_test /= (num_fold+1)

    train[pred_col] = oof_pred
    test[pred_col] = y_test
    stack_cols = [key, target, pred_col]

    df_stack = pd.concat([train[stack_cols], test[stack_cols]], ignore_index=True, axis=0)

    #========================================================================
    # Saving
    feim.to_csv(f'../valid/{start_time[4:12]}_{model_type}_SET-{target}_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-{target}_feat{len(x_train.columns)}_{comment}_CV{str(cv_score)[:7]}_LB')
    #========================================================================
