import os
import re
from collections import Counter
from multiprocessing import Pool, Manager

def process_file(file_path, shared_counter):
    """处理单个文件，统计词频"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            text = f.read().lower()
            # 简单分词，移除标点
            words = re.findall(r'\b\w+\b', text)
            # 更新共享 Counter
            shared_counter.update(words)
    except Exception as e:
        print(f"Error processing {file_path}: {e}")

def main():
    # 获取 data 目录下所有 .txt 文件
    data_dir = 'data'
    files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith('.txt')]
    
    # 使用 Manager 创建共享 Counter
    with Manager() as manager:
        shared_counter = manager.dict(Counter())
        
        # 创建进程池
        with Pool(processes=4) as pool:  # 可调整进程数
            # 分发任务给进程池
            for file_path in files:
                pool.apply_async(process_file, args=(file_path, shared_counter))
            # 关闭池并等待所有进程完成
            pool.close()
            pool.join()
        
        # 转换为普通 Counter 以获取结果
        final_counter = Counter(dict(shared_counter))
    
    # 输出前 10 高频词
    print("Top 10 高频词:")
    for word, count in final_counter.most_common(10):
        print(f"{word}: {count}")

if __name__ == '__main__':
    main()