"""
#### This is fork of https://www.kaggle.com/code1110/janestreet-faster-inference-by-xgb-with-treelite beautifull notebook on how to make faster prediction with xgb!! <br>

#### I'm using PurgedGroupTimeSeriesSplit for validation with multitarget.
"""
"""
# Install treelite
"""
!pip --quiet install ../input/treelite/treelite-0.93-py3-none-manylinux2010_x86_64.whl
!pip --quiet install ../input/treelite/treelite_runtime-0.93-py3-none-manylinux2010_x86_64.whl
"""
# Imports 🛬
"""
import numpy as np
import pandas as pd

import os, sys
import gc
import math
import random
import pathlib
from tqdm import tqdm
from typing import List, NoReturn, Union, Tuple, Optional, Text, Generic, Callable, Dict
from sklearn.preprocessing import MinMaxScaler, StandardScaler, QuantileTransformer
from sklearn.decomposition import PCA
from sklearn import linear_model
import operator
import xgboost as xgb
import lightgbm as lgb
from tqdm import tqdm

# treelite
import treelite
import treelite_runtime 

# visualize
import matplotlib.pyplot as plt
import matplotlib.style as style
import seaborn as sns
from matplotlib_venn import venn2
from matplotlib import pyplot
from matplotlib.ticker import ScalarFormatter
sns.set_context("talk")
style.use('fivethirtyeight')
pd.options.display.max_columns = None

import warnings
warnings.filterwarnings('ignore')
"""
# PurgedGroupTimeSeriesSplit
"""
from sklearn.metrics import roc_auc_score

from sklearn.model_selection import KFold
from sklearn.model_selection._split import _BaseKFold, indexable, _num_samples
from sklearn.utils.validation import _deprecate_positional_args
class PurgedGroupTimeSeriesSplit(_BaseKFold):
    """Time Series cross-validator variant with non-overlapping groups.
    Allows for a gap in groups to avoid potentially leaking info from
    train into test if the model has windowed or lag features.
    Provides train/test indices to split time series data samples
    that are observed at fixed time intervals according to a
    third-party provided group.
    In each split, test indices must be higher than before, and thus shuffling
    in cross validator is inappropriate.
    This cross-validation object is a variation of :class:`KFold`.
    In the kth split, it returns first k folds as train set and the
    (k+1)th fold as test set.
    The same group will not appear in two different folds (the number of
    distinct groups has to be at least equal to the number of folds).
    Note that unlike standard cross-validation methods, successive
    training sets are supersets of those that come before them.
    Read more in the :ref:`User Guide <cross_validation>`.
    Parameters
    ----------
    n_splits : int, default=5
        Number of splits. Must be at least 2.
    max_train_group_size : int, default=Inf
        Maximum group size for a single training set.
    group_gap : int, default=None
        Gap between train and test
    max_test_group_size : int, default=Inf
        We discard this number of groups from the end of each train split
    """

    @_deprecate_positional_args
    def __init__(self,
                 n_splits=5,
                 *,
                 max_train_group_size=np.inf,
                 max_test_group_size=np.inf,
                 group_gap=None,
                 verbose=False
                 ):
        super().__init__(n_splits, shuffle=False, random_state=None)
        self.max_train_group_size = max_train_group_size
        self.group_gap = group_gap
        self.max_test_group_size = max_test_group_size
        self.verbose = verbose

    def split(self, X, y=None, groups=None):
        """Generate indices to split data into training and test set.
        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Training data, where n_samples is the number of samples
            and n_features is the number of features.
        y : array-like of shape (n_samples,)
            Always ignored, exists for compatibility.
        groups : array-like of shape (n_samples,)
            Group labels for the samples used while splitting the dataset into
            train/test set.
        Yields
        ------
        train : ndarray
            The training set indices for that split.
        test : ndarray
            The testing set indices for that split.
        """
        if groups is None:
            raise ValueError(
                "The 'groups' parameter should not be None")
        X, y, groups = indexable(X, y, groups)
        n_samples = _num_samples(X)
        n_splits = self.n_splits
        group_gap = self.group_gap
        max_test_group_size = self.max_test_group_size
        max_train_group_size = self.max_train_group_size
        n_folds = n_splits + 1
        group_dict = {}
        u, ind = np.unique(groups, return_index=True)
        unique_groups = u[np.argsort(ind)]
        n_samples = _num_samples(X)
        n_groups = _num_samples(unique_groups)
        for idx in np.arange(n_samples):
            if (groups[idx] in group_dict):
                group_dict[groups[idx]].append(idx)
            else:
                group_dict[groups[idx]] = [idx]
        if n_folds > n_groups:
            raise ValueError(
                ("Cannot have number of folds={0} greater than"
                 " the number of groups={1}").format(n_folds,
                                                     n_groups))

        group_test_size = min(n_groups // n_folds, max_test_group_size)
        group_test_starts = range(n_groups - n_splits * group_test_size,
                                  n_groups, group_test_size)
        for group_test_start in group_test_starts:
            train_array = []
            test_array = []

            group_st = max(0, group_test_start - group_gap - max_train_group_size)
            for train_group_idx in unique_groups[group_st:(group_test_start - group_gap)]:
                train_array_tmp = group_dict[train_group_idx]
                
                train_array = np.sort(np.unique(
                                      np.concatenate((train_array,
                                                      train_array_tmp)),
                                      axis=None), axis=None)

            train_end = train_array.size
 
            for test_group_idx in unique_groups[group_test_start:
                                                group_test_start +
                                                group_test_size]:
                test_array_tmp = group_dict[test_group_idx]
                test_array = np.sort(np.unique(
                                              np.concatenate((test_array,
                                                              test_array_tmp)),
                                     axis=None), axis=None)

            test_array  = test_array[group_gap:]
            
            
            if self.verbose > 0:
                    pass
                    
            yield [int(i) for i in train_array], [int(i) for i in test_array]
"""
# Config 🔧
"""
SEED = 42 # Happy new year!
# INPUT_DIR = '../input/jane-street-market-prediction/'
START_DATE = 85
INPUT_DIR = '../input/janestreet-save-as-feather/'
TRADING_THRESHOLD = 0.502 # 0 ~ 1: The smaller, the more aggressive
"""
# Load Data and Data Preprocessing

"""
os.listdir(INPUT_DIR)
%%time

def load_data(input_dir=INPUT_DIR):
    train = pd.read_feather(pathlib.Path(input_dir + 'train.feather'))
    #features = pd.read_feather(pathlib.Path(input_dir + 'features.feather'))
    #example_test = pd.read_feather(pathlib.Path(input_dir + 'example_test.feather'))
    #ss = pd.read_feather(pathlib.Path(input_dir + 'example_sample_submission.feather'))
    return train
train = load_data(INPUT_DIR)
# reduce train
train = train.query(f'date > {START_DATE}')
train.fillna(train.mean(),inplace=True)
train = train[train['weight'] != 0]
# features
features = train.columns[train.columns.str.startswith('feature')].values.tolist()
print('{} features used'.format(len(features)))
# target
train['action'] = (train['resp'] > 0).astype('int')
f_mean = np.mean(train[features[1:]].values,axis=0)
"""
# Model💪
"""
params = {'n_estimators': 473, 'max_depth': 7, 'min_child_weight': 6, 
 'learning_rate': 0.015944928866056352, 'subsample': 0.608128483148888, 
 'gamma': 0, 'colsample_bytree': 0.643875232059528,'objective':'binary:logistic',
'eval_metric': 'auc','tree_method': 'gpu_hist', 'random_state': 42,} 

params_1 = {'n_estimators': 494, 'max_depth': 8, 'min_child_weight': 6, 'learning_rate': 0.009624384025871735, 
            'subsample': 0.8328412036014541, 'gamma': 0, 'colsample_bytree': 0.715303237773365,
           'objective':'binary:logistic', 'eval_metric': 'auc','tree_method': 'gpu_hist', 'random_state': 42,}
training = True
import pickle

if training:
    import time
    import gc
    resp_cols = ['resp_1', 'resp_2', 'resp_3', 'resp', 'resp_4']
    X = train[features].values
    #y = train['action'].values
    y = np.stack([(train[c] > 0).astype('int') for c in resp_cols]).T #Multitarget
    groups = train['date'].values
    models = []
    scores = []

    cv = PurgedGroupTimeSeriesSplit(
        n_splits=4,
        group_gap=20,
    )
    for t in tqdm(range(y.shape[1])):
        yy = y[:,t]
        for i, (train_index, valid_index) in enumerate(cv.split(
                X,
                yy,
                groups=groups)):
            print(f'Target {t} Fold {i} started at {time.ctime()}')
            X_train, X_valid = X[train_index], X[valid_index]
            y_train, y_valid = yy[train_index], yy[valid_index]
            model = xgb.XGBClassifier(**params_1, n_jobs = -1)
            model.fit(X_train, y_train, 
                    eval_set=[(X_valid, y_valid)], eval_metric='auc',
                    verbose=100, callbacks = [xgb.callback.EarlyStopping(rounds=300,save_best=True)])
            pred = model.predict(X_valid)
            score = roc_auc_score(y_valid,pred)
            model.save_model(f'my_model_{t}_{i}.model')
            pickle.dump(model, open(f'my_model_{t}_{i}.pkl', "wb"))
            models.append(model)
            scores.append(score)
            del score, model
        print(scores)
        del X_train, X_valid, y_train, y_valid
        rubbish = gc.collect()
"""
# Compile with Treelite
Simply follow the tutorial: https://treelite.readthedocs.io/en/latest/tutorials/first.html
"""
# pass to treelite
if training:
    model_0 = treelite.Model.load('my_model_0_3.model', model_format='xgboost')
    model_1 = treelite.Model.load('my_model_1_3.model', model_format='xgboost')
    model_2 = treelite.Model.load('my_model_2_3.model', model_format='xgboost')
    model_3 = treelite.Model.load('my_model_3_3.model', model_format='xgboost')
    model_4 = treelite.Model.load('my_model_4_3.model', model_format='xgboost')
if training:
    m = [model_0,model_1,model_2,model_3,model_4]
    for j,i in enumerate(m):
        toolchain = 'gcc'
        i.export_lib(toolchain=toolchain, libpath=f'./mymodel_{j}.so',
                     params={'parallel_comp': 32}, verbose=True)
# predictor from treelite
if training:
    predictor_0 = treelite_runtime.Predictor(f'./mymodel_{0}.so', verbose=True)
    predictor_1 = treelite_runtime.Predictor(f'./mymodel_{1}.so', verbose=True)
    predictor_2 = treelite_runtime.Predictor(f'./mymodel_{2}.so', verbose=True)
    predictor_3 = treelite_runtime.Predictor(f'./mymodel_{3}.so', verbose=True)
    predictor_4 = treelite_runtime.Predictor(f'./mymodel_{4}.so', verbose=True)
"""
# 🏹Submission🎯
"""
import janestreet
env = janestreet.make_env() # initialize the environment
iter_test = env.iter_test() # an iterator which loops over the test set
f = np.median 
index_features = [n for n in range(1,(len(features) + 1))]
for (test_df, pred_df) in tqdm(iter_test):
    
    if test_df['weight'].item() > 0:
        
        
        x_tt = test_df.values[0][index_features].reshape(1,-1)
            
        if np.isnan(x_tt[:, 1:].sum()):
            x_tt[:, 1:] = np.nan_to_num(x_tt[:, 1:]) + np.isnan(x_tt[:, 1:]) * f_mean
        
        # inference with treelite
        batch = treelite_runtime.Batch.from_npy2d(x_tt)
        pred_0 = predictor_0.predict(batch)
        pred_1 = predictor_1.predict(batch)
        pred_2 = predictor_2.predict(batch)
        pred_3 = predictor_3.predict(batch)
        pred_4 = predictor_4.predict(batch)
        
        # Prediction
        pred = np.stack([pred_0,pred_1,pred_2,pred_3,pred_4],axis=0).T
        pred = f(pred)
        pred_df.action = int(pred >= TRADING_THRESHOLD)
        
    else:
         pred_df['action'].values[0] = 0
    env.predict(pred_df)
"""
# If this notebook helped Please do Upvote 💘💙✅

## Part 2 getting ready with feature selction!!! 
"""