import keras
import numpy as np

from keras.layers import Dense, LSTM
from keras.utils import to_categorical
from keras.models import Sequential
from keras.models import load_model
from keras.callbacks import ModelCheckpoint
from keras.utils.np_utils import to_categorical
from keras.layers import Masking, Embedding, Flatten
from sklearn import preprocessing
from transferAPI import wordModel
import os
import pandas as pd

from config import *

le = preprocessing.LabelEncoder()

sequenceLength = 5000
nb_lstm_outputs = 30  # 神经元个数
nb_time_steps = sequenceLength  # 时间序列长度

modelname = 'models/thread_lstm_model.h5'
outClasses = 6

nb_input_vector = 16
def transformLabel(label):
    return to_categorical(label,outClasses)


def generate_arrays(df,batch_size):
    x,y = df.shape
    t =  np.random.RandomState(1675)
    cnt = int(t.rand()*x)
    while 1:
        X = []
        Y = []
        for index,row in  df.iloc[cnt:cnt+batch_size,:].iterrows():
            Y.append(transformLabel(row['label']))
            tempx = []
            bList = eval(row['apilist'])
            for api in bList:
                if api in wordModel:
                    tempx.append(wordModel[api])
            X.append(tempx)

        X = keras.preprocessing.sequence.pad_sequences(X, maxlen=sequenceLength, dtype='int32', value=-1)
        X = X.reshape(batch_size, sequenceLength, nb_input_vector)
        cnt=(batch_size+cnt)%x
        yield (np.array(X), np.array(Y))

def str2list(apistr):
    result = eval(apistr)
    return result



def checkData():
    df = pd.read_csv('data/train_thread_list.csv')
    model = Sequential()
    model.add(Masking(mask_value=-1, input_shape=(sequenceLength, nb_input_vector)))

    model.add(LSTM(units=nb_lstm_outputs, input_shape=(nb_time_steps, nb_input_vector)))
    # model.add(Flatten())
    model.add(Dense(outClasses, activation='softmax'))

    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    modelcheck = ModelCheckpoint(filepath=modelname)
    model.fit_generator(generate_arrays(df,batch_size=128),steps_per_epoch=1000,epochs=10,verbose=1,nb_worker=1,callbacks=[modelcheck])
    # model.fit(x, y, epochs=30, batch_size=128, verbose=1, callbacks=[modelcheck])


checkData()
