# -*- coding:utf-8 -*-
import os
import numpy as np
from numpy import random
read_filePath = 'D:/MalwareDB/tokenized/'  # this is the root path
write_filePath = 'D:/MalwareDB/MalwareDB/'
 
if __name__ == "__main__":
    sentences = []
    sentence = []
    file_num = 0
    for i in os.listdir(read_filePath):
        print(i)
        file_num = file_num + 1
        read_file = open(read_filePath + i, 'r', encoding='UTF-8')
        for line in read_file:
            line = line.strip()
            if line == '':
                if not sentence:
                    continue
                sentences.append(sentence)  # 全部的句子
                sentence = []
            else:
                sentence.append(line.split()[0] + ' ' + line.split()[-1])  # 代表每个句子

    print('There are %d files' % file_num)
    sentence_num = len(sentences)  # 所有句子的数量
    print('There is %d sentences data' % sentence_num)
    piece = sentence_num // 10
    sentences = np.array(sentences)  # 将句子信息转化为numpy变量
    sentence_idx = np.arange(sentence_num)  # 构造索引列表
    random.shuffle(sentence_idx)  # 打乱索引列表

    train_indx = sentence_idx[0: 8 * piece]
    print('The train sentence_num is %d' % (8 * piece))
    train_sentences = sentences[train_indx].tolist()

    valid_indx = sentence_idx[8 * piece: 9 * piece]
    print('The valid sentence_num is %d' % piece)
    valid_sentences = sentences[valid_indx].tolist()

    test_indx = sentence_idx[9 * piece:]
    print('The test sentence_num is ', sentence_num - 9*piece)
    test_sentences = sentences[test_indx].tolist()

    for sentence_item, file_str in zip([train_sentences, valid_sentences, test_sentences], ['train.txt', 'valid.txt', 'test.txt']):
        file_wrt = write_filePath + file_str
        write = open(file_wrt, "a", encoding='utf-8')
        for each_sentence in sentence_item:
            for word_label in each_sentence:
                write.write(word_label+'\n')
            write.write('\n')
        write.close()






