import argparse
from encodings import utf_8
from importlib.resources import path
import thulac
import json
import os

from tqdm import tqdm
from keras.preprocessing.text import Tokenizer

def loaddata(filepath):
    filelist = os.listdir(filepath)
    path = os.path.join(filepath,filelist[0])
    if(filelist[0].endswith(".txt")):
        print("FileProcessing:{}".format(filelist[0]))
        with open(path,"r",encoding="utf-8") as f:
            yield f.read()
    #for filename in filelist:
    #    path = os.path.join(filepath,filename)
    #    if(filename.endswith(".txt")):
    #        print("FileProcessing:{}".format(filename))
    #        with open(path,"r",encoding="utf-8") as f:
    #            yield f.read()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--raw_data_path', default='../data/', type=str, required=False, help='原始训练语料')
    parser.add_argument('--vocab_file', default='vocab_processed.txt', type=str, required=False, help='生成vocab链接')
    parser.add_argument('--vocab_size', default=50000, type=int, required=False, help='词表大小')
    args = parser.parse_args()

    tokenizer = Tokenizer(num_words=args.vocab_size)
    print('args:\n' + args.__repr__())
    print('This script is extremely slow especially for large corpus. Take a break.')

    textdata = loaddata(args.raw_data_path)
    tokenizer.fit_on_texts(textdata)

    vocab = list(tokenizer.index_word.values())
    #pre = ['[SEP]', '[CLS]', '[MASK]', '[PAD]', '[UNK]']
    pre = ['[SEP]']
    vocab = pre + vocab
    with open(args.vocab_file, 'w') as f:
        for word in vocab[:args.vocab_size + 5]:
            f.write(word + '\n')


if __name__ == "__main__":
    main()
