# -*- coding: utf-8 -*-
"""
Created on Thu Sep 30 16:46:40 2021

@author: arash
"""

import numpy as np
import pickle
import torch
from torch.utils.data import TensorDataset, Dataset, IterableDataset, DataLoader
from torch.nn.utils.rnn import pad_packed_sequence, pad_sequence
from random import sample



device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load('model2N12.pt').eval().cuda()
# model=nn.DataParallel(model, device_ids=[0,1,2,3], output_device=[0])
# model.to(device)

def preprocess(text, N, vocab):

    ### Do something with text here
    #text_pp = text.lower().strip()
    ###
    temp=[]
    
    for s in text:
        a=np.zeros((N,))
        ss=s.split(' ')
        for y in ss:
            try:
                a[int(vocab[y])]=1.0
            except:
                print(y)
        temp.append(a) 

    return torch.Tensor(temp)




vocab, index_word=pickle.load(open('vocab2.p','rb'))

#vulner = pickle.load(open('training_lineddetection_LLVM.p','rb')) 


N=20086

data=[]
temp=[]
vul_counter=0
good_counter=0
count=0
print(len(vulner))
co=0
for X in vulner:
    if co%5000==0:
        print(co)
        print(len(data))
    
    if len(X[1])<251:
        co+=1
        err=X[0]
        L=np.zeros(len(X[1])).astype(int)
        L[err]=1
        L=torch.Tensor(L).int()
        local_batch=preprocess(X[1], N, vocab)
        local_batch=local_batch.view(-1,local_batch.size()[0],local_batch.size()[1])
        #print(local_batch.size())
        model.module.window=local_batch.size()[0]
        H=model.module.input(local_batch.cuda())
        y,(h, c)=model.module.rnn(H, (model.module.initHidden(),model.module.initHidden()))
        
        for i in range(len(X[1])):
            #print(L[i].item())
            if L[i].item()==1:
                #print(vul_counter)
                vul_counter+=1
                data.append([X[1][i],y[0,i,:],L[i]])
            elif L[i].item()==0:
                good_counter+=1
                temp.append([X[1][i],y[0,i,:],L[i]])
            
        #print('****************',len(data))
        if len(temp)>8000:
            for _ in range(2000):
                temp.pop(np.random.randint(0,len(temp)))
                
                            
            
        if len(data)>2000:
            ind=np.random.randint(0,len(temp),5000).tolist()
            [data.insert(np.random.randint(0,5000),temp[x]) for x in ind]
            print('***********save')
            nam='./line/training_lined' + str(count) + '.p'
            pickle.dump(data,open(nam,'wb'))   
            count+=1
            data=[]
            temp=[]
            
        #print(y[0,:,:].size())
        #print(L.size())
        
        #torch.cat((y[:,0,500:],y[:,-1,:500]),1)

nam='./line/training_lined' + str(count) + '.p'
pickle.dump(data,open(nam,'wb'))   
count+=1
data=[]
    
    
    
    
    
    
    
    
    
    
    
    
    


