# coding: utf-8
# In[1]:
try:
    import cudf as gd
    import nvstrings
    from librmm_cffi import librmm
    from nvstring_workaround import get_unique_tokens,on_gpu,get_token_counts,is_in
    from cudf_workaround import unique,rename_col,to_pandas
except:
    print('cudf not imported')


# In[2]:
from scipy import sparse
from sklearn.preprocessing import StandardScaler 
import time
from ml_robot import timer,bag_xgb_model,xgb_model,lgb_model,bag_lgb_model,write_log,sk_lgb_model,bag_sk_lgb_model
from ml_robot.metrics import get_score
import os
import pandas as pd
import numpy as np
from collections import OrderedDict,Counter
import sys
from ffm_model import FFM
import re
# In[3]:


WORK = '/raid/data/ml/malware'
CACHE = '%s/code/cache'%WORK
PATH = '%s/input'%WORK

MODE = None
METRIC = 'auc'
GPU = 0
if len(sys.argv)==2 and sys.argv[0].endswith('.py'):
    GPU = int(sys.argv[1])
    print('Reset GPU',GPU)
os.environ['CUDA_VISIBLE_DEVICES'] = str(GPU)
print("GPU:",GPU)


IDCOL = 'MachineIdentifier'
YCOL = 'HasDetections'


# # Functions

# In[5]:

def remove_bad(data,names,bad = []):
    bad.append(IDCOL)
    bad.append(YCOL)
    
    good = [c for c,i in enumerate(names) if i not in bad]
    names = [i for c,i in enumerate(names) if i not in bad]
    print(names)
    return data[:,good],names

def mkdir(path):
    if os.path.exists(path)==0:
        os.mkdir(path)

@timer
def get_dtype(path,nrows=10):
    """get data type for cudf.read_csv
    by using pandas's read_csv with a small number of rows.

    Parameters
    ----------
    path : str, path of the csv file
    nrows : int, default: 10
        number of rows to read for pd.read_csv

    Returns
    ----------
    col2dtype : dictionary, column name => data type
    """
    if nrows is not None:
        train = pd.read_csv(path,nrows=nrows)
    else:
        train = pd.read_pickle(path.replace('.csv','.pkl'))
    col2dtype = OrderedDict()
    for col in train.columns:
        if train[col].dtype=='O':
            col2dtype[col] = 'str'
        elif train[col].dtype==np.int64:
            col2dtype[col] = 'int32'
        else:
            col2dtype[col] = 'float32'
    return col2dtype

@timer
def read_csv_hash_nvstring(path,cols=None):
    """read a csv file to a dictionary of nvstring objects
    Parameters
    ----------
    path : str, path of the csv file
    cols : list, a list of column names to be read

    Returns
    ----------
    str_cols : list of string column name 
    df: gd.dataframe
    """
    dic = get_dtype(path,1000)
    col_names = [i for i in dic]
    dtypes = [dic[i] for i in col_names]

    gd_cols = gd.io.csv.read_csv_strings(path,
                        names=col_names, dtype=dtypes,
                        skiprows=1)
    str_cols = [] 
    df = gd.DataFrame()
    for name,dtype,ds in zip(col_names,dtypes,gd_cols):
        if cols is not None and name not in cols:
            continue
        if dtype =='str':
            df[name] = on_gpu(ds,'hash')
            str_cols.append(name)
        else:
            df[name] = ds
            
    del gd_cols
    return df,str_cols

@timer
def read_csv_with_nvstring(path,cols=None,allstr=False):
    """read a csv file to a dictionary of nvstring objects
    Parameters
    ----------
    path : str, path of the csv file
    cols : list, a list of column names to be read

    Returns
    ----------
    str_cols : dictionary, column name => nvstring object
    df: gd.dataframe
    """
    dic = get_dtype(path,1000)
    col_names = [i for i in dic]
    dtypes = [dic[i] for i in col_names]
    if allstr:
        dtypes = ['str' for i in col_names]
    gd_cols = gd.io.csv.read_csv_strings(path,
                        names=col_names, dtype=dtypes,
                        skiprows=1)
    str_cols = {} # column name => nvstring object
    df = gd.DataFrame()
    for name,dtype,ds in zip(col_names,dtypes,gd_cols):
        if cols is not None and name not in cols:
            continue
        if dtype =='str':
            str_cols[name] = ds
        else:
            df[name] = ds
    del gd_cols
    return df,str_cols

@timer
def read_csv_hash(path):
    """read a csv file by hashing the str cols
    Parameters
    ----------
    path : str, path of the csv file

    Returns
    ----------
    df: gd.dataframe
    """
    dic = get_dtype(path,1000)
    col_names = [i for i in dic]
    dtypes = [dic[i] for i in col_names]
    str_cols = [i for i in col_names if dic[i]=='str'] 
    dtypes = ['int32' if i=='str' else i for i in dtypes]

    gdf = gd.read_csv(path,names=col_names,dtype=dtypes,skiprows=1)
    return gdf,str_cols

@timer
def split():
    mkdir(CACHE)
    out = '%s/tr.csv'%CACHE
    if os.path.exists(out):
        return
    gdf,ds = read_csv_with_nvstring('%s/train.csv'%PATH)
    col = 'AvSigVersion'
    nv = ds[col]
    nvs = nv.split_column('.')
    c = 1 # use the 2nd digit (0 based) to split train data.
    xcol = '%s_%d'%(col,c)    
    gdf[xcol] = on_gpu(nvs[c],'stoi')
    df = gdf.to_pandas()
    tmp = pd.read_csv('%s/train.csv'%PATH,nrows=10)
    cols = [i for i in tmp.columns]
    for col in ds:
        df[col] = ds[col].to_host()
    df = df[cols+[xcol]] # reorder the columns
    mask = df[xcol]>=275
    df['random'] = np.random.random(df.shape[0])
    random_mask = df.random<0.13
    print(df[mask].shape,df[~mask].shape)
    mask = mask | random_mask
    print(df[mask].shape,df[~mask].shape)
    del ds
    del gdf
    del nv,nvs
    df[~mask].drop('random',axis=1).to_csv(out,index=False)
    df[mask].drop('random',axis=1).to_csv(out.replace('tr','va'),index=False)


# In[6]:


@timer
def get_id_y(mode):
    if mode == 'cv':
        tr,te = 'tr','va'
    else:
        tr,te = 'train','test'
    out = '%s/%s_y.pkl'%(CACHE,tr)
    if os.path.exists(out):
        y = pd.read_pickle(out)[YCOL].values
        if mode == 'cv':
            yt = pd.read_pickle(out.replace(tr,te))[YCOL].values
        else:
            yt = None
        out = out.replace('_y','_id')
        tr_id = pd.read_pickle(out)[IDCOL].values
        te_id = pd.read_pickle(out.replace(tr,te))[IDCOL].values
        return y,yt,tr_id,te_id

    tr_path,te_path = get_tr_te_paths(mode)
    gtr_y,gtr_id = read_csv_with_nvstring(tr_path,[IDCOL,YCOL])
    gte_y,gte_id = read_csv_with_nvstring(te_path,[IDCOL,YCOL])
    
    y = gtr_y[YCOL].to_array()
    pd.DataFrame({YCOL:y}).to_pickle(out)
    #pd.DataFrame({YCOL:y}).to_csv(out,index=False,compression='gzip')
    del gtr_y
    
    if mode == 'cv':
        yt = gte_y[YCOL].to_array()
        pd.DataFrame({YCOL:yt}).to_pickle(out.replace(tr,te))
        #pd.DataFrame({YCOL:yt}).to_csv(out.replace(tr,te),index=False,compression='gzip')
    else:
        yt = None
    del gte_y
    
    out = out.replace('_y','_id')
    
    tr_id = np.array(gtr_id[IDCOL].to_host())
    pd.DataFrame({IDCOL:tr_id}).to_pickle(out)
    #pd.DataFrame({IDCOL:tr_id}).to_csv(out,index=False,compression='gzip')
    del gtr_id
    
    te_id = np.array(gte_id[IDCOL].to_host())
    pd.DataFrame({IDCOL:te_id}).to_pickle(out.replace(tr,te))
    #pd.DataFrame({IDCOL:te_id}).to_csv(out.replace(tr,te),index=False,compression='gzip')
    del gte_id
    
    return y,yt,tr_id,te_id


# In[7]:


def get_tr_te_paths(mode):
    if mode == 'cv':
        tr_path = '%s/tr.csv'%CACHE
        te_path = '%s/va.csv'%CACHE
    else:
        tr_path = '%s/train.csv'%PATH
        te_path = '%s/test.csv'%PATH
    return tr_path,te_path

def rm_cols(gdf,cols):
    gcols = [i for i in gdf.columns]
    cols = set(cols)
    for col in cols:
        if col in gcols:
            del gdf[col]
    return gdf

def exist(i,mode):
    out = '%s/fea%s_%d_tr.pkl'%(CACHE,mode,i)
    return os.path.exists(out)

def load(i,mode):
    out = '%s/fea%s_%d_tr.pkl'%(CACHE,mode,i)
    x = pd.read_pickle(out)
    xt = pd.read_pickle(out.replace('tr','te'))
    cols = [i for i in x.columns]
    return x.values,xt.values,cols

def rm(i,mode):
    out = '%s/fea%s_%d_tr.pkl'%(CACHE,mode,i)
    if os.path.exists(out):
        os.remove(out)
    out = out.replace('tr','te')
    if os.path.exists(out):
        os.remove(out)

def dump(i,mode,x,xt,cols):
    out = '%s/fea%s_%d_tr.pkl'%(CACHE,mode,i)
    pd.DataFrame(x,columns=cols).to_pickle(out)
    pd.DataFrame(xt,columns=cols).to_pickle(out.replace('tr','te'))

def has_sparse(X):
    for i in X:
        if isinstance(i,sparse.csr_matrix):
            return True
    return False    
    
@timer
def build(mode,build_list=[1],cache=[0]):
    assert len(build_list) == len(cache)
    y,yt,tr_id,te_id = get_id_y(mode)
    if len(build_list)==0:
        return None,None,y,yt,None,tr_id,te_id
    X,Xt,names = [],[],[]
    for i,s in zip(build_list,cache):
        if s and exist(i,mode):
            x,xt,cols = load(i,mode)
        else:
            x,xt,cols = eval('build%d(mode)'%i)
            if s:
                dump(i,mode,x,xt,cols)
            else:
                rm(i,mode)
        X.append(x)
        Xt.append(xt)
        names.append(cols)
    if True:
        if has_sparse(X):
           X = sparse.hstack(X,format='csr')
           Xt = sparse.hstack(Xt,format='csr')
        else:
           X = np.hstack(X)
           Xt = np.hstack(Xt)
    names = [i for cols in names for i in cols]
    return X,Xt,y,yt,names,tr_id,te_id   
      
@timer 
def gdf_factorize(gtr,strcols):
    for col in strcols:
        if col in gtr.columns:
            gtr[col],_ = gtr[col].fillna(0).factorize()
    return gtr

@timer
def build1(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,strcols = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)
    badcols = ['Census_OSVersion','AVProductsInstalled','OsVer','AvSigVersion','PuaMode','AvSigVersion_1','OsSuite','ProductName','SmartScreen','SkuEdition', 'Census_IsPortableOperatingSystem','Census_ThresholdOptIn','EngineVersion']#,'AppVersion','CityIdentifier']
    #badcols += [i for i in gte.columns if 'Is' in i]
    gtr = rm_cols(gtr,[IDCOL,YCOL]+badcols)
    gte = rm_cols(gte,[IDCOL,YCOL]+badcols)
    print("build1",len(gtr),len(gte))
    return post_gdf(gtr,gte)   

def mtr(tr,col,te):
    #print('0')
    tr[col] = tr[col].fillna(-999)
    #print('1')
    tr[YCOL] = tr[YCOL].astype('float32')
    tr[col] = tr[col].astype('int32')
    te[col] = te[col].astype('int32')
    #print('2')
    dg = tr.groupby(col).agg({YCOL:'mean'})
    #print('3')
    dg = rename_col(dg,'mean_%s'%YCOL,'mtr_%s'%col)
    #print('4')
    if len(dg)>2:
        #tr = tr.merge(dg,on=[col],how='left')
        te = te.merge(dg,on=[col],how='left')
        #ptr = dg.to_pandas()
        #pte = te.to_pandas()
        #del dg,te
        #dg = pd.DataFrame.from_pandas(ptr)
        #te = pd.DataFrame.from_pandas(pte)
        
    del dg
    return te

@timer
def mtr_encode(tr,te,cols):
    tr = rm_cols(tr,[i for i in tr.columns if i not in cols+[YCOL]])
    te = rm_cols(te,[i for i in te.columns if i not in cols+[YCOL]])
    tr['idx'] = np.arange(len(tr))
    tr['random'] = np.random.random(len(tr))
    tr = tr.sort_values(by='random')
    tr.drop_column('random') 
    N = len(tr)//2
    tr1 = tr[:N]
    tr2 = tr[N:]
    te['idx'] = np.arange(len(te))
    for col in cols:
        if col not in tr.columns:
            continue
        #print('mtr tr1')
        tr2 = mtr(tr1,col,tr2)
        #tr2.to_pandas().to_pickle('tr2.pkl')
        #tr1.to_pandas().to_pickle('tr1.pkl')
        #print('mtr tr2')
        tr1 = mtr(tr2,col,tr1)
        #print('mtr te')
        te = mtr(tr,col,te)
    if len(tr1.columns)!=len(tr2.columns):
        del tr1,tr2
        return rm_cols(tr,[i for i in tr.columns]),rm_cols(te,[i for i in te.columns])
    del tr
    tr = gd.concat([tr1,tr2])
    tr = tr.sort_values(by='idx')
    tr = rm_cols(tr,[i for i in tr.columns if i.startswith('mtr_')==0])
    del tr1,tr2
    te = te.sort_values(by='idx')
    te = rm_cols(te,[i for i in te.columns if i.startswith('mtr_')==0])
    return tr,te

@timer
def count_encode(df,cols,keep=[]):
    df['idx'] = np.arange(len(df))
    for col in cols:
        if col not in df.columns:
            continue
        df[col] = df[col].fillna(0)
        dg = df.groupby(col).agg({col:'count'})
        if True:#len(dg)>2:
            df = df.merge(dg,on=[col],how='left')
            #df.drop_column(col)
        del dg
    df = df.sort_values(by='idx')
    df.drop_column('idx')
    df = rm_cols(df,[i for i in df.columns if i.startswith('count_')==0 and i not in keep])
    return df

@timer
def count_encode2way(df,cols,keep=[]):
    df['idx'] = np.arange(len(df))
    for col in cols:
        if col[0] not in df.columns or col[1] not in df.columns:
            continue
        for c in col:
            df[c] = df[c].fillna(0)
        dg = df.groupby(col).agg({col[0]:'count'})
        tag = '_'.join(col)
        dg = rename_col(dg,'count_%s'%col[0],'count_%s'%(tag))
        if len(dg)>2:
            df = df.merge(dg,on=col,how='left')
        del dg
    df = df.sort_values(by='idx')
    df.drop_column('idx')
    df = rm_cols(df,[i for i in df.columns if i.startswith('count_')==0  and i not in keep])
    return df

def reset_col_dtype(gtr,gte):
    for col in gtr.columns:
        a,b = str(gtr[col].dtype),str(gte[col].dtype)
        if a!=b:
            gtr[col] = gtr[col].astype('float32')
            gte[col] = gte[col].astype('float32')
    return gtr,gte

@timer
def build0(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)
    gtr = rm_cols(gtr,[IDCOL,YCOL])
    gte = rm_cols(gte,[IDCOL,YCOL])
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    num_cols = get_true_numerical_columns()    
    str_cols = [i for i in df.columns if i not in num_cols]
    df = df[str_cols+num_cols]
    df = gdf_factorize(df,str_cols)     
    data = df.to_pandas()
    del df
    cols = [i for i in data.columns]
    data = data.values
    print("# of ",data.shape)
    return data[:N],data[N:],cols

@timer
def build2(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)
    cols = ['Census_OSVersion']
    #count_cols = ['Census_ProcessorModelIdentifier', 'Census_ActivationChannel', 'Census_PrimaryDiskTotalCapacity']
    #count_cols = ['Census_InternalPrimaryDiagonalDisplaySizeInInches', 'Census_ProcessorModelIdentifier', 'Census_ActivationChannel', 'Wdft_IsGamer', 'Census_PrimaryDiskTotalCapacity', 'Census_PowerPlatformRoleName', 'Census_OSVersion', 'DefaultBrowsersIdentifier', 'Census_TotalPhysicalRAM', 'Census_GenuineStateName', 'IeVerIdentifier','Census_OSWUAutoUpdateOptionsName', 'SkuEdition', 'Census_InternalBatteryNumberOfCharges', 'Processor']
    #count_cols = [i for i in gte.columns if i not in [IDCOL,YCOL,'AvSigVersion','AvSigVersion_1']]
    #count_cols = ['CityIdentifier', 'Census_OSBuildRevision', 'Census_TotalPhysicalRAM']#,'AvSigVersion']
    bad = [i for i in gtr.columns if i not in count_cols]
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    #badcols = [i for i in gtr.columns]
    #count_cols = ['Census_OSBuildRevision', 'Census_PrimaryDiskTotalCapacity', 'Census_MDC2FormFactor']
    print("# of count cols",len(count_cols))
    df = count_encode(df,count_cols)
    data = df.to_pandas()
    del df
    cols = [i for i in data.columns]
    data = data.values
    print("count shape",data.shape)
    return data[:N],data[N:],cols

@timer
def build6(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)
    c1 = ['AvSigVersion']#,'AVProductStatesIdentifier']
    c2 = ['Census_OSBuildRevision','EngineVersion','Census_OSWUAutoUpdateOptionsName']#,'Census_TotalPhysicalRAM','CityIdentifier','Census_ProcessorModelIdentifier','AVProductStatesIdentifier']#,'OsVer']
    c2 += [i for i in gte.columns if 'Name' in i]
    c2 += get_true_numerical_columns()
    c2 = list(set(c2))
    bad = ['Census_PowerPlatformRoleName','Census_InternalBatteryNumberOfCharges','Census_OSSkuName']
    bad += ['Census_TotalPhysicalRAM', 'Census_SystemVolumeTotalCapacity', 'GeoNameIdentifier']
    c2 = [i for i in c2 if i not in bad] 
    #c2 = ['Census_OSBuildRevision', 'IsBeta', 'OsBuild', 'RtpStateBitfield', 'AutoSampleOptIn']#,'OsPlatformSubRelease', 'AVProductsEnabled']
    #c2 = [i for i in gte.columns if i not in [IDCOL,YCOL,'AvSigVersion','AvSigVersion_1']]
    count_cols = []
    for i in c1:
        for j in c2:
            count_cols.append([i,j])
    bad = [i for i in gtr.columns if i not in c1 and i not in c2]
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    print("# of count cols",len(count_cols))
    df = count_encode2way(df,count_cols,keep=c1)
    print(df.columns)
    df = count_encode(df,c1)
    print(df.columns)
    data = to_pandas(df)
    print(data.head())
    del df
    for i in c1:
        for j in c2:
            data['count_%s_%s'%(i,j)] = data['count_%s_%s'%(i,j)]/data['count_%s'%(i)]
    data = data.drop(['count_%s'%i for i in c1],axis=1)
    print(data.head())
    for i in data.columns:
        print(i,data[i].max())
        assert data[i].max()<=1.0
    cols = [i for i in data.columns]
    data = data.values
    print("count shape",data.shape)
    return data[:N],data[N:],cols

@timer
def build14(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)
    tr,te = load_AvSig(mode=mode)
    tr = AvSig_to_timestamp(tr,unit='day')
    te = AvSig_to_timestamp(te,unit='day')
    gtr['AvDay'] = np.ascontiguousarray(tr)
    gte['AvDay'] = np.ascontiguousarray(te)
    c1 = ['AvDay'] 
    #c1 = ['AvSigVersion']#,'AVProductStatesIdentifier']
    #c2 = ['Census_OSBuildRevision','Census_OSBuildNumber','RtpStateBitfield']#,'Census_OSBranch']#'Census_TotalPhysicalRAM','AppVersion']#,'OsVer']
    c2 = ['Census_OSBuildRevision', 'IsBeta', 'OsBuild', 'RtpStateBitfield', 'AutoSampleOptIn']#,'OsPlatformSubRelease', 'AVProductsEnabled']
    #c2 = [i for i in gte.columns if i not in [IDCOL,YCOL,'AvSigVersion','AvSigVersion_1']]
    count_cols = []
    for i in c1:
        for j in c2:
            count_cols.append([i,j])
    bad = [i for i in gtr.columns if i not in c1 and i not in c2]
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    print("# of count cols",len(count_cols))
    df = count_encode2way(df,count_cols,keep=c1)
    print(df.columns)
    df = count_encode(df,c1)
    print(df.columns)
    data = df.to_pandas()
    print(data.head())
    del df
    for i in c1:
        for j in c2:
            data['count_%s_%s'%(i,j)] = data['count_%s_%s'%(i,j)]/data['count_%s'%(i)]
    data = data.drop(['count_%s'%i for i in c1],axis=1)

    print(data.head())
    cols = [i for i in data.columns]
    data = data.values
    print("count shape",data.shape)
    return data[:N],data[N:],cols

@timer
def build13(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)

    acol = 'AvSigVersion'
    count_cols = []
    c1 = get_true_numerical_columns()
    c2 = ['OrganizationIdentifier','CityIdentifier','Census_ProcessorModelIdentifier']
    #count_cols.extend([[acol,i,j] for i in c1 for j in c2])

    c1 = ['Census_OSInstallTypeName']
    c2 = ['Processor']#,'Platform','CountryIdentifier']
    count_cols.extend([[acol,i,j] for i in c1 for j in c2])

    c1 = ['EngineVersion']
    c2 = ['DefaultBrowsersIdentifier','AVProductStatesIdentifier']#,'UacLuaenable']
    count_cols.extend([[acol,i,j] for i in c1 for j in c2])

    c1 = ['OrganizationIdentifier']
    c2 = ['IsBeta']#'Census_OSEdition']#'Census_OEMModelIdentifier']
    count_cols.extend([[acol,i,j] for i in c1 for j in c2])

    c1 = ['Census_ProcessorModelIdentifier','Census_FirmwareManufacturerIdentifier']
    c1 += ['Census_PrimaryDiskTypeName','Wdft_RegionIdentifier']
    #c1 = [i for i in gte.columns if 'Identifier' in i]
    #c1 = [i for i in gte.columns if 'Name' in i]
    c2 = ['AppVersion']#,'Census_OEMModelIdentifier']
    count_cols.extend([[acol,i,j] for i in c1 for j in c2])


    allc = [j for i in count_cols for j in i]
    allc = set(allc) 
    bad = [i for i in gtr.columns if i not in allc]
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    print("# of count cols",len(count_cols))
    df = count_encode(df,[acol],keep=allc)
    print(df.columns)
    df = count_encode2way(df,count_cols,keep=['count_%s'%acol])
    print(df.columns)
    data = to_pandas(df)
    print(data.head())
    del df
    for col in count_cols:
        tag = '_'.join(col)
        tag = 'count_%s'%tag
        data[tag] = data[tag]/data['count_%s'%acol]
    print(data.head())
    data.drop(['count_%s'%acol],axis=1,inplace=True)
    for i in data.columns:
        print(i,data[i].max())
        assert data[i].max()<=1.0
    cols = [i for i in data.columns]
    data = data.values
    print("count shape",data.shape)
    return data[:N],data[N:],cols

@timer
def build12(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)
    c1 = ['AvSigVersion','CityIdentifier']#,'AVProductStatesIdentifier']
    c2 = ['Census_ProcessorModelIdentifier','Census_IsPenCapable']#,'Census_ProcessorManufacturerIdentifier']
    #c2 = ['CityIdentifier']#,'AppVersion','Census_IsTouchEnabled']#,'Processor']#,'AppVersion']#,'Census_ProcessorClass']
    #c2 = ['Census_OSBuildRevision', 'IsBeta', 'OsBuild', 'RtpStateBitfield', 'AutoSampleOptIn','OsPlatformSubRelease', 'AVProductsEnabled']
    #c2 = [i for i in gte.columns if i not in [IDCOL,YCOL,'AvSigVersion','AvSigVersion_1']]
    count_cols = []
    for j in c2:
        count_cols.append(c1+[j])
    bad = [i for i in gtr.columns if i not in c1 and i not in c2]
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    print("# of count cols",len(count_cols))
    df = count_encode2way(df,count_cols)
    print(df.columns)
    data = df.to_pandas()
    print(data.head())
    del df

    print(data.head())
    cols = [i for i in data.columns]
    data = data.values
    print("count shape",data.shape)
    return data[:N],data[N:],cols




@timer
def build11(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)

    count_cols = [
        #'EngineVersion-Census_OSWUAutoUpdateOptionsName'.split('-'),
        #'SmartScreen-Census_ProcessorClass'.split('-'),
        'AppVersion-Census_OSInstallTypeName'.split('-'),
    ]
    cc = set([j for i in count_cols for j in i])
    bad = [i for i in gtr.columns if i not in cc]
    #bad = [i for i in gtr.columns if i not in c1 and i not in c2]
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    print("# of count cols",len(count_cols))
    df = count_encode2way(df,count_cols)
    print(df.columns)
    data = df.to_pandas()
    del df
    print(data.head())
    cols = [i for i in data.columns]
    data = data.values
    print("count shape",data.shape)
    return data[:N],data[N:],cols

def load_AvSig(mode):
    tr_out = '%s/%s_tr_Avsig.pkl'%(CACHE,mode)
    te_out = '%s/%s_te_Avsig.pkl'%(CACHE,mode)
    if os.path.exists(tr_out):
        return pd.read_pickle(tr_out),pd.read_pickle(te_out)
    tr_path,te_path = get_tr_te_paths(mode)
    try: 
        tr = pd.read_pickle(tr_path.replace('.csv','.pkl')) 
        te = pd.read_pickle(te_path.replace('.csv','.pkl'))  
    except:
        tr,te = pd.read_csv(tr_path),pd.read_csv(te_path) 
    tr[['AvSigVersion']].to_pickle(tr_out)
    te[['AvSigVersion']].to_pickle(te_out)
    return tr[['AvSigVersion']],te[['AvSigVersion']]

def AvSig_to_timestamp(df,unit='day'):
    col = 'DateAS'
    datedictAS = np.load('%s/AvSigVersionTimestamps.npy'%CACHE)[()]
    df[col] = df['AvSigVersion'].map(datedictAS) 
    df[col] = pd.to_datetime(df[col])
    day = df[col].dt.day.values
    month = df[col].dt.month.values
    year = df[col].dt.year.values
    if unit=='day':
        return (year*12+month)*31+day
    assert 0

def build7(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    c1 = ['AvSigVersion']#,'AVProductStatesIdentifier']
    c2 = ['CityIdentifier', 'Census_OSBuildRevision', 'Census_TotalPhysicalRAM']
    gtr,strs = read_csv_with_nvstring(tr_path,c1+c2)
    gte,stes = read_csv_with_nvstring(te_path,c1+c2)

    gtr['av'] = on_gpu(strs['AvSigVersion'].split_column('.')[1],'stoi')
    gte['av'] = on_gpu(stes['AvSigVersion'].split_column('.')[1],'stoi')
    for col in c2:
        if col not in gtr.columns:
            gtr[col] = on_gpu(strs[col],'hash')
            gte[col] = on_gpu(stes[col],'hash') 	
    c1 = ['av']
    del strs,stes
    count_cols = []
    for i in c1:
        for j in c2:
            count_cols.append([i,j])
    bad = [i for i in gtr.columns if i not in c1 and i not in c2]
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    print("# of count cols",len(count_cols))
    df = count_encode2way(df,count_cols,keep=c1)
    df = count_encode(df,c1)
    data = df.to_pandas()
    del df
    for i in c1:
        for j in c2:
            data['count_%s_%s'%(i,j)] = data['count_%s_%s'%(i,j)]/data['count_%s'%(i)]
    data = data.drop(['count_%s'%i for i in c1],axis=1)

    print(data.head())
    cols = [i for i in data.columns]
    data = data.values
    print("count shape",data.shape)
    return data[:N],data[N:],cols

def concat_str(strs,cols):
    s = strs[cols[0]]
    for col in cols[1:]:
        s = s.cat(strs[col],sep=':', na_rep='_')
    return on_gpu(s,'hash')

@timer
def build10(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    cols = [
        #"Wdft_IsGamer,Census_OSInstallLanguageIdentifier".split(','),
        #"Census_ProcessorCoreCount,Wdft_RegionIdentifier,Census_OSBuildNumber".split(','),
        #"Census_ProcessorCoreCount,Census_OEMNameIdentifier,CityIdentifier".split(','),
        #"GeoNameIdentifier,Census_OEMNameIdentifier,Census_OSBuildRevision".split(',')
        "AVProductStatesIdentifier-Firewall".split('-'),
        'EngineVersion-Census_OSWUAutoUpdateOptionsName'.split('-'),
    ]
    #cols += [i+['SmartScreen'] for i in cols]
    ucols = list(set([x for i in cols for x in i]))

    _,strs = read_csv_with_nvstring(tr_path,allstr=True,cols=ucols)
    _,stes = read_csv_with_nvstring(te_path,allstr=True,cols=ucols)

    gtr = gd.DataFrame()
    gte = gd.DataFrame()

    for c,cs in enumerate(cols):
        gtr['-'.join(cs)] = concat_str(strs,cs)
        gte['-'.join(cs)] = concat_str(stes,cs)

    ptr = gtr.to_pandas()
    pte = gte.to_pandas()
    cols = [i for i in ptr.columns]
    del gtr,gte,strs,stes
    return ptr.values,pte.values,cols

@timer
def build9(mode):
    path = '../bench'
    if mode in ['stackpre_cv','cv']:
        tr = 'tr.npz'
        te = 'va.npz'
    elif mode in ['sub']:
        tr = 'train.npz'
        te = 'test.npz'
    else:
        assert 0
    train = sparse.load_npz('%s/%s'%(path,tr))
    test = sparse.load_npz('%s/%s'%(path,te))
    train = sparse.csr_matrix(train, dtype='float32')
    test = sparse.csr_matrix(test, dtype='float32')
    return train,test,['sparse%d'%i for i in range(train.shape[1])]

def av_encode(ds):
    nvs = ds['AvSigVersion'].split_column('.')
    s = nvstrings.to_device(["0." for i in range(nvs[0].size())])
    for nv in nvs:
        s = s.cat(nv,sep='', na_rep='')
    del nvs
    return on_gpu(s,'stof',dtype=np.float32)

def build8(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    c1 = ['AvSigVersion']#,'AVProductStatesIdentifier']
    c2 = ['CityIdentifier', 'Census_OSBuildRevision', 'Census_TotalPhysicalRAM']
    gtr,strs = read_csv_with_nvstring(tr_path,c1+c2)
    gte,stes = read_csv_with_nvstring(te_path,c1+c2)

    for col in c2:
        if col not in gtr.columns:
            gtr[col] = on_gpu(strs[col],'hash')
            gte[col] = on_gpu(stes[col],'hash')
    bad = [i for i in gtr.columns if i not in c1 and i not in c2]
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)

    gtr['av'] = av_encode(strs)
    gte['av'] = av_encode(stes)
    del strs,stes
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    N = len(gtr)
    del gtr,gte 
 
    dg = df.groupby('av').agg({'av':'count'})
    dg.drop_column('count_av')
    dg = dg.sort_values(by='av')
    dg['av_values'] = np.arange(len(dg))    
    df = df.merge(dg,on=['av'],how='left')
    df.drop_column('av')
    df1 = df.copy()
    del dg

    count_cols = c2
    print("# of count cols",len(count_cols))
    df = count_encode_next_prev(df,count_cols)
    data = df.to_pandas()
    del df
    df1 = count_encode2way(df1,[['av_values',i] for i in c2])
    data1 = df1.to_pandas()
    print(data1.columns)
    del df1 
    for col in c2:
        #data['diff_count_%s'%(col)] = data['next_count_%s'%(col)] - data['prev_count_%s'%(col)]
        data['prev_count_%s'%(col)] = data1['count_av_values_%s'%(col)] - data['prev_count_%s'%(col)]
        data['prev_count_%s'%(col)] /= data1['count_av_values_%s'%(col)]
        data['next_count_%s'%(col)] = data1['count_av_values_%s'%(col)] - data['next_count_%s'%(col)]
        data['next_count_%s'%(col)] /= data1['count_av_values_%s'%(col)]  
    #cols = ['diff_count_%s'%(col) for col in c2]
    #data = data[cols]
    print(data.head())
    cols = [i for i in data.columns]
    data = data.values
    print("count shape",data.shape)
    return data[:N],data[N:],cols

@timer
def count_encode_next_prev(df,cols):
    df['idx'] = np.arange(len(df))
    df['av_values'] = df['av_values'].fillna(0)
    df['av_next'] = df['av_values']+1
    df['av_prev'] = df['av_values']-1
    for col in cols:
        df[col] = df[col].fillna(0)
        dg = df.groupby(['av_values',col]).agg({col:'count'})

        dg = rename_col(dg,'count_%s'%col,'prev_count_%s'%(col))
        dg = rename_col(dg,'av_values','av_prev')
        
        df = df.merge(dg,on=['av_prev',col],how='left')

        dg = rename_col(dg,'prev_count_%s'%(col),'next_count_%s'%(col))
        dg = rename_col(dg,'av_prev','av_next')
        df = df.merge(dg,on=['av_next',col],how='left')
 
        del dg
    df = df.sort_values(by='idx')
    df = rm_cols(df,[i for i in df.columns if i.startswith('next')==0 and i.startswith('prev')==0 ])
    return df

@timer
def pd_count_encode_next_prev(df,cols):
    df['av_next'] = df['av_values']+1
    df['av_prev'] = df['av_values']-1
    ds = df[['av_next','av_prev']]
    for col in cols:
        dg = df.groupby(['av_values',col]).agg({'av_next':'count'}).reset_index()
        dg.drop(col,axis=1,inplace=True)
        dg.columns = ['av_prev','prev_count_%s'%(col)]

        ds = ds.merge(dg,on=['av_prev'],how='left')

        dg.columns = ['av_next','next_count_%s'%(col)]
        ds = ds.merge(dg,on=['av_next'],how='left')

        del dg
    del df
    bad = [i for i in ds.columns if  i.startswith('next')==0 and i.startswith('prev')==0]
    return ds.drop(bad,axis=1)

@timer
def build5(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)
    bad = [i for i in gtr.columns if i!='SmartScreen']
    gtr = rm_cols(gtr,bad)
    gte = rm_cols(gte,bad)
    N = len(gtr)
    gtr,gte = reset_col_dtype(gtr,gte)
    df = gd.concat([gtr,gte])
    del gtr,gte
    data = df.to_pandas()
    del df
    data = pd.get_dummies(data['SmartScreen'])
    #cols = [re.sub(r'[\W_]+', '', i) for i in data.columns]
    cols = ['SmartScreen_%d'%c for c in range(data.shape[1])]
    print(cols)
    data = data.values
    print("onehot shape",data.shape)
    return data[:N],data[N:],cols

@timer
def build4(mode):
    name = '%s/%s_embed_tr.pkl'%(CACHE,mode)
    if os.path.exists(name):
        x = pd.read_pickle(name)
        xt = pd.read_pickle(name.replace('_tr','_te'))
        names = [i for i in x.columns]
        x = x.values
        xt = xt.values
    else:
        x,xt,names = get_pretrained_embedding(mode)
        pd.DataFrame(x,columns=names).to_pickle(name)
        pd.DataFrame(xt,columns=names).to_pickle(name.replace('_tr','_te'))
    #return x,xt,names 
    #good = ['CountryIdentifier','GeoNameIdentifier', 'AppVersion']
    #good = ['Census_OSEdition', 'CountryIdentifier', 'LocaleEnglishNameIdentifier', 'Census_OSBranch','SmartScreen','AVProductStatesIdentifier','AppVersion']
    #good = ['SmartScreen', 'CountryIdentifier', 'AVProductStatesIdentifier', 'AppVersion', 'Census_HasOpticalDiskDrive']
    #good = ['SmartScreen', 'CountryIdentifier', 'Census_OSInstallTypeName', 'Wdft_RegionIdentifier', 'AVProductStatesIdentifier', 'AppVersion', 'LocaleEnglishNameIdentifier', 'Census_FirmwareManufacturerIdentifier']
    good = ['SmartScreen', 'CountryIdentifier', 'AVProductStatesIdentifier', 'LocaleEnglishNameIdentifier', 'Census_OSInstallTypeName']
    cols = []
    newnames = []
    for c,col in enumerate(names):
        xcol = '_'.join(col.split('_')[1:-1]) 
        if xcol in good:
            cols.append(c)
            newnames.append(col)
    print('build4',x[:,cols].shape,xt[:,cols].shape)
    return x[:,cols],xt[:,cols],newnames

@timer
def build3(mode):
    tr_path,te_path = get_tr_te_paths(mode)
    gtr,_ = read_csv_hash_nvstring(tr_path)
    gte,_ = read_csv_hash_nvstring(te_path)
    gtr = rm_cols(gtr,[IDCOL])
    gte = rm_cols(gte,[IDCOL,YCOL])
    mtr_cols = ['SmartScreen','AVProductStatesIdentifier','AppVersion']
    print("# of mtr cols",len(mtr_cols))
    gtr,gte = mtr_encode(gtr,gte,mtr_cols)
    ptr = gtr.to_pandas()
    pte = gte.to_pandas()
    cols = [i for i in ptr.columns]
    del gtr,gte
    return ptr.values,pte.values,cols

def post_gdf(gtr,gte):
    tr = to_pandas(gtr).values
    te = to_pandas(gte).values
    cols = [i for i in gtr.columns]
    del gtr,gte
    return tr,te,cols

    """
    for col in gtr.columns:
        a,b = str(gtr[col].dtype),str(gte[col].dtype)
        if a!=b:
            gtr[col] = gtr[col].astype('float32')
            gte[col] = gte[col].astype('float32')
    """


# In[8]:


def get_model(names,num_class,model='nn'):
    model = eval('get_%s_model(names,num_class)'%model)
    return model

def get_categorical_columns(names):
    true_numerical_columns = get_true_numerical_columns()
    true_numerical_columns += get_new_numerical_columns()
    #true_numerical_columns += get_binary_columns()
    return [i for i in names if i not in true_numerical_columns and i.startswith('mtr_')==False and i.startswith('count_')==False]

def get_true_numerical_columns():
    true_numerical_columns = [
        'Census_ProcessorCoreCount',
        'Census_PrimaryDiskTotalCapacity',
        'Census_SystemVolumeTotalCapacity',
        'Census_TotalPhysicalRAM',
        'Census_InternalPrimaryDiagonalDisplaySizeInInches',
        'Census_InternalPrimaryDisplayResolutionHorizontal',
        'Census_InternalPrimaryDisplayResolutionVertical',
        'Census_InternalBatteryNumberOfCharges'
    ]
    return true_numerical_columns

def get_new_numerical_columns():
    cols = 'AvSigVersion_1,'.split(',')
    return cols

def get_xgb_cpu_params(names,num_class):
    print("# clases",num_class)
    #categorical = get_categorical_columns(names)
    params =  {
        'objective': 'binary:logistic', 
        'tree_method': 'hist', 
        'early_stopping_rounds':100,#None,
        'eta':0.1,
        'nthread': 16, 
        'folds':100,
        'watch':True,#False,
        'num_class':num_class,
        'num_round':12000,
        'max_depth': 10, 
        'silent':1,
        'subsample':0.8,
        'colsample_bytree': 0.7,
        'min_child_weight':10,
        'feature_names':names,
        'maximize':True,
        'eval_metric':METRIC,
        'verbose_eval':100,
    }
    return params

def get_xgb_gpu_params(names,num_class):
    params = get_xgb_cpu_params(names,num_class)
    params.update({'tree_method': 'gpu_hist',
    })
    return params

def get_xgb_params(names,num_class):
    #params = get_xgb_cpu_params(names,num_class)
    params = get_xgb_gpu_params(names,num_class)
    return params

def get_xgb_model(names,num_class):
    xgb_params = get_xgb_params(names,num_class)
    model = xgb_model(**xgb_params)
    return model

def get_bag_xgb_model(names,num_class):
    xgb_params = get_xgb_params(names,num_class)
    xgb_params.update({
        'watch':False,
        'early_stopping_rounds':None,
    })
    xgb_params['folds'] = 4
    xgb_params['num_round'] = 500
    #xgb_params['ranknorm'] = True
    print(xgb_params)
    model = bag_xgb_model(**xgb_params)
    return model

def get_bag_sk_lgb_model(names,num_class):
    lgb_params = get_sk_lgb_params(names,num_class)
    lgb_params['folds'] = 4
    #lgb_params['n_estimators'] =30000 
    model = bag_sk_lgb_model(**lgb_params)
    return model

def get_sk_lgb_model(names,num_class):
    lgb_params = get_sk_lgb_params(names,num_class)
    model = sk_lgb_model(**lgb_params)
    return model

def get_sk_lgb_params(names,num_class):
    params = {
        "max_depth":-1,
        "n_estimators":30000,
        "learning_rate":0.05,
        "num_leaves":2**11-1,
        "colsample_bytree":0.3,
        "objective":'binary',
        "eval_metric":'auc',
        "n_jobs":16,
        "verbose":100,
        "early_stopping_rounds":200
    }
    return params

def get_lgb_model(names,num_class):
    lgb_params = get_lgb_params(names,num_class)
    model = lgb_model(**lgb_params)
    return model

def get_bag_lgb_model(names,num_class):
    lgb_params = get_lgb_params(names,num_class)
    lgb_params['folds'] = 4
    lgb_params['num_round'] = 1000
    model = bag_lgb_model(**lgb_params)
    return model

def get_lgb_params(names,num_class):
    print("# clases",num_class)
    categorical = get_categorical_columns(names)
    print("categorical",categorical)
    lgbm_params =  {
        'num_threads':16,
        'boosting_type': 'gbdt',
        'objective': 'binary',
        'metric': METRIC,#None,#'multi_logloss',
        'num_round': 1000,
        'folds':1000,
        'num_class':num_class,
        "early_stopping_rounds":100,
        "feature_name": names,
        #'num_leaves': 16,
        'seed': 3,
        #'load_model':'backup/2018-05-29-19-56-01_cv_0.2215/lgb_sub.model',
    
        'learning_rate': 0.1,
        'max_bin':32,
        'subsample': .9,
        'colsample_bytree': 0.7,
        'reg_alpha': .01,
        'reg_lambda': .01,
        'min_split_gain': 0.01,
        'min_child_weight': 10,
        #'silent': -1,
        'max_depth': None,    

        'verbosity':-1,
        'importance_type': 'gain',
        'stratified':False,
        'verbose': -1,
        'verbose_eval':100,
        'categorical': categorical
    }
    return lgbm_params

def get_nn_model(names,num_class):
    from malnet import MalNN
    H = 256
    catfeas = len(get_categorical_columns(names))
    params = get_nn_params(H,catfeas,num_class)
    return MalNN(**params)

def get_nn_params(H,catfeas,num_class):
    params = {
        'classes':num_class,
        'embedding_size':4,
        'catfeas':catfeas,
        'stratified':True,
        'metric':METRIC,
        'save_path':'weights',
        #'load_path':'weights/r67.npy',
        'epochs':25 if MODE=='sub' else 100,
        'Hs':[H,H//2,H//4],
        'drop_prob':0.5,
        'early_stopping_epochs':5,
        'learning_rate':0.01,
        'batch_size':102400,
        #'verbosity':10,
        'folds':1000
    }
    return params

def get_sparse_nn_model(names,num_class):
    from sparse_nn import SparseNN
    H = 256
    catfeas = len(get_categorical_columns(names))
    params = get_nn_params(H,catfeas,num_class)
    #params['batch_size'] = 102400*2
    if MODE == 'sub':
        params['epochs']=30
    return SparseNN(**params)


# In[9]:
def get_ffm_model(names=None,num_class=None):
    params = {
        'mode': MODE,
        'cache':'%s/ffm1'%CACHE,
        'path': CACHE if MODE=='cv' else PATH,
        'fmap':'fmap.txt',
        'autostop':3,
        'metric':'auc',
        'threads':64,
        'embedding_size':4,
        'lr':0.1,
        'epochs':100,
        'lambda':0.0001,
    }
    if MODE == 'sub':
        params['autostop'] = params['epochs']
    model = FFM(**params)
    return model

@timer
def fit_predict(names,X,y,Xt,yt=None,model_name='lgb'):
    num_class = 1
    model = get_model(names,num_class,model_name)
    if model_name == 'ffm':
        yp = model.fit_predict()
        return yp,model
    if model_name in ['nn']:
        X,Xt = preprocess(X,Xt,names)
    if model_name=='nn' and MODE == 'sub':
         model.fit(X,y,va=[X,y]) 
    elif 'bag' not in model_name and yt is not None:
    #if yt is not None:
         model.fit(X,y,va=[Xt,yt])
    else:
         model.fit(X,y)#,va=None)
    yp = model.predict(Xt)
    return yp,model

@timer
def preprocess(X,Xt,names):
    scaler = StandardScaler()
    N = len(get_categorical_columns(names)) 
    X = pd.DataFrame(X).fillna(0).values
    Xt = pd.DataFrame(Xt).fillna(0).values
    X[:,N:] = scaler.fit_transform(X[:,N:])
    Xt[:,N:] = scaler.transform(Xt[:,N:]) 
    return X,Xt

def get_kmeans(mode,K=10):
    out_path = "%s/kmeans"%CACHE
    mkdir(out_path)
    name = '%s/%s_embed_tr.pkl'%(CACHE,mode)
    x = pd.read_pickle(name)
    xt = pd.read_pickle(name.replace('_tr','_te'))  

    gx = gd.DataFrame.from_pandas(x)
    gxt = gd.DataFrame.from_pandas(xt)

    name = '%s/fea%s_0_tr.pkl'%(CACHE,mode)
    x0 = pd.read_pickle(name)
    xt0 = pd.read_pickle(name.replace('_tr','_te'))

    #num_cols = 
    #for col in x0.columns:
    

def get_pretrained_embedding(mode):
    model = 'nn'
    b1 = [0,2]
    c1 = [1,1]
    X,Xt,y,yt,names,tr_id,te_id = build(mode,build_list=b1,cache=c1)
    model = get_model(names,1,'nn')
    X,Xt = preprocess(X,Xt,names) 
    model.params['get_embedding'] = 1
    if mode == 'cv':
        model.bestw = 'backup/2019-01-24-13-45-38_cv_0.7039/2019-01-24-13-33-57_MalNN_1_18_best.npy'
        #model.bestw = 'backup/2019-02-11-20-00-25_cv_0.7111/2019-02-11-17-48-23_MalNN_1_26_best.npy'
    else:
        model.bestw = 'backup/2019-01-24-13-55-48_cv_0.7039/2019-01-24-13-43-41_MalNN_0_23_best.npy'
        #model.bestw = 'backup/2019-02-12-00-17-05_cv_0.7111/2019-02-11-20-21-53_MalNN_5_25_best.npy'
    y = model.predict(X)
    yt = model.predict(Xt)
    print(X.shape,y.shape)
    print(Xt.shape,yt.shape)
    cats = names[:model.params['catfeas']]
    assert len(cats) == len(model.emd_size)
    names = []
    for c,e in zip(cats,model.emd_size):
        names.extend(['emb_%s_%d'%(c,i) for i in range(e)])
    assert len(names) == y.shape[1]
    return y,yt,names 

def main(mode,model='lgb'):
    cname = sys.argv[0].replace('.py','_%d.py'%GPU)
    cmd = 'cp %s %s'%(sys.argv[0],cname)
    os.system(cmd)
    global MODE
    MODE = mode 
    start = time.time()
    tag = '%s_%s_%d'%(mode,model,GPU)
    out = '%s.csv.gz'%tag
    #bl = [13,6,4,2]
    bl = [13,6,1,4]
    cl = [0]*len(bl)
    if mode == 'sub':
        cl = [0 if i in [4,9] else 1 for i in bl]
    #cl = [1,1,1,0]
    if model == 'ffm':
        bl,cl = [],[]
    if model == 'nn':
        bl = [0,6]
        cl = [1]*len(bl)
    if model == 'sparse_nn':
        bl,cl = [9],[0]
    X,Xt,y,yt,names,tr_id,te_id = build(mode,build_list=bl,cache=cl)
    yp,model = fit_predict(names,X,y,Xt,yt,model_name=model)
    try:
        model.get_importance("%s_importance.csv"%tag)
    except:
        pass
    if mode == 'cv':
        score = get_score(yt,yp,METRIC)
    else:
        score = -1 #None
    print(METRIC,score)
    if mode == 'sub':
        s = pd.DataFrame({IDCOL:te_id,YCOL:yp})
    else:
        s = pd.DataFrame({YCOL:yp})
        
    s.to_csv(out,index=False,compression='gzip',float_format='%.6f')
    duration = time.time()-start
    files = [i for i in os.listdir('.') if i.startswith(tag)]+[cname]
    write_log(duration,score,"{}".format(tag),mfiles=files)

@timer
def tohost(df):
    return df.to_pandas()

@timer
def test_tohost(mode='cv'):
    N=10000000
    M = 20
    gtr = gd.DataFrame()
    for i in range(M):
        gtr[i] = np.arange(N)
    start = time.time()
    df = gtr.to_pandas()
    duration = time.time()-start
    print('Dataframe shape: %s to_pandas: %.2f seconds'%(str(df.shape),duration))

if __name__ == '__main__':
    model = ['xgb','lgb','bag_xgb','bag_lgb','ffm','nn','sparse_nn','sk_lgb','bag_sk_lgb'][2]
    #main('cv',model)
    main('sub',model)
    #X,Xt,y,yt,names,tr_id,te_id = build(mode='sub',build_list=[4],cache=[0])
    #tr,te = load_AvSig(mode='sub')
    #tr = AvSig_to_timestamp(tr) 
    #te = AvSig_to_timestamp(te)
    #test_tohost()
