
# coding: utf-8

# In[1]:


import os
import gc
import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder

data_path = './data'


# In[3]:


train = pd.read_csv(os.path.join(data_path, 'security_train.csv'))
train_data = train[['file_id','label']].drop_duplicates()


# # Uni-Gram Features

# In[4]:


### aggregate over api.
# File_id (Api): count,nunique
agg_func = ['count','nunique'] 
for func in agg_func:
    tmp = train.groupby(['file_id'])['api'].agg({'fileid_api_' + func: func}).reset_index() 
    train_data = pd.merge(train_data, tmp, how='left', on='file_id')
    
### aggregate over tid
# File_id (Tid): count,nunique,max,min,quantile(20,40,50,60,80),std,range
agg_func = ['count','nunique','max','min','median','std'] 
for func in agg_func:
    tmp = train.groupby(['file_id'])['tid'].agg({'fileid_tid_' + func: func}).reset_index() 
    train_data = pd.merge(train_data,tmp,how='left', on='file_id') 

secs = [0.2,0.4,0.6,0.8]
for sec in secs: 
    train_data['fileid_tid_quantile_' + str(sec * 100)] = train.groupby(['file_id'])['tid'].quantile(sec).values
 
train_data['fileid_tid_range'] = train.groupby(['file_id'])['tid'].quantile(0.975).values - train.groupby(['file_id'])['tid'].quantile(0.0125).values

### aggregate over index
# File_id (Index): count,nunique,max,min,quantile(20,40,50,60,80),std,range
index_opt = ['count','nunique','max','min','median','std'] 
for opt in index_opt:
    print(opt)
    tmp = train.groupby(['file_id'])['index'].agg({'fileid_index_' + opt: opt}).reset_index() 
    train_data = pd.merge(train_data,tmp,how='left', on='file_id') 
    
secs = [0.2,0.4,0.6,0.8]
for sec in secs: 
    train_data['fileid_index_quantile_' + str(sec * 100)] = train.groupby(['file_id'])['index'].quantile(sec).values
 
train_data['fileid_index_range'] = train.groupby(['file_id'])['index'].quantile(0.975).values - train.groupby(['file_id'])['index'].quantile(0.0125).values


# In[5]:


def groupby_pivot_features(data_merge, data_orig , groupby_features, col1 = None, col2 = None, agg_func = None):
    for func in agg_func:
        train_split = data_orig.groupby(['file_id',col1])[col2].agg({'fileid_' + col1 + '_'+col2+'_'+ str(func):func}).reset_index() 
        
        train_split_ =  pd.pivot_table(train_split, values = 'fileid_' + col1 + '_'+col2+'_'+ str(func), index=['file_id'],columns=[col1])
        new_cols = [ 'fileid_' + col1 + '_'+col2+  '_' + func + '_' + str(col) for col in train_split_.columns]
        
        groupby_features.append(new_cols)
        train_split_.columns = new_cols 

        train_split_.reset_index(inplace = True)
        
        data_merge = pd.merge(data_merge, train_split_, how='left', on='file_id') 
    return data_merge, groupby_features 


# In[6]:


### 
# File_id + Api (tid): count,nunique
groupby_features = []
agg_func = ['count', 'nunique']
train_data_,groupby_features = groupby_pivot_features(train_data, train, groupby_features, 
                                                      col1 = 'api', col2 = 'tid', agg_func = agg_func)

###
# File_id + Api (index): nunique, max, min, median, std
agg_func = ['nunique','max','min','median','std']
train_data_,groupby_features = groupby_pivot_features(train_data_, train, groupby_features, 
                                                      col1 = 'api', col2 = 'index', agg_func = agg_func) 

# File_id + Api (index_diff): 'nunique','max','min','median','std'
train_diff = train.groupby(['file_id','tid'])['index'].diff().fillna(-999).values
train['index_diff'] = train_diff
train_diff = train.loc[train.index_diff!=-999] 

agg_func = ['nunique','max','min','median','std']
train_data_,groupby_features = groupby_pivot_features(train_data_, train_diff, groupby_features, 
                                                      col1 = 'api', col2 = 'index_diff', agg_func = agg_func) 


# # Bi-Gram Features

# In[7]:


# File_id（Api_2）:count,nunique
train['api_shift'] = train['api'].shift(-1)
train['api_2'] = train['api'] +'_' + train['api_shift']
train.drop(['api_shift'],axis=1,inplace=True)
api_count = train['api_2'].value_counts()

agg_func = ['count','nunique'] 
for func in agg_func:
    tmp = train.groupby(['file_id'])['api_2'].agg({'fileid_api_2_' + func: func}).reset_index() 
    train_data_ = pd.merge(train_data_,tmp,how='left', on='file_id')  
    
###
# File_id + tid (Api_2): count
api_value_counts = pd.DataFrame(api_count).reset_index()
api_value_counts.columns = ['api_2','api_2_count']
train = pd.merge(train, api_value_counts, on ='api_2' , how='left')

agg_func = ['count']
groupby_features =  []
train_data_,groupby_features = groupby_pivot_features(train_data_, train.loc[train.api_2_count>=20], groupby_features, 
                                                      col1 = 'api_2', col2 = 'tid', agg_func = agg_func)


# # Feature selection

# In[8]:


def lgb_logloss(preds,data):
    labels_ = data.get_label()
    classes_ = np.unique(labels_) 
    preds_prob = []
    for i in range(len(classes_)):
        preds_prob.append(preds[i*len(labels_):(i+1) * len(labels_)])
    preds_prob_ = np.vstack(preds_prob) 
    
    loss = [] 
    for i in range(preds_prob_.shape[1]): 
        sum_ = 0  
        for j in range(preds_prob_.shape[0]): 
            pred = preds_prob_[j,i] 
            if  j == labels_[i]:
                sum_ += np.log(pred)
            else:
                sum_ += np.log(1 - pred) 
             
        loss.append(sum_)  
         
    return 'loss is: ' ,-1 * (np.sum(loss) / preds_prob_.shape[1]),False


# In[ ]:


train_features = [col for col in train_data_.columns if col!='label' and col!='file_id']
print(len(train_features))
train_X, test_X, train_Y, test_Y = train_test_split( train_data_[train_features],train_data_['label'].values, test_size = 0.33) 

dtrain = lgb.Dataset(train_X,train_Y) 
dval   = lgb.Dataset(test_X,test_Y, reference = dtrain) 

params = {
        'task':'train', 
        'num_leaves': 255,
        'objective': 'multiclass',
        'num_class':8,
        #'min_data_in_leaf': 40,
        'min_data_in_leaf': 10,
        'learning_rate': 0.05,
        'feature_fraction': 0.85,
        'bagging_fraction': 0.9,
        'bagging_freq': 5, 
        'max_bin':128,
        'num_threads': 64,
        'random_state':100,
        'device': 'gpu',
        'device_id': 0,
        'gpu_platform_id': 0
    }  
lgb_model = lgb.train(params, dtrain, num_boost_round=500,valid_sets=[dtrain,dval], early_stopping_rounds=50, feval=lgb_logloss)  


# In[ ]:


fea_imp = pd.DataFrame({'feature':train_features, 'imp':lgb_model.feature_importance()}).sort_values('imp')
important_features = fea_imp.loc[fea_imp.imp >=1, 'feature'].values
important_features = list(important_features)

important_features.append('file_id')
important_features.append('label')


# In[ ]:


agg_func = ['max','min']
train_data_,groupby_features = groupby_pivot_features(train_data_[important_features], train.loc[train.api_2_count>=20], groupby_features, col1 = 'api_2', col2 = 'index', agg_func = agg_func)


# In[ ]:


train_features = [col for col in train_data_.columns if col!='label' and col!='file_id' and 'std' not in col and 'quantile' not in col]

dtrain = lgb.Dataset(train_data_.loc[train_X.index,train_features],train_data_.loc[train_X.index,'label'].values) 
dval   = lgb.Dataset(train_data_.loc[test_X.index,train_features],train_data_.loc[test_X.index,'label'].values, reference = dtrain) 

params = {
        'task':'train', 
        'num_leaves': 255,
        'objective': 'multiclass',
        'num_class':8,
        'min_data_in_leaf': 10,
        #'min_data_in_leaf': 1,
        'learning_rate': 0.05,
        'feature_fraction': 0.85,
        'bagging_fraction': 0.9,
        'bagging_freq': 5, 
        'max_bin':128,
        'num_threads': 64,
        'random_state':100,
        'device': 'gpu',
        'device_id': 0,
        'gpu_platform_id': 0
    }  
lgb_model = lgb.train(params, dtrain, num_boost_round=500,valid_sets=[dtrain,dval], early_stopping_rounds=50, feval=lgb_logloss)  


# In[ ]:


fea_imp = pd.DataFrame({'feature':train_features, 'imp':lgb_model.feature_importance()}).sort_values('imp')
important_features = fea_imp.loc[fea_imp.imp >=1, 'feature'].values
important_features = list(important_features)

important_features.append('file_id')
important_features.append('label')

train_data_[important_features].to_csv('./feats/train_feats_stats.csv',index = None)
 

