import os
import lightgbm as lgb
from config import *
from tfidfModel import *

# staticFeatureFile = '/home/malware/tc/malwaretc/file_features.csv'
staticFeatureFile = '/home/malware/tc/test/file_features.csv'


def fileFeatures(inputfile):
    df = pd.read_csv(inputfile)
    x = df[['max_thread_len', 'min_thread_len', 'mean_thread_len', 'thread_num']]
    fileIDs = df['file_id']
    return x, fileIDs


def predict():
    x, fileIds = fileFeatures(staticFeatureFile)
    row,col = x.shape
    tfFeatures = loadTFFeatures()
    print(row)
    x = pd.concat([x, tfFeatures.iloc[-row:,:]], axis=1)
    lgbm = lgb.Booster(model_file='model.txt')
    y_pred = lgbm.predict(x) #pd.DataFrame()
    out1 = open('predict.csv','w+')
    out2 = open('predict_detail.csv', 'w+')
    out1.write('file_id,prob0,prob1,prob2,prob3,prob4,prob5\n')
    for i in range(row):
        y = y_pred[i]
        result = [1e-7]*6
        maxIndex=0
        maxVal = y[0]
        for index,val in enumerate(y):
            if maxVal<val:
                maxIndex = index
                maxVal = val
        result[maxIndex] = 1-1e-7
        out1.write('%d,%s\n'%(fileIds[i],','.join(['%.7f'%s for s in result])))
        out2.write('%d,%s\n'%(fileIds[i],','.join(['%.7f'%s for s in y])))


    # df = pd.concat([fileIds,y_pred],axis=1)


if __name__ == '__main__':
    predict()
