import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import TruncatedSVD
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.patches import Ellipse

# 设置中文字体和样式
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
sns.set_style("whitegrid")

# 示例文档集
documents = [
    "机器学习 深度学习 神经网络 人工智能",
    "Python 编程 代码 算法 数据结构",
    "数学 统计 概率 线性代数 微积分",
    "机器学习 Python 算法 统计 数据科学",
    "深度学习 神经网络 Python 编程",
    "数学 线性代数 算法 数据结构"
]

# 创建文档-词语矩阵 (TF-IDF)
vectorizer = TfidfVectorizer()
doc_term_matrix = vectorizer.fit_transform(documents)
terms = vectorizer.get_feature_names_out()

print("原始文档-词语矩阵维度:", doc_term_matrix.shape)
print("词语:", terms)

# 使用截断SVD进行潜在语义分析
n_components = 3  # 选择3个主题
svd = TruncatedSVD(n_components=n_components, random_state=42)
doc_topic_matrix = svd.fit_transform(doc_term_matrix)
topic_term_matrix = svd.components_

print(f"\n奇异值(主题重要性): {svd.singular_values_}")
print(f"解释方差比例: {svd.explained_variance_ratio_}")
print(f"累计解释方差: {np.cumsum(svd.explained_variance_ratio_)}")

# 显示文档在主题空间中的坐标
doc_topic_df = pd.DataFrame(doc_topic_matrix, 
                           columns=[f'主题{i+1}' for i in range(n_components)],
                           index=[f'文档{i+1}' for i in range(len(documents))])

# 图片1：主题重要性分析图
plt.figure(figsize=(10, 6))
singular_values = svd.singular_values_
explained_variance = svd.explained_variance_ratio_
cumulative_variance = np.cumsum(explained_variance)

x = range(1, len(singular_values) + 1)
plt.bar(x, singular_values, alpha=0.7, color='skyblue', label='奇异值')
plt.plot(x, cumulative_variance * max(singular_values), 'ro-', linewidth=2, markersize=8, label='累计解释方差')
plt.xlabel('主题编号')
plt.ylabel('奇异值 / 解释方差')
plt.title('SVD主题重要性分析：奇异值分布与能量累积', fontsize=14, fontweight='bold')
plt.legend()
plt.grid(True, alpha=0.3)

# 为右侧y轴添加解释方差比例
ax2 = plt.gca().twinx()
ax2.plot(x, cumulative_variance, 'go--', linewidth=2, markersize=6, label='累计解释方差比例')
ax2.set_ylabel('解释方差比例')
ax2.set_ylim(0, 1.1)
ax2.legend(loc='lower right')

plt.tight_layout()
plt.show()

# 图片2：文档-主题关联热力图
plt.figure(figsize=(10, 6))
sns.heatmap(doc_topic_df, annot=True, cmap='YlOrRd', center=0,
            cbar_kws={'label': '主题权重'}, fmt='.3f')
plt.xlabel('潜在主题')
plt.ylabel('文档编号')
plt.title('文档-主题关联热力图：发现文档的潜在语义结构', fontsize=14, fontweight='bold')
plt.tight_layout()
plt.show()

# 图片3：文档在主题空间中的分布
plt.figure(figsize=(10, 8))
colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown']
document_labels = [f'文档{i+1}' for i in range(len(documents))]

for i, (x, y) in enumerate(zip(doc_topic_matrix[:, 0], doc_topic_matrix[:, 1])):
    plt.scatter(x, y, color=colors[i], s=150, alpha=0.7, label=document_labels[i])
    plt.annotate(document_labels[i], (x, y), xytext=(8, 8), 
                textcoords='offset points', fontsize=11, 
                bbox=dict(boxstyle="round,pad=0.3", facecolor=colors[i], alpha=0.2))

plt.axhline(y=0, color='gray', linestyle='--', alpha=0.5)
plt.axvline(x=0, color='gray', linestyle='--', alpha=0.5)
plt.xlabel('主题1 (技术主题)', fontsize=12)
plt.ylabel('主题2 (数学主题)', fontsize=12)
plt.title('文档在潜在语义空间中的分布：可视化文档聚类', fontsize=14, fontweight='bold')
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.grid(True, alpha=0.3)

# 添加聚类椭圆示意
ellipse1 = Ellipse(xy=(0.6, 0.1), width=0.4, height=0.2, 
                   angle=0, color='red', alpha=0.1, label='技术文档簇')
ellipse2 = Ellipse(xy=(-0.2, 0.7), width=0.3, height=0.15, 
                   angle=0, color='blue', alpha=0.1, label='数学文档簇')
plt.gca().add_patch(ellipse1)
plt.gca().add_patch(ellipse2)

plt.tight_layout()
plt.show()

# 图片4：主题-词语关联图
plt.figure(figsize=(12, 8))

# 选择每个主题最重要的3个词语
top_terms_per_topic = {}
for topic_idx in range(n_components):
    top_indices = topic_term_matrix[topic_idx].argsort()[-3:][::-1]
    top_terms = [(terms[i], topic_term_matrix[topic_idx, i]) for i in top_indices]
    top_terms_per_topic[f'主题{topic_idx+1}'] = top_terms

# 创建主题-词语关联条形图
topic_data = []
for topic, terms_list in top_terms_per_topic.items():
    for term, score in terms_list:
        topic_data.append({'主题': topic, '词语': term, '关联度': score})

topic_df = pd.DataFrame(topic_data)
sns.barplot(data=topic_df, x='关联度', y='词语', hue='主题', palette='viridis', 
            dodge=False)
plt.xlabel('词语-主题关联度', fontsize=12)
plt.ylabel('')
plt.title('主题-词语关联分析：解读潜在主题的语义含义', fontsize=14, fontweight='bold')
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

# 添加数值标注
for i, (index, row) in enumerate(topic_df.iterrows()):
    plt.text(row['关联度'] + 0.01, i, f'{row["关联度"]:.3f}', 
             va='center', fontsize=10)

plt.tight_layout()
plt.show()

# 图片5：3D文档分布图
plt.figure(figsize=(12, 8))
ax = plt.subplot(111, projection='3d')

colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown']
document_labels = [f'文档{i+1}' for i in range(len(documents))]

for i, (x, y, z) in enumerate(doc_topic_matrix[:, :3]):
    ax.scatter(x, y, z, color=colors[i], s=100, alpha=0.7, label=document_labels[i])
    ax.text(x, y, z, document_labels[i], fontsize=11)

ax.set_xlabel('主题1\n(技术主题)', fontsize=11, linespacing=1.5)
ax.set_ylabel('主题2\n(数学主题)', fontsize=11, linespacing=1.5)
ax.set_zlabel('主题3\n(算法主题)', fontsize=11, linespacing=1.5)
ax.set_title('3D语义空间：文档在多主题维度中的分布', fontsize=14, fontweight='bold')
ax.legend(bbox_to_anchor=(1.1, 1), loc='upper left')

plt.tight_layout()
plt.show()

# 图片6：主题解释方差饼图
plt.figure(figsize=(10, 8))
exploded = [0.1 if i == np.argmax(explained_variance) else 0 for i in range(len(explained_variance))]
colors_pie = ['#ff9999', '#66b3ff', '#99ff99'][:n_components]

# 创建主题标签（包含解释方差信息）
topic_labels = [f'主题{i+1}\n({explained_variance[i]:.1%})' 
               for i in range(n_components)]

wedges, texts, autotexts = plt.pie(explained_variance, labels=topic_labels,
                                  autopct='%1.1f%%', startangle=90, 
                                  colors=colors_pie, explode=exploded,
                                  textprops={'fontsize': 12})

for autotext in autotexts:
    autotext.set_color('white')
    autotext.set_fontweight('bold')
    autotext.set_fontsize(11)

plt.title('主题解释方差分布：各主题的信息贡献比例', fontsize=14, fontweight='bold')

# 添加总结文本框
summary_text = f"主题分析总结:\n• 3个主题累计解释方差: {np.sum(explained_variance):.1%}\n• 主题1贡献最大: {explained_variance[0]:.1%}\n• 信息压缩效果显著"
plt.figtext(0.5, 0.02, summary_text, ha='center', fontsize=12, 
            bbox=dict(boxstyle="round,pad=0.5", facecolor="lightgray", alpha=0.8))

plt.tight_layout()
plt.show()