# mount drive

from google.colab import drive
drive.mount('/content/drive')

# 必要なライブラリのインストール
!pip install -q transformers > /dev/null

# カレントディレクトリを変更
import os
os.chdir('/content/drive/My Drive/Colab Notebooks/probspace/kiva/')
print(os.getcwd())

class Config():
    root_path = './'
    input_path = os.path.join(root_path, 'data')
    output_path = os.path.join(root_path, 'output')
    bert_model_name = 'bert-base-uncased'
    seed = 42
    debug = False

# create dirs

for dir in [Config.output_path]:
    os.makedirs(dir, exist_ok=True)

import datetime
import itertools
import matplotlib.pylab as plt
import numpy as np
import pandas as pd
import random
import pickle
import seaborn as sns
from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import mean_absolute_error, mean_squared_error
from sklearn.model_selection import KFold
import torch
from tqdm import tqdm

# NLP
from sklearn.decomposition import PCA
import transformers
import ssl
ssl._create_default_https_context = ssl._create_unverified_context

# Model
import lightgbm as lgb

pd.set_option('max_columns', None)
pd.options.display.float_format = '{:.5f}'.format

def seed_everything(seed=2021):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    

seed_everything(Config.seed)

train_df = pd.read_csv(os.path.join(Config.input_path, 'train.csv'))
if Config.debug:
    train_df = train_df[:1000]
print(train_df.shape)
display(train_df)

test_df = pd.read_csv(os.path.join(Config.input_path, 'test.csv'))
if Config.debug:
    test_df = test_df[:1000]
print(test_df.shape)
display(test_df)

class BertSequenceVectorizer:
    """
    事前学習済み BERT モデルを使ったテキスト特徴抽出
    """
    def __init__(self, model_name='bert-base-uncased', max_len=128):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.model_name = model_name
        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_name)
        self.model = transformers.AutoModel.from_pretrained(self.model_name)
        self.model = self.model.to(self.device)
        self.max_len = max_len

    def vectorize(self, sentence : str) -> np.array:
        inp = self.tokenizer.encode(sentence)
        len_inp = len(inp)

        if len_inp >= self.max_len:
            inputs = inp[:self.max_len]
            masks = [1] * self.max_len
        else:
            inputs = inp + [0] * (self.max_len - len_inp)
            masks = [1] * len_inp + [0] * (self.max_len - len_inp)

        inputs_tensor = torch.tensor([inputs], dtype=torch.long).to(self.device)
        masks_tensor = torch.tensor([masks], dtype=torch.long).to(self.device)

        output = self.model(inputs_tensor, masks_tensor)
        seq_out = output['last_hidden_state']

        if torch.cuda.is_available():    
            return seq_out[0][0].cpu().detach().numpy()
        else:
            return seq_out[0][0].detach().numpy()

def get_bert_feature(input_df):
    vectorizer = BertSequenceVectorizer(model_name=Config.bert_model_name)
    texts = input_df['DESCRIPTION_TRANSLATED'].fillna('')
    text_vecs = np.array([vectorizer.vectorize(x) for x in texts])
    pca = PCA(n_components=64)
    text_vecs = pca.fit_transform(text_vecs)

    output_df = pd.DataFrame(text_vecs, columns=[f'bert_pca_vecs={i:03}' for i in range(text_vecs.shape[1])])
    output_df.index = input_df.index
    return output_df

train_bert = get_bert_feature(train_df)
test_bert = get_bert_feature(test_df)

model_params = {
    'n_estimators': 1000,
    'objective': 'mae',
    'learning_rate': 0.1,
    'random_state': Config.seed,
    'n_jobs': -1,
}

fit_params = {
    'early_stopping_rounds': 100,
    'verbose': False
}

def make_kf(X, y, n_splits=5):
    kf = KFold(n_splits=n_splits, shuffle=True, random_state=Config.seed)
    return list(kf.split(X))

def train_cv(X, y, model, model_params, fit_params, cv, folds):
    oof = []; va_idxes = []; scores = []; models = {}

    train_x, train_y = X.values, y.values
    fold_idx = cv(train_x, train_y, n_splits=folds)

    for fold, (tr_idx, va_idx) in enumerate(fold_idx):
        tr_x, va_x = train_x[tr_idx], train_x[va_idx]
        tr_y, va_y = train_y[tr_idx], train_y[va_idx]
        va_idxes.append(va_idx)

        est = model(**model_params)

        est.fit(tr_x, np.log1p(tr_y),
               eval_set=[[va_x, np.log1p(va_y)]],
               **fit_params)
        
        model_name = f'LGBM_FOLD{fold}'
        models[model_name] = est

        preds = est.predict(va_x)
        preds = np.expm1(preds)
        oof.append(preds)

        score = mean_absolute_error(va_y, preds)
        scores.append(score)
        print(f'FOLD: {fold}, SCORE: {score:.5f}')

    va_idxes = np.concatenate(va_idxes)
    oof = np.concatenate(oof)
    order = np.argsort(va_idxes)
    oof = oof[order]

    oof_score = mean_absolute_error(train_y, oof)
    print(f'oof score: {oof_score:.5f}\n')
    return oof, models

oof, models = train_cv(train_bert, train_df['LOAN_AMOUNT'],
                        lgb.LGBMModel,
                        model_params,
                        fit_params,
                        cv=make_kf,
                        folds=5)

def predict(test_x, models):
    preds = []
    for i, (name, est) in enumerate(models.items()):
        print(f'{name}')
        preds.append(est.predict(test_x.values))
    preds = np.mean(preds, axis=0)
    preds = np.expm1(preds)

    return preds

preds = predict(test_bert, models)

submit_df = pd.DataFrame({
    'LOAN_ID': test_df['LOAN_ID'],
    'LOAN_AMOUNT': preds.reshape(-1)
})

display(submit_df)

submit_df.to_csv(os.path.join(Config.output_path, f'submission_bert_tutorial.csv'), index=False)