import jieba
from wordcloud import WordCloud
import matplotlib.pyplot as plt
from collections import Counter
import re
from PIL import Image
import numpy as np

# 1. 读取政府工作报告文本
def read_report(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        text = f.read()
    return text

# 2. 文本预处理和分词
def preprocess_and_segment(text):
    # 去除标点符号和数字
    text = re.sub(r'[^\w\s]', '', text)
    text = re.sub(r'\d+', '', text)
    
    # 使用jieba进行分词
    words = jieba.lcut(text)
    
    # 加载停用词表
    stopwords = set()
    with open('stopwords.txt', 'r', encoding='utf-8') as f:
        for line in f:
            stopwords.add(line.strip())
    
    # 过滤停用词和单字词
    filtered_words = [word for word in words if len(word) > 1 and word not in stopwords]
    
    return filtered_words

# 3. 统计词频
def count_word_frequency(words, top_n=50):
    word_counts = Counter(words)
    return word_counts.most_common(top_n)

# 4. 生成词云
def generate_word_cloud(word_freq, output_path='wordcloud.png', mask_image=None):
    if mask_image:
        # 使用蒙版图片
        mask = np.array(Image.open(mask_image))
        wc = WordCloud(
            font_path='simhei.ttf',  # 使用黑体
            background_color='white',
            max_words=200,
            mask=mask,
            contour_width=3,
            contour_color='steelblue'
        )
    else:
        wc = WordCloud(
            font_path='simhei.ttf',
            background_color='white',
            width=800,
            height=600,
            max_words=200
        )
    
    # 生成词云
    wc.generate_from_frequencies(dict(word_freq))
    
    # 显示词云
    plt.figure(figsize=(10, 8))
    plt.imshow(wc, interpolation='bilinear')
    plt.axis('off')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.show()

# 主程序
def main():
    # 文件路径
    report_file = 'government_report.txt'  # 政府工作报告文本
    stopwords_file = 'stopwords.txt'      # 停用词表
    mask_image = 'china_map.png'          # 可选:中国地图轮廓图片
    
    # 1. 读取报告
    report_text = read_report(report_file)
    
    # 2. 预处理和分词
    words = preprocess_and_segment(report_text)
    
    # 3. 统计词频
    word_freq = count_word_frequency(words, top_n=50)
    
    print("Top 50 热词及频率:")
    for word, freq in word_freq:
        print(f"{word}: {freq}")
    
    # 4. 生成词云
    generate_word_cloud(word_freq, mask_image=mask_image)

if __name__ == '__main__':
    main()