# -*- coding: utf-8 -*-
"""
Created on Fri Sep 10 14:28:06 2021

@author: arash
"""
import numpy as np
import pickle
import torch, glob
from torch.utils.data import TensorDataset, Dataset, IterableDataset, DataLoader
from torch.nn.utils.rnn import pad_packed_sequence, pad_sequence
from random import sample

def my_collate(batch):
    # batch contains a list of tuples of structure (sequence, target)
    #print(batch[0])
    #[print(item[1]) for item in batch]
    data = [item[0] for item in batch]
    #[print(np.shape(item[0])) for item in batch]
    data = pad_sequence(data)
    targets = [item[1] for item in batch]
    #print(targets)
    return [data, torch.Tensor(targets)]

class CustomIterableDataset(IterableDataset):

    def __init__(self, data_link):

        #Store the filename in object's memory
        self.leng_max=6444
        self.vocab, self.index_word=pickle.load(open('vocab2.p','rb'))
        self.N=len(self.vocab)
        self.data_link = data_link
        da=pickle.load(open(self.data_link,'rb'))
        print(len(da))
        self.file_itr =iter(da)
        #And that's it, we no longer need to store the contents in the memory

    def preprocess(self, text):

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

        return torch.Tensor(temp)

    def line_mapper(self, line):
        
        #Splits the line into text and label and applies preprocessing to the text
        #print(line[0])
        if line[0]==0:
            label=0
        else:
            label=1
        text=line[1]
        text = self.preprocess(text)
        #if len(text)>0:
        #print('*',label)
        #print(torch.Tensor([label]))
        return text, torch.Tensor([label])#.long().cuda()


    def __iter__(self):

        #Create an iterator
        
        
        
        #Map each element using the line_mapper
        mapped_itr = map(self.line_mapper, self.file_itr)
        #print(list(mapped_itr))   
        return mapped_itr

class CustomDataset(Dataset):

    def __init__(self, data_link = 'training_processed_LLVM3.p'):

        #Store the filename in object's memory
        self.leng_max=6444
        self.vocab, self.index_word=pickle.load(open('vocab2.p','rb'))
        self.N=len(self.vocab)
        self.data_link=data_link
        self.data = pickle.load(open(self.data_link,'rb'))
        self.len = len(self.data)
        #And that's it, we no longer need to store the contents in the memory

    def randdata(self):
        None
    def __getitem__(self, index):
        
        line=self.data[index]
        #print(line[0])
        if line[0][0]==0:
            label=0
        else:
            label=1
        text=line[1]
        text = self.preprocess(text)
        return text, torch.Tensor([label])

    def __len__(self):
        return self.len


    def preprocess(self, text):

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

        return torch.Tensor(temp)



class CustomDatasetBalanced(Dataset):

    def __init__(self, data_link = 'training_processed_LLVM4.p'):

        #Store the filename in object's memory
        self.leng_max=6444
        self.vocab, self.index_word=pickle.load(open('vocab2.p','rb'))
        self.N=len(self.vocab)
        self.data_link=data_link
        self.vulner, self.good = pickle.load(open(self.data_link,'rb'))        
        self.len = 2*len(self.vulner)
        #And that's it, we no longer need to store the contents in the memory

    def randdata(self):
        self.data=self.vulner.copy()
        #X=self.good[np.random.randint(0,len(self.good),len(self.vulner))]
        X=sample(self.good,len(self.vulner))
        [self.data.append(x) for x in X]
    
    def __getitem__(self, index):
        line=self.data[index]
        if line[0][0]==0:
            label=0
        else:
            label=1
        text=line[1]
        text = self.preprocess(text)
        return text, torch.Tensor([label])

    def __len__(self):
        return self.len


    def preprocess(self, text):

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

        return torch.Tensor(temp)
    
    
    
    
    
class LineDataset(Dataset):

    def __init__(self, datalink):

        #Store the filename in object's memory
        self.leng_max=6444
        self.vocab, self.index_word=pickle.load(open('vocab2.p','rb'))
        self.N=len(self.vocab)
        #self.data_link=data_link
        #self.files=glob.glob('training_lined*.p')
        #print(self.files)
        # self.len=0
        # for f in self.files:
        #     d = pickle.load(open(f,'rb'))
        #     self.len+=len(d)
        #     del d
        # self.data = pickle.load(open(self.files[0],'rb'))
        self.data = pickle.load(open(datalink,'rb'))
        self.len=len(self.data)
        self.file_counter=1
        self.current_len=len(self.data)
        
        #And that's it, we no longer need to store the contents in the memory

    def randdata(self):
        None
    def __getitem__(self, index):
        #print('next')
        # if index>self.current_len:
        #     print('****************************new file')
        #     del self.data
        #     self.data = pickle.load(open(self.files[self.file_counter],'rb'))
        #     self.file_counter+=1
        #     self.current_len+=len(self.data)
        
        line=self.data[index]
        text = self.preprocess(line[0])
        return text, line[1], line[2]
        

    def __len__(self):
        return self.len


    def preprocess(self, s):

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

        return torch.Tensor(a)
    


class LineDatasetBalanced(Dataset):

    def __init__(self, data_link = 'training_processed_LLVM4.p'):

        #Store the filename in object's memory
        self.leng_max=6444
        self.vocab, self.index_word=pickle.load(open('vocab2.p','rb'))
        self.N=len(self.vocab)
        self.data_link=data_link
        self.data = pickle.load(open(self.data_link,'rb'))        
        self.len = len(self.data)
        #And that's it, we no longer need to store the contents in the memory

    # def randdata(self):
    #     self.data=self.vulner.copy()
    #     #X=self.good[np.random.randint(0,len(self.good),len(self.vulner))]
    #     X=sample(self.good,len(self.vulner))
    #     [self.data.append(x) for x in X]
    
    def __getitem__(self, index):
        line=self.data[index]
        # if line[0][0]==0:
        #     label=0
        # else:
        #     label=1
        text=line[0]
        #print(text)
        text = self.preprocess(text)
        #print(text.size(), line[1].size(), line[2].view(-1).size())
        return text, line[1], line[2].view(-1)

    def __len__(self):
        return self.len


    def preprocess(self, text):

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

        return torch.Tensor(a)
    
    
    





    
class LineIterableDataset(IterableDataset):

    def __init__(self, data_link):

        #Store the filename in object's memory
        self.leng_max=6444
        self.vocab, self.index_word=pickle.load(open('vocab2.p','rb'))
        self.N=len(self.vocab)
        self.data_link = data_link
        da=pickle.load(open(self.data_link,'rb'))
        print(len(da))
        self.file_itr =iter(da)
        #And that's it, we no longer need to store the contents in the memory

    def preprocess(self, text):

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

        return torch.Tensor(temp)

    def line_mapper(self, line):
        
        #Splits the line into text and label and applies preprocessing to the text
        #print(line[0])
        if line[0]==0:
            label=0
        else:
            label=1
        text=line[1]
        text = self.preprocess(text)
        #if len(text)>0:
        #print('*',label)
        #print(torch.Tensor([label]))
        return text, torch.Tensor([label])#.long().cuda()


    def __iter__(self):

        #Create an iterator
        
        
        
        #Map each element using the line_mapper
        mapped_itr = map(self.line_mapper, self.file_itr)
        #print(list(mapped_itr))   
        return mapped_itr
    