import pickle
import os, argparse
parser = argparse.ArgumentParser(description='Chinese NER task')
parser.add_argument('--train_data', type=str, default='data_path', help='train data source')
args = parser.parse_args()


def read_corpus(corpus_path):
    data = []
    with open(corpus_path, encoding='utf-8') as fr:
        lines = fr.readlines()
    sent_, tag_ = [], []
    for line in lines:
        if line != '\n':
            [char, label] = line.strip().split()
            sent_.append(char)
            tag_.append(label)
        else:
            data.append((sent_, tag_))
            sent_, tag_ = [], []
    return data

def vocab_build(corpus_path, save_path,min_count):
    data = read_corpus(corpus_path)
    word2id = {}
    for sent_, tag_ in data:
        for word in sent_:
            if word.isdigit():
                word = '<NUM>'
            elif ('\u0041' <= word <='\u005a') or ('\u0061' <= word <='\u007a'):
                word = '<ENG>'
            if word not in word2id:
                word2id[word] = [len(word2id)+1, 1]
            else:
                word2id[word][1] += 1
    low_freq_words = []
    for word, [word_id, word_freq] in word2id.items():
        if word_freq < min_count and word != '<NUM>' and word != '<ENG>':
            low_freq_words.append(word)
    for word in low_freq_words:
        del word2id[word]

    new_id = 1
    for word in word2id.keys():
        word2id[word] = new_id
        new_id += 1
    word2id['<UNK>'] = new_id
    word2id['<PAD>'] = 0

    print(len(word2id))
    with open(save_path, 'wb') as fw:
        pickle.dump(word2id, fw)

if __name__ == '__main__':
    file_list = os.path.join('.', args.train_data, 'word_npy')
    save_path = os.path.join('.', args.train_data, 'word2id_2.pkl')
    vocab_build(file_list, save_path, 1)