"""
## Load library
"""
import numpy as np 
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import lightgbm as lgb
from category_encoders.ordinal import OrdinalEncoder
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import f1_score
"""
## Load data
"""
train = pd.read_csv('/kaggle/input/kakr-4th-competition/train.csv')
test = pd.read_csv('/kaggle/input/kakr-4th-competition/test.csv')
sample_submission = pd.read_csv('/kaggle/input/kakr-4th-competition/sample_submission.csv')
train.head()
"""
## EDA & pre-process
"""
sns.distplot(train.age)
sns.distplot(np.log1p(train.age))
train.age = np.log1p(train.age)
test.age = np.log1p(test.age)
sns.distplot(train.hours_per_week)
sns.distplot(np.log1p(train.hours_per_week))
train.hours_per_week = np.log1p(train.hours_per_week)
test.hours_per_week = np.log1p(test.hours_per_week)
sns.distplot(train.fnlwgt)
sns.distplot(np.log1p(train.fnlwgt))
train.fnlwgt = np.log1p(train.fnlwgt)
test.fnlwgt = np.log1p(test.fnlwgt)
train['capital_d'] = train.capital_gain - train.capital_loss
test['capital_d'] = test.capital_gain - test.capital_loss
train.capital_gain = np.log1p(train.capital_gain)
test.capital_gain = np.log1p(test.capital_gain)
train.capital_loss = np.log1p(train.capital_loss)
test.capital_loss = np.log1p(test.capital_loss)
train.capital_d = np.log1p(train.capital_d)
test.capital_d = np.log1p(test.capital_d)
"""
## Label data convert
"""
target = train['income'] != '<=50K'
train.drop(['income'], axis=1, inplace=True)
train_le = train
test_le = test
# LE_encoder = OrdinalEncoder(list(train.columns))
# train_le = LE_encoder.fit_transform(train, target)
# test_le = LE_encoder.transform(test)
cat_features = ['workclass', 'education', 'marital_status', 'occupation', 'relationship', 'race', 'sex', 'native_country']
for f in cat_features:
    train_le[f] = train_le[f].astype('category')
    test_le[f] = test_le[f].astype('category')
skf = StratifiedKFold(n_splits=5)
# train_le['fold'] = 0
fold_id_list = list()
for fold_id, (trn_id, val_id) in enumerate(skf.split(X=train_le.index, y=target)):
    fold_id_list.append(val_id)
params = {
        'objective':'binary',
        "boosting": "gbdt",
        'num_leaves': 100,
        'max_depth': 6,#16
        'learning_rate': 0.1,#0.1
#         'min_data_in_leaf': 16, 
#         'min_child_samples': 30,
#         'min_child_weight': 0.5,
#         'min_split_gain': 0.1,
#         "feature_fraction": 0.9,
#         "bagging_fraction": 0.9,
#         "bagging_freq": 2,
        "bagging_seed": 42,
        "colsample_bytree": 0.8,
        "metric": 'lgb_f1_score',
        "lambda_l1": 0.2,
        "lambda_l2": 0.5,
        'verbose':-1
    }
def custom_round(predict, threshold):
    data = predict.copy()
    try:
        data.loc[data>=threshold] = 1
        data.loc[data<threshold] = 0
    except:
        data[data>=threshold] = 1
        data[data<threshold] = 0
    return data

def lgb_f1_score(y_hat, data):
    y_true = data.get_label()
#     y_hat = custom_round(y_hat, np.quantile(y_hat, 0.7))
    y_hat = custom_round(y_hat, 0.5) #  
    return 'f1', f1_score(y_true, y_hat), True
C_FOLD = 0
result_pred = list()
for C_FOLD in range(5):
    valid_fold_id = fold_id_list[C_FOLD]
    
    train_dataset = train_le.loc[train_le.index.isin(valid_fold_id) == False]
    valid_dataset = train_le.loc[train_le.index.isin(valid_fold_id)]
    
    train_target = target[train_le.index.isin(valid_fold_id) == False]
    valid_target = target[train_le.index.isin(valid_fold_id)]
    
    # LightGBM
    trn_df = lgb.Dataset(train_dataset, label=train_target)
    val_df = lgb.Dataset(valid_dataset, label=valid_target)
    
    model = lgb.train(params, trn_df, 1000, 
                            valid_sets = [trn_df, val_df], 
                            early_stopping_rounds = 100, 
                            verbose_eval=100, 
                            feval=lgb_f1_score)
    pred = model.predict(test_le)
    result_pred.append(pred)
    
    # RandomForest
#     RF_clf = RandomForestClassifier()
#     RF_clf.fit(train_dataset, train_target)
    
#     pred = RF_clf.predict_proba(test_le)
#     result_pred.append(pred[:, 1])
result_pred = np.stack(result_pred)
# result_pred = custom_round(result_pred.mean(axis=0), np.quantile(result_pred.mean(axis=0), 0.7))
result_pred = (result_pred.mean(axis=0) >= 0.5)
sample_submission['prediction'] = result_pred.astype(int)
sample_submission.to_csv('submission.csv', index=False)