# -*- coding: utf-8 -*-
import pandas as pd 
import numpy as np 
import matplotlib.pyplot as plt 
from sklearn.preprocessing import LabelEncoder
import re 
from sklearn.feature_extraction import DictVectorizer
from sklearn.model_selection import train_test_split
import pandas as pd
from pandas.tseries.holiday import USFederalHolidayCalendar as calendar 
import pandas as pd
import numpy as np
import lightgbm as lgb
#import xgboost as xgb
from scipy.sparse import vstack, csr_matrix, save_npz, load_npz
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.model_selection import StratifiedKFold
#from sklearn.metrics import roc_auc_score
import gc
import sys 
gc.enable()

    
### FUNC ########################################################################
def get_xgb_imp(xgb, feat_names):
    from numpy import array
    imp_vals = xgb.booster().get_fscore()
    imp_dict = {feat_names[i]:float(imp_vals.get('f'+str(i),0.)) for i in range(len(feat_names))}
    total = array(imp_dict.values()).sum()
    return {k:v/total for k,v in imp_dict.items()}


def xgb_feat_importance(model,cols,file_name):
    print('-----> Feature importance ... ')
    feature_importance_dict = model.get_fscore()
    fs = ['f%i' % i for i in range(len(cols))]
    f1 = pd.DataFrame({'f': list(feature_importance_dict.keys()), 'importance': list(feature_importance_dict.values())})
    f2 = pd.DataFrame({'f': fs, 'feature_name': cols})
    feature_importance = pd.merge(f1, f2, how='right', on='f')
    feature_importance = feature_importance.fillna(0)
    feature_importance.sort_values(by='importance', ascending=False)
    print(feature_importance.sort_values)
    feature_importance.to_csv(file_name, index=False)

def add_avg_per(df,what_to_avg,on,new_name,include_delta=True,include_perc=True):
    if type(on) == str:
        _full = [on,what_to_avg]
        _fulla = [on,new_name]
    elif type(on) == list:
        _full = on.copy()
        _full.append(what_to_avg)
        _fulla = on.copy()
        _fulla.append(new_name)
    else:
        raise Exception('what type is on!')
    _avg = df.groupby(on)[_full].mean()
    _avg.columns = _fulla
    prev_len = len(df)
    df = df.merge(_avg,how='inner' , on=on)
    assert len(df) == prev_len
    if include_delta:
        df[str(new_name+'_delta')] = df[what_to_avg] - df[new_name]
    if include_perc:
        df[str(new_name+'_perc')] = (df[what_to_avg] - df[new_name])/df[new_name]
    return df 
        
def encode_dataset(train,test,meta,target_model='xgb'):
    y_train = train[meta['target']]
    train = train.drop([meta['target']],axis=1)
    assert train.shape[1] == test.shape[1]
    for i in range(train.shape[1]):
        assert train.columns[i] == test.columns[i]
    train_obs = len(train)
    #
    all_data = pd.concat([train,test],axis=0)
    for i,f in enumerate(meta['cols'].keys()):
        print(i,f,meta['cols'][f])
        if meta['cols'][f] == 'CAT':
            all_data[f] = all_data[f].astype('str')
            encoder = LabelEncoder()
            encoder.fit(np.unique(all_data[f].unique().tolist()))
            if target_model == 'xgb':
                all_data[f] = encoder.transform(all_data[f])
            else:
                all_data[f] = encoder.transform(all_data[f]).astype(int)
        elif meta['cols'][f] == 'NUM':
            all_data[f] = all_data[f].fillna(-1)
        elif meta['cols'][f] == 'REM':
            all_data = all_data.drop(f,axis=1)
        else:
            raise Exception(str(meta['cols'][f])+":unknown mapping")
    assert train_obs == len(y_train)
    return all_data , y_train
#################################################################################    
dtypes = {
        'MachineIdentifier':                                    'category',
        'ProductName':                                          'category',
        'EngineVersion':                                        'category',
        'AppVersion':                                           'category',
        'AvSigVersion':                                         'category',
        'IsBeta':                                               'int8',
        'RtpStateBitfield':                                     'float16',
        'IsSxsPassiveMode':                                     'int8',
        'DefaultBrowsersIdentifier':                            'float16',
        'AVProductStatesIdentifier':                            'float32',
        'AVProductsInstalled':                                  'float16',
        'AVProductsEnabled':                                    'float16',
        'HasTpm':                                               'int8',
        'CountryIdentifier':                                    'int16',
        'CityIdentifier':                                       'float32',
        'OrganizationIdentifier':                               'float16',
        'GeoNameIdentifier':                                    'float16',
        'LocaleEnglishNameIdentifier':                          'int8',
        'Platform':                                             'category',
        'Processor':                                            'category',
        'OsVer':                                                'category',
        'OsBuild':                                              'int16',
        'OsSuite':                                              'int16',
        'OsPlatformSubRelease':                                 'category',
        'OsBuildLab':                                           'category',
        'SkuEdition':                                           'category',
        'IsProtected':                                          'float16',
        'AutoSampleOptIn':                                      'int8',
        'PuaMode':                                              'category',
        'SMode':                                                'float16',
        'IeVerIdentifier':                                      'float16',
        'SmartScreen':                                          'category',
        'Firewall':                                             'float16',
        'UacLuaenable':                                         'float32',
        'Census_MDC2FormFactor':                                'category',
        'Census_DeviceFamily':                                  'category',
        'Census_OEMNameIdentifier':                             'float16',
        'Census_OEMModelIdentifier':                            'float32',
        'Census_ProcessorCoreCount':                            'float16',
        'Census_ProcessorManufacturerIdentifier':               'float16',
        'Census_ProcessorModelIdentifier':                      'float16',
        'Census_ProcessorClass':                                'category',
        'Census_PrimaryDiskTotalCapacity':                      'float32',
        'Census_PrimaryDiskTypeName':                           'category',
        'Census_SystemVolumeTotalCapacity':                     'float32',
        'Census_HasOpticalDiskDrive':                           'int8',
        'Census_TotalPhysicalRAM':                              'float32',
        'Census_ChassisTypeName':                               'category',
        'Census_InternalPrimaryDiagonalDisplaySizeInInches':    'float16',
        'Census_InternalPrimaryDisplayResolutionHorizontal':    'float16',
        'Census_InternalPrimaryDisplayResolutionVertical':      'float16',
        'Census_PowerPlatformRoleName':                         'category',
        'Census_InternalBatteryType':                           'category',
        'Census_InternalBatteryNumberOfCharges':                'float32',
        'Census_OSVersion':                                     'category',
        'Census_OSArchitecture':                                'category',
        'Census_OSBranch':                                      'category',
        'Census_OSBuildNumber':                                 'int16',
        'Census_OSBuildRevision':                               'int32',
        'Census_OSEdition':                                     'category',
        'Census_OSSkuName':                                     'category',
        'Census_OSInstallTypeName':                             'category',
        'Census_OSInstallLanguageIdentifier':                   'float16',
        'Census_OSUILocaleIdentifier':                          'int16',
        'Census_OSWUAutoUpdateOptionsName':                     'category',
        'Census_IsPortableOperatingSystem':                     'int8',
        'Census_GenuineStateName':                              'category',
        'Census_ActivationChannel':                             'category',
        'Census_IsFlightingInternal':                           'float16',
        'Census_IsFlightsDisabled':                             'float16',
        'Census_FlightRing':                                    'category',
        'Census_ThresholdOptIn':                                'float16',
        'Census_FirmwareManufacturerIdentifier':                'float16',
        'Census_FirmwareVersionIdentifier':                     'float32',
        'Census_IsSecureBootEnabled':                           'int8',
        'Census_IsWIMBootEnabled':                              'float16',
        'Census_IsVirtualDevice':                               'float16',
        'Census_IsTouchEnabled':                                'int8',
        'Census_IsPenCapable':                                  'int8',
        'Census_IsAlwaysOnAlwaysConnectedCapable':              'float16',
        'Wdft_IsGamer':                                         'float16',
        'Wdft_RegionIdentifier':                                'float16',
        'HasDetections':                                        'int8'
        }


print('Download Train and Test Data.\n')
DEBUG = False 
if DEBUG:
    print(">>> Debug mode ... ")
    train = pd.read_csv('data/train_sample.csv', dtype=dtypes, low_memory=True)
    train['MachineIdentifier'] = train.index.astype('uint32')
    test  = pd.read_csv('data/test_sample.csv',  dtype=dtypes, low_memory=True)
    test['MachineIdentifier']  = test.index.astype('uint32')
    gc.collect()
else:
    print(">>> Production mode ... ")
    train = pd.read_csv('data/train.csv', dtype=dtypes, low_memory=True)
    train['MachineIdentifier'] = train.index.astype('uint32')
    test  = pd.read_csv('data/test.csv',  dtype=dtypes, low_memory=True)
    test['MachineIdentifier']  = test.index.astype('uint32')
    gc.collect()


print('Transform almost all features to category.\n')
meta = {'target': 'HasDetections',
                'test_id': 'MachineIdentifier',
               'cols': {
                             'MachineIdentifier':                                  'REM',
                           'ProductName':                                          'CAT',
                           'EngineVersion':                                        'CAT',
                           'AppVersion':                                           'CAT',
                           'AvSigVersion':                                         'CAT',
                           'IsBeta':                                               'CAT',
                           'RtpStateBitfield':                                     'CAT',
                           'IsSxsPassiveMode':                                     'CAT',
                           'DefaultBrowsersIdentifier':                            'CAT',
                           'AVProductStatesIdentifier':                            'CAT',
                           'AVProductsInstalled':                                  'CAT',
                           'AVProductsEnabled':                                    'CAT',
                           'HasTpm':                                               'CAT',
                           'CountryIdentifier':                                    'CAT',
                           'CityIdentifier':                                       'CAT',
                           'OrganizationIdentifier':                               'CAT',
                           'GeoNameIdentifier':                                    'CAT',
                           'LocaleEnglishNameIdentifier':                          'CAT',
                           'Platform':                                             'CAT',
                           'Processor':                                            'CAT',
                           'OsVer':                                                'CAT',
                           'OsBuild':                                              'CAT',
                           'OsSuite':                                              'CAT',
                           'OsPlatformSubRelease':                                 'CAT',
                           'OsBuildLab':                                           'CAT',
                           'SkuEdition':                                           'CAT',
                           'IsProtected':                                          'CAT',
                           'AutoSampleOptIn':                                      'CAT',
                           'PuaMode':                                              'CAT',
                           'SMode':                                                'CAT',
                           'IeVerIdentifier':                                      'CAT',
                           'SmartScreen':                                          'CAT',
                           'Firewall':                                             'CAT',
                           'UacLuaenable':                                         'CAT',
                           'Census_MDC2FormFactor':                                'CAT',
                           'Census_DeviceFamily':                                  'CAT',
                           'Census_OEMNameIdentifier':                             'CAT',
                           'Census_OEMModelIdentifier':                            'CAT',
                           'Census_ProcessorCoreCount':                            'CAT',
                           'Census_ProcessorManufacturerIdentifier':               'CAT',
                           'Census_ProcessorModelIdentifier':                      'CAT',
                           'Census_ProcessorClass':                                'CAT',
                           'Census_PrimaryDiskTotalCapacity':                      'NUM',
                           'Census_PrimaryDiskTypeName':                           'CAT',
                           'Census_SystemVolumeTotalCapacity':                     'NUM',
                           'Census_HasOpticalDiskDrive':                           'CAT',
                           'Census_TotalPhysicalRAM':                              'CAT',
                           'Census_ChassisTypeName':                               'CAT',
                           'Census_InternalPrimaryDiagonalDisplaySizeInInches':    'NUM',
                           'Census_InternalPrimaryDisplayResolutionHorizontal':    'CAT',
                           'Census_InternalPrimaryDisplayResolutionVertical':      'CAT',
                           'Census_PowerPlatformRoleName':                         'CAT',
                           'Census_InternalBatteryType':                           'CAT',
                           'Census_InternalBatteryNumberOfCharges':                'NUM',
                           'Census_OSVersion':                                     'CAT',
                           'Census_OSArchitecture':                                'CAT',
                           'Census_OSBranch':                                      'CAT',
                           'Census_OSBuildNumber':                                 'NUM',
                           'Census_OSBuildRevision':                               'CAT',
                           'Census_OSEdition':                                     'CAT',
                           'Census_OSSkuName':                                     'CAT',
                           'Census_OSInstallTypeName':                             'CAT',
                           'Census_OSInstallLanguageIdentifier':                   'CAT',
                           'Census_OSUILocaleIdentifier':                          'CAT',
                           'Census_OSWUAutoUpdateOptionsName':                     'CAT',
                           'Census_IsPortableOperatingSystem':                     'CAT',
                           'Census_GenuineStateName':                              'CAT',
                           'Census_ActivationChannel':                             'CAT',
                           'Census_IsFlightingInternal':                           'CAT',
                           'Census_IsFlightsDisabled':                             'CAT',
                           'Census_FlightRing':                                    'CAT',
                           'Census_ThresholdOptIn':                                'CAT',
                           'Census_FirmwareManufacturerIdentifier':                'CAT',
                           'Census_FirmwareVersionIdentifier':                     'CAT',
                           'Census_IsSecureBootEnabled':                           'CAT',
                           'Census_IsWIMBootEnabled':                              'CAT',
                           'Census_IsVirtualDevice':                               'CAT',
                           'Census_IsTouchEnabled':                                'CAT',
                           'Census_IsPenCapable':                                  'CAT',
                           'Census_IsAlwaysOnAlwaysConnectedCapable':              'CAT',
                           'Wdft_IsGamer':                                         'CAT',
                           'Wdft_RegionIdentifier':                                'CAT'
                          }}




all_data , y_train = encode_dataset(train=train,test=test,meta=meta,target_model='lightgbm')
print("*****************************")
print(all_data.head())
train_obs = len(y_train)
train = all_data[:train_obs]
test = all_data[train_obs:]
train_ids = train.index
test_ids  = test.index

skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
skf.get_n_splits(train_ids, y_train)

lgb_test_result  = np.zeros(test_ids.shape[0])
#lgb_train_result = np.zeros(train_ids.shape[0])
#xgb_test_result  = np.zeros(test_ids.shape[0])
#xgb_train_result = np.zeros(train_ids.shape[0])
counter = 0
#Transform data using small groups to reduce memory usage
m = 100000
print('\nLightGBM\n')

for train_index, test_index in skf.split(train_ids, y_train):
    print('Fold {}\n'.format(counter + 1))
    print("**************************") 
    print("train_index:",train_index)
    print("**************************")
    print("test_index:",test_index)
    sys.stdout.flush()
    ##train = load_npz('train.npz')
    #X_fit = vstack([train.iloc[train_index[i*m:(i+1)*m]] for i in range(train_index.shape[0] // m + 1)])
    #X_val = vstack([train.iloc[test_index[i*m:(i+1)*m]]  for i in range(test_index.shape[0] //  m + 1)])
    X_fit = train.iloc[train_index]
    X_val = train.iloc[test_index]
    X_fit, X_val = csr_matrix(X_fit, dtype='float32'), csr_matrix(X_val, dtype='float32')
    y_fit, y_val = y_train[train_index], y_train[test_index]
    lgb_model = lgb.LGBMClassifier(max_depth=-1,
                                   n_estimators=30000,
                                   learning_rate=0.01,
                                   num_leaves=2**12-1,
                                   colsample_bytree=0.28,
                                   objective='binary', 
                                   n_jobs=30)
    lgb_model.fit(X_fit, y_fit, eval_metric='auc', 
                  eval_set=[(X_val, y_val)], 
                  verbose=100, early_stopping_rounds=100)
    ##test = load_npz('test.npz')
    test = csr_matrix(test, dtype='float32')
    lgb_test_result += lgb_model.predict_proba(test)[:,1]
    #xgb_test_result += xgb_model.predict_proba(test)[:,1]
    counter += 1



feature_imp = pd.DataFrame(sorted(zip(lgb_model.feature_importances_,train.columns)), columns=['Value','Feature'])
feature_imp.to_csv('base_lightgbm_feat_imp.csv', index=False)
submission = pd.read_csv('data/sample_submission.csv')
submission['HasDetections'] = lgb_test_result / counter
submission.to_csv('lgb_submission.csv', index=False)

print('\nDone.')
