import numpy as np
import openpyxl
import os.path
import matplotlib.pyplot as plt


def draw_cbf_graph(cbf_mean_dict, data_type, save_path):
    plt.rcParams["figure.figsize"] = [32, 18]
    plt.rcParams["figure.autolayout"] = True
    plt.rcParams["font.sans-serif"] = ["SimHei"]  # 设置字体
    subplot = plt.subplot(1, 1, 1, projection='polar')

    count = len(cbf_mean_dict)
    theta = np.linspace(0.0, 2 * np.pi, count, endpoint=False)
    data = np.array(list(cbf_mean_dict.values()))
    width = 1.8 * np.pi / count
    subplot.set_theta_zero_location('N')
    subplot.set_theta_direction(-1)
    subplot.xaxis.set_ticks(theta)
    subplot.yaxis.set_ticks([])

    bars = subplot.bar(x=theta, height=data, width=width)
    for data_value, bar in zip(data, bars):
        bar.set_facecolor(plt.cm.turbo(data_value / 60))
        bar.set_alpha(0.8)
    for theta_value, data_value in zip(theta, data):
        margin = data.max() * 0.08
        subplot.annotate(str(round(data_value, 2)), xy=(theta_value, data_value + margin))

    subplot.set_xticklabels(list(cbf_mean_dict.keys()))
    for label, angle in zip(subplot.get_xticklabels(), theta):
        x, y = label.get_position()
        lab = subplot.text(x, y - 0.1, label.get_text(), transform=label.get_transform(), ha=label.get_ha(),
                           va=label.get_va())
        rotation = - np.rad2deg(angle) + 90
        if rotation < -90:
            rotation += 180
        lab.set_rotation(rotation)
    subplot.set_xticklabels([])

    plt.title(f"221例健康人西门子5延迟数据各脑区{data_type}")
    plt.savefig(save_path, dpi=200)
    plt.cla()
    plt.close('all')


def main():
    # 请尝试添加att参数，并调整色阶范围使得CBF=40时为红色，60时为蓝色
    data_types = ['corr-CBF']
    atlas_list = ['AnImage_AAL3']
    output_dir = r'..\Data\polar-graph'
    os.makedirs(output_dir, exist_ok=True)

    input_file = r'..\Data\csv-simple-stat.xlsx'
    workbook = openpyxl.load_workbook(input_file, data_only=True)

    for data_type in data_types:
        for atlas in atlas_list:
            cbf_mean_dict = {}
            worksheet = workbook[f'{atlas}_{data_type}']
            for column_num in range(2, worksheet.max_column + 1):
                name = worksheet.cell(row=1, column=column_num).value
                mean = worksheet.cell(row=worksheet.max_row - 8, column=column_num).value
                print(atlas, data_type, name, mean)
                if type(mean) is float:
                    cbf_mean_dict[name] = mean

            draw_cbf_graph(cbf_mean_dict, data_type, os.path.join(output_dir, f'{atlas}_{data_type}.png'))


main()