# Library
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import seaborn as sns
from tqdm.auto import tqdm

import warnings
warnings.simplefilter('ignore')

# mount
from google.colab import drive
if not os.path.isdir('/content/drive'):
    drive.mount('/content/drive')

# Config
DRIVE_PATH = "/content/drive/MyDrive/ML/PROBSPACE/pollen_counts"
INPUT = os.path.join(DRIVE_PATH, "input")
OUTPUT = os.path.join(DRIVE_PATH, "output")

TRAIN_FILE = os.path.join(INPUT, "train_v2.csv")
TEST_FILE = os.path.join(INPUT, "test_v2.csv")
SUB_FILE = os.path.join(INPUT, "submission.csv")

exp_name = 'trial_seed_effect'
seed = 42

# plot style
pd.set_option('display.max_rows', 1000)
pd.set_option('display.max_columns', 1000)
plt.rcParams['axes.facecolor'] = 'EEFFFE'

# Data
train = pd.read_csv(TRAIN_FILE)
test = pd.read_csv(TEST_FILE)
sub = pd.read_csv(SUB_FILE)

target_col = ['pollen_utsunomiya', 'pollen_chiba', 'pollen_tokyo']
temp_col = ['temperature_utsunomiya', 'temperature_chiba', 'temperature_tokyo']
windd_col = ['winddirection_utsunomiya', 'winddirection_chiba', 'winddirection_tokyo']
winds_col = ['windspeed_utsunomiya', 'windspeed_chiba', 'windspeed_tokyo']
ppt_col = ['precipitation_utsunomiya', 'precipitation_chiba', 'precipitation_tokyo']

# 降雪かつ他の地域での飛散量も0以下の時0を代入
train.loc[((train['pollen_utsunomiya']==-9998)|(train['pollen_chiba']==-9998)|(train['pollen_tokyo']==-9998))&\
          (((train['pollen_utsunomiya']<=0)&(train['pollen_chiba']<=0)&(train['pollen_tokyo']<=0))), target_col] = 0
train = train[(train['pollen_utsunomiya']>=0)&(train['pollen_chiba']>=0)&(train['pollen_tokyo']>=0)]

# object(欠測) -> float
import lightgbm as lgb
from lightgbm import LGBMRegressor
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer

train_df = train.replace('欠測', np.nan)
lgb_imp = IterativeImputer(
                       estimator=LGBMRegressor(num_boost_round=1000, random_state=seed),
                       max_iter=10, 
                       initial_strategy='mean',
                       imputation_order='ascending',
                       verbose=1,
                       random_state=seed)

train_df = pd.DataFrame(lgb_imp.fit_transform(train_df), columns=train_df.columns)
train_df[['winddirection_chiba', 'winddirection_tokyo']] = train_df[['winddirection_chiba', 'winddirection_tokyo']].round().astype(int)
train_df[['precipitation_tokyo', 'temperature_chiba', 'temperature_tokyo', 'windspeed_chiba', 'windspeed_tokyo']] = train_df[['precipitation_tokyo', 'temperature_chiba', 'temperature_tokyo', 'windspeed_chiba', 'windspeed_tokyo']].round(1)
train_df['datetime'] = train_df['datetime'].astype(int)
train = train_df
train

from sklearn.model_selection import (
    StratifiedKFold, 
    KFold, 
    GroupKFold,
    StratifiedGroupKFold,
)

from sklearn.metrics import mean_absolute_error as mae

import lightgbm as lgb

import os
import random
import tensorflow as tf
from tqdm.notebook import tqdm

import warnings
warnings.filterwarnings('ignore')

# param
seed = 42
plot_mode=False

def set_seed(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)

# LightGBM
class ModelLgb:
    def __init__(self, plot: bool, params: dict):
        self.model = None
        self.plot = plot
        self.params = params

    def fit(self, tr_x, tr_y, va_x=None, va_y=None):
        num_round = 10000
        early_stopping_rounds=50
        
        # validation
        if va_x is not None:
            lgb_train = lgb.Dataset(tr_x.values, tr_y)
            lgb_eval = lgb.Dataset(va_x.values, va_y)
            self.model = lgb.train(self.params, lgb_train, valid_sets=lgb_eval, num_boost_round=num_round, verbose_eval=0,
                                  callbacks=[lgb.early_stopping(stopping_rounds=early_stopping_rounds, verbose=False)]
                                  )
        # No validation
        else:
            lgb_train = lgb.Dataset(tr_x, tr_y)
            self.model = lgb.train(self.params, lgb_train, num_boost_round=100, verbose_eval=0)

        # plot feature importance
        if self.plot:
            f_importance = np.array(self.model.feature_importance())
            df_importance = pd.DataFrame({'feat': tr_x.columns, 'importance': f_importance})
            df_importance = df_importance.sort_values('importance', ascending=True)
            plt.figure(figsize=(8,12))
            plt.barh('feat', 'importance', data=df_importance.iloc[-30:])
            plt.show()   
        
    def predict(self, x):
        pred = self.model.predict(x, num_iteration=self.model.best_iteration)
        return pred

# Ref : https://comp.probspace.com/competitions/pollen_counts/discussions/saru_da_mon-Post5943fd8142f960c070d7
def zero_count(input_df, alpha = 0.05):
    df_count = []
    n_count = 0
    for i in range(len(input_df)):
        if input_df[i] < 0.5:
            n_count += 1
        else:
            n_count = 0
        df_count.append(n_count)
    df_count = np.tanh(np.array(df_count)*alpha)
    return df_count

def run_trial_feat(train, test):
    # 連結して全データに対して処理
    df = pd.concat([train, test]).reset_index(drop=True)

    # 時間特徴
    df['time'] = pd.to_datetime(df.datetime.astype(str).str[:-2])
    df['year'] = df['time'].dt.year
    df['month'] = df['time'].dt.month
    df['hour'] = df.datetime.astype(str).str[-2:].astype(int)

    # 降水量の変換
    for c in ppt_col:
        df[c] = zero_count(df[c])

    # train/testに再分割、欠損処理
    train_df = df[:len(train)]
    test_df = df[len(train):]
    train_df = train_df.dropna().reset_index(drop=True)

    return train_df, test_df

# run
train_df, test_df = run_trial_feat(train, test)

print(train_df.shape)
display(train_df.head(3))
print(test_df.shape)
display(test_df.head(3))

from sklearn.model_selection import train_test_split

def run_trial(test_size=0.25, seed=42, plot_mode=False):
    set_seed(seed)
    vq = {'pollen_utsunomiya':20, 'pollen_chiba':36, 'pollen_tokyo':24}
    params = {
            'boosting':'gbdt',
            'objective':'fair',
            'metric':'fair',
            'seed': 42,
            'verbosity':-1,
            'learning_rate':0.1,
            }

    results = dict()
    score = []
    for tcol in target_col:
        train_tmp = train_df.copy()
        test_tmp = test_df.copy()

        qth = vq[tcol]
        train_tmp = train_tmp[train_tmp[tcol] <= qth].reset_index(drop=True)

        del_columns = target_col+['datetime', 'time']
        train_x = train_tmp.drop(del_columns, axis=1)
        train_y = np.log1p(train_tmp[tcol]/4).values
        test_x = test_tmp.drop(del_columns, axis=1)

        # seed値によって分割されたデータが異なる
        tr_x, va_x, tr_y, va_y = train_test_split(train_x, train_y, test_size=test_size, random_state=seed)

        # training
        model = ModelLgb(plot=plot_mode, params=params)
        model.fit(tr_x, tr_y, va_x, va_y)

        # valid / test predict
        va_pred = model.predict(va_x.values)
        va_pred = np.where(va_pred < 0, 0, va_pred) # post-processing
        test_pred = model.predict(test_x.values)
        test_pred = np.where(test_pred < 0, 0, test_pred) # post-processing

        # valid loss
        va_loss = mae(va_y, va_pred)

        # plot valid / pred
        if plot_mode:
            plt.figure(figsize=(50,5))
            plt.plot(va_y, label='original', linestyle='-')
            plt.plot(va_pred, label='pred', linestyle='-')
            plt.title(f'{tcol} : {va_loss}')
            plt.legend()
            plt.show()

        # save per target
        results[tcol] = np.expm1(test_pred)
        score.append(va_loss)

    return results, np.array(score).mean()

%%time
seed_array = range(9)
test_rate = 0.33

result_list = list()
loss_list = list()
for s in tqdm(seed_array):
    # run
    result, loss = run_trial(test_rate, s)
    result_list.append(result)
    loss_list.append(loss)
    # submit
    results_df = pd.DataFrame(result)
    results_df = results_df.round()*4
    sub[target_col] = results_df
    sub.to_csv(os.path.join(OUTPUT, f'{exp_name}{s}.csv'), index=False)

# Blending
for s in seed_array:
    blend_list = [pd.DataFrame(result) for result in result_list]
result_blend = np.array(blend_list).mean(axis=0)
results_ens_df = pd.DataFrame(result_blend, columns=target_col)
resultsr_ens_df = results_ens_df.round()*4
sub[target_col] = results_df
sub.to_csv(os.path.join(OUTPUT, f'{exp_name}_blend.csv'), index=False)

LB_loss_list = [
    12.34328, 
    12.68159, 
    12.58209, 
    12.82090, 
    12.78109, 
    12.96020, 
    12.32338, 
    12.58209, 
    12.76119
    ]

fig = plt.figure()
ax1 = fig.add_subplot(111)
ln1=ax1.plot(pd.DataFrame(np.array(loss_list), columns=['Score']), label='Validation Score', color='blue')
ax2 = ax1.twinx()
ln2=ax2.plot(pd.DataFrame(np.array(LB_loss_list), columns=['PublicLB Score']), label='PublicLB Score', color='red')

h1, l1 = ax1.get_legend_handles_labels()
h2, l2 = ax2.get_legend_handles_labels()
ax1.legend(h1+h2, l1+l2, loc='upper left')

ax1.grid(True)
ax1.set_xlabel('Seed')
ax1.set_ylabel('Validation Score')
ax2.set_ylabel('PublicLB Score')