from sklearn.decomposition import PCA  # 导入PCA模型
import matplotlib.pyplot as plt  # 导入绑图库

plt.rcParams['font.sans-serif'] = ['Source Han Serif SC']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

pca_model = PCA()  # 初始化PCA模型（保留所有主成分）
pca_scores = pca_model.fit_transform(standardized_features)  # 拟合PCA并获取主成分得分
explained_variance_ratio = pca_model.explained_variance_ratio_  # 提取各主成分的方差解释比
cumulative_variance_ratio = np.cumsum(explained_variance_ratio)  # 计算累积方差解释比

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))  # 创建1行2列子图

# 左图：碎石图
axes[0].bar(range(1, len(explained_variance_ratio) + 1), explained_variance_ratio, color='#008080', alpha=0.8)  # 柱状图
axes[0].set_xlabel('主成分编号')  # 设置x轴标签
axes[0].set_ylabel('方差解释比')  # 设置y轴标签
axes[0].set_title('碎石图 (Scree Plot)')  # 设置标题

# 右图：累积方差解释比
axes[1].plot(range(1, len(cumulative_variance_ratio) + 1), cumulative_variance_ratio, 'o-', color='#E3120B')  # 折线图
axes[1].axhline(y=0.80, color='gray', linestyle='--', alpha=0.5, label='80% 阈值')  # 80%阈值线
axes[1].set_xlabel('主成分个数')  # 设置x轴标签
axes[1].set_ylabel('累积方差解释比')  # 设置y轴标签
axes[1].set_title('累积方差解释比')  # 设置标题
axes[1].legend()  # 显示图例
plt.tight_layout()  # 自动调整间距
plt.show()  # 显示图表

print(f'PC1 方差解释比: {explained_variance_ratio[0]:.2%}')  # 输出PC1方差解释比
print(f'PC1+PC2 累积: {cumulative_variance_ratio[1]:.2%}')  # 输出前两个PC累积方差
print(f'PC1+PC2+PC3 累积: {cumulative_variance_ratio[2]:.2%}')  # 输出前三个PC累积方差

fig, ax = plt.subplots(figsize=(8, 6))  # 创建单张图表

# 提取前两个主成分的载荷
loadings_pc1 = pca_model.components_[0]  # PC1的载荷向量
loadings_pc2 = pca_model.components_[1]  # PC2的载荷向量

# 绘制载荷箭头
for i, feature_name in enumerate(pca_feature_columns):  # 遍历每个特征
    ax.annotate('', xy=(loadings_pc1[i], loadings_pc2[i]), xytext=(0, 0),  # 从原点到载荷点画箭头
                arrowprops=dict(arrowstyle='->', color='#E3120B', lw=2))  # 红色箭头
    ax.text(loadings_pc1[i] * 1.15, loadings_pc2[i] * 1.15, feature_name,  # 在箭头端点标注特征名
            fontsize=11, ha='center', color='#2C3E50', fontweight='bold')  # 加粗标注

ax.axhline(y=0, color='gray', linestyle='--', alpha=0.3)  # 水平参考线
ax.axvline(x=0, color='gray', linestyle='--', alpha=0.3)  # 垂直参考线
ax.set_xlabel(f'PC1 ({explained_variance_ratio[0]:.1%})')  # x轴标签（含方差解释比）
ax.set_ylabel(f'PC2 ({explained_variance_ratio[1]:.1%})')  # y轴标签（含方差解释比）
ax.set_title('PCA 载荷图')  # 设置标题
plt.tight_layout()  # 自动调整间距
plt.show()  # 显示图表

# 合并行业信息
yrd_pca_data_with_industry = yrd_pca_data.merge(  # 合并行业分类
    stock_basic_data[['order_book_id', 'industry_name']], on='order_book_id', how='left'  # 左连接
)  # 获取每家公司的行业归属

# 统计各行业频次，选择前6大行业
top_six_industries = yrd_pca_data_with_industry['industry_name'].value_counts().head(6).index.tolist()  # 前6大行业名称
plot_mask = yrd_pca_data_with_industry['industry_name'].isin(top_six_industries)  # 筛选属于前6行业的公司

fig, ax = plt.subplots(figsize=(10, 6))  # 创建图表
industry_color_palette = ['#E3120B', '#008080', '#F0A700', '#2C3E50', '#8E9EAA', '#6A5ACD']  # 6种行业颜色

for i, industry_name in enumerate(top_six_industries):  # 遍历前6大行业
    industry_mask = yrd_pca_data_with_industry['industry_name'] == industry_name  # 筛选当前行业
    valid_indices = yrd_pca_data_with_industry[industry_mask].index  # 获取行业内公司的索引
    # 获取这些公司在PCA得分矩阵中的位置
    pca_row_indices = [yrd_pca_data.index.get_loc(idx) for idx in valid_indices if idx in yrd_pca_data.index]  # 映射索引
    ax.scatter(pca_scores[pca_row_indices, 0], pca_scores[pca_row_indices, 1],  # 绘制散点
               c=industry_color_palette[i], label=industry_name, alpha=0.6, s=30)  # 按行业着色

ax.set_xlabel(f'PC1 ({explained_variance_ratio[0]:.1%})')  # x轴标签
ax.set_ylabel(f'PC2 ({explained_variance_ratio[1]:.1%})')  # y轴标签
ax.set_title('主成分得分的行业分布')  # 设置标题
ax.legend(fontsize=9, ncol=2)  # 显示两列图例
plt.tight_layout()  # 自动调整间距
plt.show()  # 显示图表

from sklearn.cluster import KMeans  # 导入K-Means聚类模型
from sklearn.metrics import silhouette_score  # 导入轮廓系数

# 使用前3个主成分作为聚类特征
pca_features_for_clustering = pca_scores[:, :3]  # 提取前3个PC得分

k_range = range(2, 10)  # 候选K值范围：2到9
inertia_values = []  # 存储各K值的组内平方和 (SSE/Inertia)
silhouette_values = []  # 存储各K值的轮廓系数

for k in k_range:  # 遍历每个候选K值
    kmeans_model = KMeans(n_clusters=k, random_state=42, n_init=10)  # 初始化K-Means模型
    cluster_labels = kmeans_model.fit_predict(pca_features_for_clustering)  # 拟合并预测聚类标签
    inertia_values.append(kmeans_model.inertia_)  # 记录组内平方和
    silhouette_values.append(silhouette_score(pca_features_for_clustering, cluster_labels))  # 记录轮廓系数

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))  # 创建1行2列子图

axes[0].plot(list(k_range), inertia_values, 'o-', color='#E3120B', linewidth=2)  # 肘部法折线图
axes[0].set_xlabel('聚类数 K')  # x轴标签
axes[0].set_ylabel('组内平方和 (SSE)')  # y轴标签
axes[0].set_title('肘部法 (Elbow Method)')  # 标题

axes[1].plot(list(k_range), silhouette_values, 's-', color='#008080', linewidth=2)  # 轮廓系数折线图
axes[1].set_xlabel('聚类数 K')  # x轴标签
axes[1].set_ylabel('轮廓系数')  # y轴标签
axes[1].set_title('轮廓系数法 (Silhouette Method)')  # 标题

plt.tight_layout()  # 自动调整间距
plt.show()  # 显示图表

fig, ax = plt.subplots(figsize=(9, 6))  # 创建图表
cluster_colors = ['#E3120B', '#008080', '#F0A700', '#2C3E50']  # 4个聚类的颜色

for cluster_id in range(4):  # 遍历每个聚类
    cluster_mask = final_cluster_labels == cluster_id  # 筛选当前聚类的样本
    ax.scatter(pca_scores[cluster_mask, 0], pca_scores[cluster_mask, 1],  # 在PC1-PC2空间绑制散点
               c=cluster_colors[cluster_id], label=f'聚类 {cluster_id}', alpha=0.6, s=30)  # 按聚类着色

ax.set_xlabel(f'PC1 ({explained_variance_ratio[0]:.1%})')  # x轴标签
ax.set_ylabel(f'PC2 ({explained_variance_ratio[1]:.1%})')  # y轴标签
ax.set_title('K-Means 聚类结果（PCA空间投影）')  # 标题
ax.legend()  # 显示图例
plt.tight_layout()  # 自动调整间距
plt.show()  # 显示图表

from scipy.cluster.hierarchy import dendrogram, linkage  # 导入层次聚类工具

np.random.seed(42)  # 设置随机种子
sample_size = min(50, len(standardized_features))  # 取样本量（最多50家）
sample_indices = np.random.choice(len(standardized_features), sample_size, replace=False)  # 随机抽样索引
sample_features = standardized_features[sample_indices]  # 提取样本特征

# 合并行业名称用于标签
sample_stock_codes = yrd_pca_data.iloc[sample_indices]['order_book_id'].values  # 获取样本公司代码
sample_industry_names = stock_basic_data.set_index('order_book_id').loc[sample_stock_codes, 'industry_name'].values  # 获取行业名

# Ward 法层次聚类
linkage_matrix = linkage(sample_features, method='ward')  # 使用Ward方法计算链接矩阵

fig, ax = plt.subplots(figsize=(14, 6))  # 创建图表
dendrogram(linkage_matrix, labels=sample_industry_names, leaf_rotation=90, leaf_font_size=8, ax=ax)  # 绘制树状图
ax.set_title('层次聚类树状图（Ward法）')  # 标题
ax.set_ylabel('距离')  # y轴标签
plt.tight_layout()  # 自动调整间距
plt.show()  # 显示图表

from scipy.spatial.distance import squareform  # 导入距离矩阵压缩工具

# 将方阵转换为压缩距离向量
condensed_distance = squareform(distance_matrix.values, checks=False)  # 压缩距离矩阵
stock_linkage = linkage(condensed_distance, method='ward')  # Ward法层次聚类

# 构建行业标签
industry_labels = []  # 存放行业简称标签
for stock_code in distance_matrix.columns:  # 遍历每只股票
    label_found = False  # 标记是否找到行业
    for ind_name, ind_stocks in industry_stock_mapping.items():  # 遍历行业映射
        if stock_code in ind_stocks:  # 匹配股票代码
            industry_labels.append(f'{ind_name[:2]}')  # 取行业名前两个字作为简称
            label_found = True  # 标记已找到
            break  # 跳出内层循环
    if not label_found:  # 未匹配时
        industry_labels.append(stock_code[:6])  # 使用股票代码前6位作为标签

fig, ax = plt.subplots(figsize=(14, 6))  # 创建图表
dendrogram(stock_linkage, labels=industry_labels, leaf_rotation=0, leaf_font_size=10, ax=ax)  # 绘制树状图
ax.set_title('基于收益率相关性的股票聚类树状图')  # 标题
ax.set_ylabel('距离（1 - 相关系数）')  # y轴标签
plt.tight_layout()  # 自动调整间距
plt.show()  # 显示图表

import seaborn as sns  # 导入seaborn高级绑图库

# 使用seaborn的clustermap自动排列行列
cluster_heatmap = sns.clustermap(  # 创建聚类热力图
    correlation_matrix,  # 输入相关系数矩阵
    method='ward',  # 使用Ward法聚类
    cmap='coolwarm',  # 冷暖色调（蓝=低相关，红=高相关）
    vmin=-0.5, vmax=1,  # 颜色范围
    figsize=(10, 8),  # 图表大小
    linewidths=0.5  # 网格线宽度
)  # 生成聚类热力图
cluster_heatmap.ax_heatmap.set_title('股票收益率相关性聚类热力图', fontsize=14, pad=60)  # 设置标题
plt.show()  # 显示图表