from sklearn.feature_extraction import DictVectorizer
from sklearn.model_selection import train_test_split 
from sklearn.metrics import roc_auc_score
from sklearn.metrics import accuracy_score 
from sklearn.metrics import log_loss

import pandas as pd


import lightgbm as lgb 
from datetime import datetime
import os

phase1 = '/home/malware/trace1_test/dex.txt'
phase2 = '/home/OpenCode/malware/phase2/trace1_test/dex.txt'
def line2data(line):
    return line.strip().split(';')


def readFile(filename):        
    dataList = []
    shaList = []
    inf = open(filename)
    keys = line2data(inf.readline())
    for line in inf.readlines():
        vals = line2data(line)
        newData = {}
        for i in range(len(keys)):
            newData[keys[i]] = vals[i]           
        dataList.append(newData)
        shaList.append(newData['sha1'])
    return dataList,shaList

def processDex(dataList):
    train = []
    for data in dataList:
        newData = {}
        permList = data['permission_list'].split(',')
        # newData['perms']  = len(permList)        
        for key in data:
            if key not in ['sha1','permission_list','opcode_count']:
                newData[key] = int(data[key])
        train.append(newData)
    return train

def binLabel(typeList):
    result = []
    for data in typeList:
        result.append(int(data['safe_type']))
    return result

def multiLabel(typeList):
    result = []
    for data in typeList:
        result.append(int(data['family_id'])) 
    return result

def readDataFromDir(dirname,multi=False):
    dexFile = os.path.join(dirname,'dex.txt')
    typeFile = os.path.join(dirname,'safetype.txt')
    v = DictVectorizer()

    dataList,nameList = readFile(dexFile)
    trainList = processDex(dataList)
    x = v.fit_transform(trainList)
    
    typeList,_ = readFile(typeFile)
    if multi:
        y = multiLabel(typeList)
    else:
        y = binLabel(typeList)
    return x,y,nameList,trainList

def formatLabel(y_pred,multi=False):
    if multi:
        newLabel = []
        classLen = len(y_pred[0])
        for i in range(len(y_pred)):
            label = 0
            for j in range(classLen):
                if y_pred[i][j]>0.5:
                    label = j
            newLabel.append(label)
        y_pred = newLabel
    else:
        for i in range(len(y_pred)):       
            if y_pred[i]>=.5:       # setting threshold to .5
                y_pred[i]=1
            else:  
                y_pred[i]=0
    return y_pred

def checkModel(y_test,y_pred,multi=False):

    y_pred = formatLabel(y_pred)

    #calculating accuracy
    accuracy_lgbm = accuracy_score(y_pred,y_test)
    print(accuracy_lgbm)
   
    if not multi:
        auc_lgbm =  roc_auc_score(y_test,y_pred)
        print(auc_lgbm)


def checkMultiLabel():
    x,y,shaList,_ = readDataFromDir('/home/malware/trace1_train1',multi=True)

    x_train,x_test,y_train,y_test=train_test_split(x,y,test_size=.1)
    train_data=lgb.Dataset(x_train,label=y_train)

    param = {'num_leaves':150, 'objective':'multiclass','max_depth':7,'learning_rate':.05,'max_bin':200,'num_class':6}
    param['metric'] = ['auc']

    #Here we have set max_depth in xgb and LightGBM to 7 to have a fair comparison between the two.

    #training our model using light gbm
    num_round=50
    start=datetime.now()
    lgbm=lgb.train(param,train_data,num_round)
    stop=datetime.now()

    #Execution time of the model
    execution_time_lgbm = stop-start
    print('train time:',execution_time_lgbm)

    #predicting on test set
    y_pred=lgbm.predict(x_test)
    checkModel(y_test,y_pred,multi=True)
    x1,y1,_,_ = readDataFromDir('/home/malware/trace1_train2')
    checkModel(y1,lgbm.predict(x1))

def makeModel(x_train,y_train,multi=False):
    train_data=lgb.Dataset(x_train,label=y_train)

    if multi:
        param = {'num_leaves':50, 'objective':'multiclass','max_depth':7,'learning_rate':.05,'max_bin':200,'num_class':6}
        param['metric'] = ['auc']
    else:
        param = {'num_leaves':50, 'objective':'binary','max_depth':7,'learning_rate':.05,'max_bin':200}
        param['metric'] = ['auc']

    num_round=1000
    lgbm=lgb.train(param,train_data,num_round)
    lgbm.save_model('model.txt', num_iteration=lgbm.best_iteration)

    return lgbm




def predictData(filename,multi=False,append=False):
    v = DictVectorizer()
    x1,y1,_,trainX1 = readDataFromDir('/home/malware/trace1_train1',multi=multi)
    x2,y2,_,trainX2 = readDataFromDir('/home/malware/trace1_train2',multi=multi)
    x3,y3,_,trainX3 = readResultLabel(multi)
    trainX = trainX1+trainX2
    if append:
        trainX = trainX+trainX3
    x =  v.fit_transform(trainX)

    if append:
        y = y1+y2+y3
    else:
        y = y1 + y2
    print(len(trainX),len(y))
    model = makeModel(x,y,multi=multi)
    dataList, nameList = readFile(filename)
    trainList = processDex(dataList)
    x_test = v.fit_transform(trainList)
    y_pred = model.predict(x_test)
    if multi:
        outf = open('data/trace1_lgbm_output.txt', 'w+')
        outf.write('sha1,0,1,2,3,4,5\n')
        for i in range(len(nameList)):
            result = y_pred[i]
            outf.write(
                '%s,%f,%f,%f,%f,%f,%f\n' % (nameList[i], result[0], result[1], result[2], result[3], result[4], result[5]))
    return nameList,formatLabel(y_pred,multi=multi)


def solutionPhase1():
    nameList,y1 = predictData(phase2)
    _,y2 = predictData(phase2,multi=True)
    with open('/home/malware/phase1_trace1.csv','w+') as outf:
        for i in range(len(nameList)):
            outf.write(nameList[i]+','+str(int(y1[i]))+','+str(y2[i])+'\n')

def solutionPhase2():
    nameList,y1 = predictData(phase2,multi=False,append=False)
    _,y2 = predictData(phase2,multi=True,append=False)
    with open('/home/malware/trace1_test/trace1.csv','w+') as outf:
        for i in range(len(nameList)):
            outf.write(nameList[i]+','+str(int(y1[i]))+','+str(y2[i])+'\n')


def readResultLabel(multi):
    df2 = pd.read_csv('/home/malware/phase1_trace1.csv',header=None)
    df2.columns = ['sha1','safe_type','family_id']

    v = DictVectorizer()
    dataList, nameList = readFile(phase1)
    trainList = processDex(dataList)
    x = v.fit_transform(trainList)

    if multi:
        y = list(df2.family_id)
    else:
        y = list(df2.safe_type)
    return x, y, nameList, trainList


if __name__=='__main__':
    # solutionPhase1()
    solutionPhase2()
    




