#!/usr/bin/env python
# -- coding: utf-8 --
"""
Copyright (c) 2022. All rights reserved.
Created by C. L. Wang on 2023/6/25
"""
import os

import seaborn as sns

sns.set_theme(style="white")

import matplotlib.pyplot as plt
import numpy as np

from myutils.project_utils import read_excel_to_df
from root_dir import DATA_DIR


def draw_diagonal_scatter_plots(
    data_better, data_ref, label_list,
    min_scale=0.0, max_scale=1.05,
    x_label="", y_label="",
    save_name=""
):
    """
    繪制對角線散點圖

    :param data_better: 優(yōu)質(zhì)數(shù)據(jù)，數(shù)據(jù)位于右下方
    :param data_ref:    對比數(shù)據(jù)
    :param label_list:  標簽
    :param min_scale:   最小范圍
    :param max_scale:   最大范圍
    :param x_label:     x軸描述
    :param y_label:     y軸描述
    :param save_name:   文件存儲
    :return: 圖
    """
    assert len(data_ref) == len(data_better) == len(label_list)

    fig, ax = plt.subplots()
    fig.set_size_inches(10, 10)

    ax.grid(True)
    ax.plot([min_scale, max_scale], [min_scale, max_scale], ls="--", c=".3")
    ax.scatter(data_ref, data_better, s=100, edgecolors="black")

    plt.xlim(min_scale, max_scale)
    plt.ylim(min_scale, max_scale)
    plt.xticks(fontsize=15)
    plt.yticks(fontsize=15)
    plt.xlabel(x_label, fontsize=20)
    plt.ylabel(y_label, fontsize=20)

    for i, txt in enumerate(label_list):
        if data_ref[i] - data_better[i] > 0.025:
            ax.annotate(txt, (data_ref[i] * 1.03, data_better[i]), fontsize=10, fontweight='bold')

    if save_name:
        # transparent=True
        assert save_name.endswith("png") or save_name.endswith("jpg")
        plt.savefig(save_name, bbox_inches='tight', format='png')

    plt.show()


def main():
    df = read_excel_to_df(os.path.join(DATA_DIR, "Strategy-v1v2v3v4-TMScore.xls"))
    data1 = df["m0-score"]
    data2 = df["max-score"]
    label_list = df["target"]
    print(f"data1 : {round(float(np.mean(data1)) * 100, 4)}±{round(float(np.std(data1)), 4)}")
    print(f"data2 : {round(float(np.mean(data2)) * 100, 4)}±{round(float(np.std(data2)), 4)}")
    draw_diagonal_scatter_plots(data1, data2, label_list, min_scale=0.35, max_scale=1.05,
                                x_label="Our P.S.P. (TMScore)", y_label="CASP15 SOTA (TMScore)",
                                save_name="img.png")


if __name__ == '__main__':
    main()