from io import BytesIO

import pandas as pd
from matplotlib import rcParams
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt

# 设置中文字体
rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为宋体
rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


# PCA函数并绘制前两个主成分的散点图，带标签
def plot_pca_scatter(df, labels=None):
    # 如果 labels 为 None，则默认使用 DataFrame 的行名
    if labels is None:
        labels = df.index

    # 判断输入的数据行数是否只有一行
    if df.shape[0] == 1:
        plt.figure(figsize=(8, 6))
        plt.text(0.5, 0.5, '无法对单个样本进行主成分分析。', fontsize=16, ha='center', va='center')
        plt.title("主成分分析 结果", fontsize=16)
        plt.xlabel('')
        plt.ylabel('')
        # 将图像保存到字节流中
        img_bytes = BytesIO()
        plt.savefig(img_bytes, format='png')
        # 设置字节流的位置到开始
        img_bytes.seek(0)
        # 输出图像
        plt.close()
        return img_bytes

    # 进行PCA
    pca = PCA(n_components=2)
    pca_result = pca.fit_transform(df)

    # 将PCA结果转换为DataFrame
    pca_df = pd.DataFrame(pca_result, columns=['主成分 1', '主成分 2'])

    # 绘制前两个主成分的散点图
    plt.figure(figsize=(8, 6))
    plt.scatter(pca_df['主成分 1'], pca_df['主成分 2'], color='b', label='样本点')
    plt.title("主成分分析 前两主成分散点图", fontsize=16)
    plt.xlabel('主成分 1', fontsize=12)
    plt.ylabel('主成分 2', fontsize=12)
    plt.grid(True)

    # 在点旁边添加标签
    for i, label in enumerate(labels):
        plt.text(pca_df['主成分 1'][i] + 0.02, pca_df['主成分 2'][i] + 0.02, label, fontsize=9)

    # 将图像保存到字节流中
    img_bytes = BytesIO()
    plt.savefig(img_bytes, format='png')
    # 设置字节流的位置到开始
    img_bytes.seek(0)
    # 输出图像
    plt.close()
    return img_bytes


# 示例用法
if __name__ == '__main__':
    # 一个行名为样本名，列名为指标名的pd.dataframe
    data = pd.read_csv('./radartest.csv', index_col=0)

    # 生成散点图，带上标签
    plot_pca_scatter(data, labels=data.index)