import warnings
warnings.filterwarnings('ignore')

import os
import re
import gc
import sys
import time
import glob
import math
import json
import pickle
import joblib
import random
import string
import psutil

from pathlib import Path
from datetime import datetime
from tqdm.autonotebook import tqdm
from contextlib import contextmanager

import pandas as pd
import numpy as np
pd.set_option('display.max_columns', None)
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import StratifiedKFold, GroupKFold, KFold
from texthero import preprocessing as hero_prep
import texthero as hero

import matplotlib.pyplot as plt
plt.style.use('ggplot')
import japanize_matplotlib
import seaborn as sns

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from transformers import BertTokenizer, BertModel

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

INPUT_DIR = '/mnt/work/data/ProbSpace'
OUTPUT_DIR = f'../data/interim/{datetime.now():%Y_%m_%d}/'
MODEL_DIR = f'../models/{datetime.now():%Y_%m_%d}/'
# DATA_DIR = '/mnt/work/data/ProbSpace/shimizu/'

os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(MODEL_DIR, exist_ok=True)

print(f'保存先: {OUTPUT_DIR}')
print(f'モデル保存先: {MODEL_DIR}')

@contextmanager
def timer(name:str, slack:bool=True):
    t0 = time.time()
    p = psutil.Process(os.getpid())
    m0 = p.memory_info()[0] / 2. ** 30
    print(f'<< {name} >> Start')
    yield
    
    m1 = p.memory_info()[0] / 2. ** 30
    delta = m1 - m0
    sign = '+' if delta >= 0 else '-'
    delta = math.fabs(delta)
    
    print(f"<< {name} >> {m1:.1f}GB({sign}{delta:.1f}GB):{time.time() - t0:.1f}sec", file=sys.stderr)

class CONFIG:
    SEED = 42
    NFOLDS = 7
    DEVICE = ('cuda' if torch.cuda.is_available() else 'cpu')

def Load_rowjson(file_path:str, cols:list, train:bool=True) -> pd.DataFrame:
    
    df = pd.DataFrame(columns = cols)
    
    with open(file_path, 'r') as f:
        if train==True:
            for i, json_line in tqdm(enumerate(f)):
                data = json.loads(json_line)
                if 'cites' in data:
                    for col in cols:
                        df.at[i, col] = data[col]

                else:
                    continue

        else:
            cols.remove('cites')
            for i, json_line in tqdm(enumerate(f)):
                data = json.loads(json_line)
                for col in cols:
                    df.at[i, col] = data[col]

    df = df.reset_index(drop=True)
    
    return df

with timer('Load Data'):
    
    
    df_cols = ["title", "abstract", "doi_cites", "cites"]
    
    train_df = Load_rowjson(os.path.join(INPUT_DIR, 'train_data.json'), df_cols, train=True)
    test_df = pd.read_json(os.path.join(INPUT_DIR, 'test_data.json'), lines=True)
    test_df = test_df[["title", "abstract", "doi_cites"]]
    
    print('train_df', train_df.shape)
    print('test_df', test_df.shape)

# doi_citesは対数変換
train_df['cites'] = np.log1p(train_df['cites'].astype('float'))
train_df['doi_cites'] = np.log1p(train_df['doi_cites'].astype('float'))
test_df['doi_cites'] = np.log1p(test_df['doi_cites'].astype('float'))

with timer('Data Merge'):
    
    # 後で分離しやすいように、flag追加
    train_df['flag'] = 0
    test_df['flag'] = 1

    whole_df = pd.concat([train_df,
                          test_df],axis=0).reset_index(drop=True)

def title_embedding(df:pd.DataFrame, pipeline:list, col:str='title') -> np.array:

    title_transformer_mean = []
    title_transformer_max = []
    

    ## 正規表現
    tmp = pd.DataFrame(hero.clean(df[col], pipeline))

    # BERT Model
    ## GPU使用した方が速い
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f'device: {DEVICE}')
    model = BertModel.from_pretrained('bert-base-uncased')
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model.to(DEVICE)
    
    for idx, rows in tqdm(tmp.iterrows(), total=len(tmp)):
        
        title = rows[col]
        
        # Transformer
        input_ids = torch.tensor(tokenizer.encode(title, add_special_tokens=True)).unsqueeze(0)
        outputs = model(input_ids.to(DEVICE))
        
        emb_transformer = outputs[0].cpu().detach().numpy()
        emb_transformer_mean = np.mean(emb_transformer, axis=1)[0, :]
        emb_transformer_max = np.max(emb_transformer, axis=1)[0, :]
        
        title_transformer_mean.append(emb_transformer_mean)
        title_transformer_max.append(emb_transformer_max)
        
        
    title_transformer_mean = np.array(title_transformer_mean)
    title_transformer_max = np.array(title_transformer_max)
    
    print(title_transformer_mean.shape, title_transformer_max.shape)
    
    title_feat = np.concatenate([title_transformer_mean,
                                 title_transformer_max],
                                axis=1)

    del title_transformer_mean, title_transformer_max
    _ = gc.collect()
    
    return title_feat

title_pipeline = [
    hero_prep.remove_diacritics,
    hero_prep.lowercase,
    hero_prep.remove_stopwords,
    hero_prep.tokenize
]


with timer('Title Embedding'):
    print('---- Train ----')
    train_title = title_embedding(train_df, 
                                  title_pipeline,
                                  'title')
    print('---- Test ----')
    test_title = title_embedding(test_df,
                                 title_pipeline,
                                 'title')

def abstract_embedding(df:pd.DataFrame, pipe_list:list, col:str='abstract') -> np.array:

    abstract_transformer_mean = []
    abstract_transformer_max = []
    
    ## 正規表現
    tmp = pd.DataFrame(hero.clean(df[col], pipe_list))
    
    # BERT Model
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f'device: {DEVICE}')
    model = BertModel.from_pretrained('bert-base-uncased')
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model.to(DEVICE)
    
    # 特徴抽出
    for idx, rows in tqdm(tmp.iterrows(), total=len(tmp)):
        
        abstract = rows[col]

        # Transformer
        input_ids = torch.tensor(tokenizer.encode(abstract, add_special_tokens=True)).unsqueeze(0)
        outputs = model(input_ids.to(DEVICE))
        emb_transformer = outputs[0].cpu().detach().numpy()
        emb_transformer_mean = np.mean(emb_transformer, axis=1)[0, :]
        emb_transformer_max = np.max(emb_transformer, axis=1)[0, :]
        
        abstract_transformer_mean.append(emb_transformer_mean)
        abstract_transformer_max.append(emb_transformer_max)
        
        
    abstract_transformer_mean = np.array(abstract_transformer_mean)
    abstract_transformer_max = np.array(abstract_transformer_max)
    
    print(abstract_transformer_mean.shape, 
          abstract_transformer_max.shape)
    
    abstract_feat = np.concatenate([abstract_transformer_mean,
                                    abstract_transformer_max],
                                   axis=1)

    del abstract_transformer_mean, abstract_transformer_max
    _ = gc.collect()
    
    return abstract_feat

abstract_pipeline = [
    hero_prep.remove_diacritics,
    hero_prep.lowercase,
    hero_prep.remove_stopwords,
    hero_prep.tokenize
]


with timer('abstract Embedding'):
    print('---- Train ----')
    train_abstract = abstract_embedding(train_df, 
                                        abstract_pipeline,
                                        'abstract')
    
    print('---- Test ----')
    test_abstract = abstract_embedding(test_df,
                                       abstract_pipeline,
                                       'abstract')
    
    print(train_abstract.shape, test_abstract.shape)

train_doi = train_df['doi_cites'].astype('float').values
test_doi = test_df['doi_cites'].astype('float').values
target = train_df['cites'].astype('float').values

train_doi = np.reshape(train_doi, (-1, 1))
test_doi = np.reshape(test_doi, (-1, 1))
target = np.reshape(target, (-1, 1))

print('----- Train -----')
print(f'title:{train_title.shape} | abstract:{train_abstract.shape} | doi:{train_doi.shape} | target:{target.shape}')
print('----- Test -----')
print(f'title:{test_title.shape} | abstract:{test_abstract.shape} | doi:{test_doi.shape}')

from sklearn.decomposition import PCA

def detect_component(train_array:np.array, test_array:np.array, th:float=0.9) -> int:
    
    whole_array = np.concatenate([train_array, test_array], axis=0)
    
    pca = PCA(n_components=whole_array.shape[1])
    pca.fit(whole_array)
    
    plt.figure(figsize=(20, 5))
    plt.bar([n for n in range(1, len(pca.explained_variance_ratio_)+1)], pca.explained_variance_ratio_)
    plt.show()
    
    tmp = pd.DataFrame(pca.explained_variance_ratio_, columns=['寄与率'])
    tmp = tmp.cumsum()
    tmp = tmp[tmp['寄与率']>th]
    
    print(f'累積寄与率 {th*100:.1f}% を超える次元数 {tmp.index.tolist()[0]}')

    return tmp.index.tolist()[0]

title_component = detect_component(train_title, test_title, 0.9)
abs_component = detect_component(train_abstract, test_abstract, 0.9)

def pca_trans(train:np.array, test:np.array, comp:int) -> (np.array, np.array):
    
    pca = PCA(n_components=comp, random_state=CONFIG.SEED)
    pca.fit(np.concatenate([train, test], axis=0))
    
    train_pca = pca.transform(train)
    test_pca = pca.transform(test)
    
    return train_pca, test_pca

with timer('PCA Transform'):
    train_title_pca, test_title_pca = pca_trans(train_title, test_title, title_component)
    train_abs_pca, test_abs_pca = pca_trans(train_abstract, test_abstract, abs_component)
    
    print(train_title_pca.shape, test_title_pca.shape)
    print(train_abs_pca.shape, test_abs_pca.shape)

class TrainDataset:
    
    def __init__(self, features_1, features_2, target_doi, target_cites):
        
        self.features_1 = features_1
        self.features_2 = features_2
        self.target_doi = target_doi
        self.target_cites = target_cites
        
    def __len__(self):
        return self.features_1.shape[0]
    
    def __getitem__(self, idx):
        
        dct = {
            'x_1' : torch.tensor(self.features_1[idx, :],
                                 dtype=torch.float),
            'x_2' : torch.tensor(self.features_2[idx, :],
                                 dtype=torch.float),
            'doi' : torch.tensor(self.target_doi[idx, :],
                                 dtype=torch.float),
            'cites' : torch.tensor(self.target_cites[idx, :],
                                   dtype=torch.float),
        }
        
        return dct

class TestDataset:
    
    def __init__(self, features_1, features_2, target_doi):
        
        self.features_1 = features_1
        self.features_2 = features_2
        self.target_doi = target_doi
        
    def __len__(self):
        return self.features_1.shape[0]
    
    def __getitem__(self, idx):
        
        dct = {
            'x_1' : torch.tensor(self.features_1[idx, :],
                                 dtype=torch.float),
            'x_2' : torch.tensor(self.features_2[idx, :],
                                 dtype=torch.float),
            'doi' : torch.tensor(self.target_doi[idx, :],
                                 dtype=torch.float),
        }
        
        return dct

def train_func(model, optimizer, scheduler, loss_fn, dataloader, device, epoch):
    
    model.train()
    final_loss = 0
    
    total_steps = EPOCHS * len(dataloader)
    
    for step, data in enumerate(dataloader):
        
        optimizer.zero_grad()
        
        title = data['x_1'].to(device)
        abstract = data['x_2'].to(device)
        doi = data['doi'].to(device)
        targets = data['cites'].to(device)

        output_cites = model(title, abstract, doi)
        
        loss = loss_fn(output_cites, targets)
        
        loss.backward()
        optimizer.step()
        scheduler.step()
        
        final_loss += loss.item()
        
    final_loss /= len(dataloader)
    
    return final_loss

def valid_func(model, loss_fn, dataloader, device):
    
    model.eval()
    final_loss = 0
    valid_cites_preds = []
    
    for step, data in enumerate(dataloader):
        
        title = data['x_1'].to(device)
        abstract = data['x_2'].to(device)
        doi = data['doi'].to(device)
        targets = data['cites'].to(device)

        output_cites = model(title, abstract, doi)
        
        loss = loss_fn(output_cites, targets)
        final_loss += loss.item()
        
        valid_cites_preds.append(output_cites.detach().cpu().numpy())
        
    final_loss /= len(dataloader)
    valid_cites_preds = np.concatenate(valid_cites_preds)
    
    return final_loss, valid_cites_preds

def test_func(model, dataloader, device):
    
    model.eval()
    preds = []
    
    
    for data in dataloader:
        
        title = data['x_1'].to(device)
        abstract = data['x_2'].to(device)
        doi = data['doi'].to(device)
        
        with torch.no_grad():
            outputs = model(title, abstract, doi)
        
        preds.append(outputs.detach().cpu().numpy())
                
    preds = np.concatenate(preds)
    
    return preds

class TwoHeadModel(nn.Module):
    
    def __init__(self, num_features_1, num_features_2, num_doi, num_cites, hidden_size):
        
        
        super(TwoHeadModel, self).__init__()
        
        self.batch_norm_1 = nn.BatchNorm1d(num_features_1)
        self.dense_1 = nn.utils.weight_norm(nn.Linear(num_features_1, 
                                                      hidden_size))
        self.activation_1 = torch.nn.PReLU(num_parameters=hidden_size, 
                                           init=1.0)
        
        self.batch_norm_2 = nn.BatchNorm1d(num_features_2)
        self.dense_2 = nn.utils.weight_norm(nn.Linear(num_features_2, 
                                                      hidden_size))
        self.activation_2 = torch.nn.PReLU(num_parameters=hidden_size, init=1.0)
        
        self.batch_norm_3 = nn.BatchNorm1d(hidden_size*2 + num_doi)
        self.dropout_3 = nn.Dropout(0.3)
        self.dense_3 = nn.utils.weight_norm(nn.Linear(hidden_size*2 + num_doi,
                                                      num_cites))
        
        
    def forward(self, x1, x2, x3):
        
        x1 = self.batch_norm_1(x1)
        x1 = self.activation_1(self.dense_1(x1))
        
        x2 = self.batch_norm_2(x2)
        x2 = self.activation_2(self.dense_2(x2))
        
        x = torch.cat([x1, x2, x3], dim=1)
        
        x = self.batch_norm_3(x)
        x = self.dropout_3(x)
        x = self.dense_3(x)
    
        return x

## RMSLEで学習させる
def RMSELoss(yhat,y):
    return torch.sqrt(torch.mean((yhat-y)**2))

## Cross Validationはdoi_citesの整数部分でStratifiedKFold

DOI_INT = pd.Series(train_df['cites'].astype('float')).astype(int)
plt.hist(DOI_INT, bins=DOI_INT.nunique())
plt.show()

DEVICE = ('cuda' if torch.cuda.is_available() else 'cpu')

EPOCHS = 20
BATCH_SIZE = 64
LEARNING_RATE = 1e-3
WEIGHT_DECAY = 1e-5
NFOLDS = 7


num_features_1 = train_title_pca.shape[1]
num_features_2 = train_abs_pca.shape[1]
num_doi = train_doi.shape[1]
num_cites = target.shape[1]

print(f'num_features_1: {num_features_1} | num_features_2: {num_features_2}')
print(f'num_doi: {num_doi} | num_cites: {num_cites}')

hidden_size = 128

CV = StratifiedKFold(n_splits=CONFIG.NFOLDS, 
                     shuffle=True,
                     random_state=CONFIG.SEED)

oof = np.zeros((train_df.shape[0]))
preds = np.zeros((test_df.shape[0]))

for fold, (tr, te) in enumerate(CV.split(train_df, DOI_INT)):
    
    print('★'*40)
    print(f'Fold: {fold+1}')
    
    X_title_tr = train_title_pca[tr]
    X_title_te = train_title_pca[te]
    
    X_abstract_tr = train_abs_pca[tr]
    X_abstract_te = train_abs_pca[te]
    
    X_doi_tr = train_doi[tr]
    X_doi_te = train_doi[te]
    
    y_tr = target[tr]
    y_te = target[te]

    
    train_dataset = TrainDataset(X_title_tr, X_abstract_tr, X_doi_tr, y_tr)
    valid_dataset = TrainDataset(X_title_te, X_abstract_te, X_doi_te, y_te)
    
    trainloader = torch.utils.data.DataLoader(train_dataset, 
                                              batch_size=BATCH_SIZE, 
                                              shuffle=True)
    
    validloader = torch.utils.data.DataLoader(valid_dataset, 
                                              batch_size=BATCH_SIZE, 
                                              shuffle=False)
    
    model = TwoHeadModel(
        num_features_1 = num_features_1,
        num_features_2 = num_features_2,
        num_doi = num_doi,
        num_cites = num_cites,
        hidden_size = hidden_size,
    )
    
    model.to(DEVICE)
    optimizer = torch.optim.Adam(model.parameters(),
                                 lr=LEARNING_RATE,
                                 weight_decay=WEIGHT_DECAY)
    
    scheduler = optim.lr_scheduler.OneCycleLR(optimizer=optimizer,
                                              pct_start=0.1, 
                                              div_factor=1e3, 
                                              max_lr=1e-2, 
                                              epochs=EPOCHS, 
                                              steps_per_epoch=len(trainloader))
    
    loss_fn = RMSELoss
    best_loss = np.inf
    
    
    for epoch in range(EPOCHS):
        
        start_time = time.time()
        train_loss = train_func(model,
                                optimizer,
                                scheduler, 
                                loss_fn, 
                                trainloader,
                                DEVICE, 
                                epoch)
        
        
        valid_loss, valid_cites_preds = valid_func(model,
                                                   loss_fn,
                                                   validloader,
                                                   DEVICE)
        
            
        end_time = time.time()
        print(f"FOLD: {fold+1} | EPOCH:{epoch+1:02d} | train_loss:{train_loss:.6f} | valid_loss:{valid_loss:.6f} | time:{end_time-start_time:.1f}s ")

        if valid_loss < best_loss:
            best_loss = valid_loss
            oof[te] = valid_cites_preds[:, 0]
            torch.save(model.state_dict(),
                       f"{MODEL_DIR}SimpleMLP_{fold+1}.pth")
            
        else:
            continue
            
    
    testdataset = TestDataset(test_title_pca, test_abs_pca, test_doi)
    testloader = torch.utils.data.DataLoader(testdataset,
                                             batch_size=BATCH_SIZE,
                                             shuffle=False)
    
    model = TwoHeadModel(
        num_features_1 = num_features_1,
        num_features_2 = num_features_2,
        num_doi = num_doi,
        num_cites = num_cites,
        hidden_size = hidden_size,
    )
    
    model.load_state_dict(torch.load(f"{MODEL_DIR}SimpleMLP_{fold+1}.pth"))
    model.to(DEVICE)
    
    preds += test_func(model, testloader, DEVICE)[:, 0]/NFOLDS

rmse = np.sqrt(mean_squared_error(target, oof))
print(f'CV : {rmse:.6f} ')

fig, ax = plt.subplots(figsize=(8, 5))
sns.histplot(target[:, 0], label='正解値', kde=True, color='orange',
             stat="density", common_norm=False, alpha=0.3)
sns.histplot(oof, label='OOF', kde=True,
             stat="density", common_norm=False, alpha=0.3)
ax.legend()
ax.grid()
plt.show()

### 予測値の統計量を確認
pd.DataFrame(np.expm1(preds)).describe().T

def make_oof_sub(oof:np.array, pred:np.array, name:str):
    
    oof_res = np.expm1(oof)
    oof_res = pd.DataFrame(oof_res, columns=['OOF'])
    oof_res.to_csv(f'{OUTPUT_DIR}{name}_oof.csv', index=False)
    print(f'{OUTPUT_DIR}{name}_oof.csv is saved.')
    
    preds_res = np.expm1(pred)
    test_df = pd.read_json(os.path.join(INPUT_DIR, 'test_data.json'), lines=True)
    preds_res = pd.concat([pd.DataFrame({ 'id': test_df['id'] }), 
                           pd.DataFrame({ 'cites': preds_res })], 
                          axis=1)
    preds_res.to_csv(f'{OUTPUT_DIR}{name}_sub.csv', index=False)
    print(f'{OUTPUT_DIR}{name}_sub.csv is saved.')

with timer('make oof sub'):
    make_oof_sub(oof, preds, 'Simple_NN')