from io import BytesIO

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd


class Radar(object):

    def __init__(self, figure, title, labels, epoch, rect=None):
        if rect is None:
            rect = [0.1, 0.05, 0.75, 0.9]

        self.n = len(title)
        self.angles = np.arange(0, 360, 360.0 / self.n)
        self.labels = labels

        self.axes = [figure.add_axes(rect, projection='polar', label='axes%d' % i) for i in range(self.n)]

        self.ax = self.axes[0]
        self.ax.set_thetagrids(self.angles, labels=title, fontsize=12)
        self.ax.xaxis.set_tick_params(pad=22)  # Increase padding to move labels outward

        for ax in self.axes[1:]:
            ax.patch.set_visible(False)
            ax.grid(False)
            ax.xaxis.set_visible(False)

        for ax, angle, label, i in zip(self.axes, self.angles, labels, epoch):
            ax.set_rgrids(i[1:], angle=angle, labels=label)
            ax.spines['polar'].set_visible(False)
            ax.set_ylim(i[0], i[-1])

    def plot(self, values, *args, **kw):
        angle = np.deg2rad(np.r_[self.angles, self.angles[0]])

        # 第一个值保持不变
        adjusted_values = [values[0]]

        # 获取 labels 对应的最后一位的差值
        base_last_value = float(self.labels[0][-1])  # 第一个 labels 列表的最后一位

        # 计算新的 values
        for i in range(1, len(values)):
            current_label_last_value = float(self.labels[i][-1])  # 当前 labels 列表的最后一位
            difference = current_label_last_value - base_last_value
            adjusted_value = values[i] - difference
            adjusted_values.append(adjusted_value)

        # 将第一个值追加到最后，闭合图形
        adjusted_values.append(adjusted_values[0])

        # 绘制雷达图
        self.ax.plot(angle, adjusted_values, *args, **kw)


def generate_lab(max_values, length=5, step=1):
    """
    根据最大值向量生成lab（标签）列表。
    """
    lab = []
    for max_val in max_values:
        labels = [str(round(val, 2)) for val in np.arange(max_val - (length - 1) * step, max_val + step, step)]
        lab.append(labels)
    return lab


def generate_epo(max_values, length=6, step=1):
    """
    根据最大值向量生成epo（刻度范围）列表。
    """
    epo = []
    for max_val in max_values:
        range_values = [round(val, 2) for val in np.arange(max_val - (length - 1) * step, max_val + step, step)]
        epo.append(range_values)
    return epo


def plot_radar(data, max_values):
    """
    绘制雷达图，根据数据和最大值生成图表。
    """
    fig = plt.figure(figsize=(10, 8))
    tit = data.columns.tolist()
    lab = generate_lab(max_values, length=5, step=1)
    epo = generate_epo(max_values, length=6, step=1)

    plt.rcParams['font.sans-serif'] = ['SimHei']
    radar = Radar(fig, tit, lab, epo)

    # 遍历 DataFrame 的每一行
    colors = plt.colormaps.get_cmap('tab20')
    for i, (idx, row) in enumerate(data.iterrows()):
        radar.plot(row.values.tolist(), '-', lw=2, color=colors(i), alpha=0.4, label=idx)

    radar.ax.legend(bbox_to_anchor=(1.2, 1.05))

    # plt.savefig('fig.png', dpi=300)

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

    # plt.show()
    plt.close()
    return img_bytes


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

    # 任务设置的每个指标的最大值
    max_values = [10, 15, 15, 10, 15, 5, 10, 10, 5, 5]  # 示例最大值向量
    # 使用这个函数
    plot_radar(data, max_values)