from config import *
import argparse
import json
from collections import Counter
from pathlib import Path
from utils import *

#生成词表
def generate_vocab(
    train_sample_path=TRAIN_SAMPLE_PATH,
    zh_vocab_path=ZH_VOCAB_PATH,
    en_vocab_path=EN_VOCAB_PATH,
    min_freq=2,
):
    if min_freq < 1:
        raise ValueError("min_freq 必须大于或等于1。")

    en_vocab = ['<pad>', '<unk>', '<sos>', '<eos>']
    zh_vocab = ['<pad>', '<unk>', '<sos>', '<eos>']
    
    zh_vocab_list = []
    en_vocab_list = []

    # 解析json文件
    with open(train_sample_path, encoding='utf-8') as file:
        lines = json.load(file)
        for en_sent, zh_sent in lines:
            # en_vocab_list += divided_en(en_sent)
            # zh_vocab_list += divided_zh(zh_sent)
            # 先统一为小写再统计，避免 The、the 等写入词表后变成重复项。
            # 同时过滤 Jieba 分词时可能产生的空白 token。
            en_vocab_list += [
                word.lower()
                for word in divided_en(en_sent)
                if word.strip()
            ]
            zh_vocab_list += [
                word.lower()
                for word in divided_zh(zh_sent)
                if word.strip()
            ]

    print('train_sample count:', len(lines))
    # 按次数生成词表，如果语料库够大，可以按最小次数，过滤掉生僻词。
    # MIN_FREQ = 2

    #统计列表中每个元素出现的次数，并按出现频率从高到低排序返回。
    # [('我', 3), ('爱', 2), ('你', 1)]
    zh_vocab_kv = Counter(zh_vocab_list).most_common()
    # zh_vocab += [k.lower() for k,v in zh_vocab_kv]
    zh_vocab += [
        word
        for word, count in zh_vocab_kv
        if count >= min_freq
    ]

    en_vocab_kv = Counter(en_vocab_list).most_common()
    # en_vocab += [k.lower() for k,v in en_vocab_kv]
    en_vocab += [
        word
        for word, count in en_vocab_kv
        if count >= min_freq
    ]

    print('en_vocab count:', len(en_vocab))
    print('zh_vocab count:', len(zh_vocab))
    print('min_freq:', min_freq)

    # 生成词表文件
    zh_vocab_path = Path(zh_vocab_path)
    en_vocab_path = Path(en_vocab_path)
    zh_vocab_path.parent.mkdir(parents=True, exist_ok=True)
    en_vocab_path.parent.mkdir(parents=True, exist_ok=True)

    with zh_vocab_path.open('w', encoding='utf-8') as file:
        file.write('\n'.join(zh_vocab))
        
    with en_vocab_path.open('w', encoding='utf-8') as file:
        file.write('\n'.join(en_vocab))

    print('en_vocab saved to:', en_vocab_path)
    print('zh_vocab saved to:', zh_vocab_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='根据训练集生成中英文词表。')
    parser.add_argument(
        '--train-path',
        default=TRAIN_SAMPLE_PATH,
        help='训练数据 JSON 路径。',
    )
    parser.add_argument(
        '--en-output',
        default=EN_VOCAB_PATH,
        help='英文词表输出路径。',
    )
    parser.add_argument(
        '--zh-output',
        default=ZH_VOCAB_PATH,
        help='中文词表输出路径。',
    )
    parser.add_argument(
        '--min-freq',
        type=int,
        default=2,
        help='进入词表所需的最低出现次数，默认值为2。',
    )
    args = parser.parse_args()

    generate_vocab(
        train_sample_path=args.train_path,
        zh_vocab_path=args.zh_output,
        en_vocab_path=args.en_output,
        min_freq=args.min_freq,
    )