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

import matplotlib.pyplot as plt

from myutils.project_utils import read_excel_to_df
from root_dir import DATA_DIR


def diverging_bar(
  df, value_col, y_ticks_col, x_label="", y_label="", title="",
  is_grid=False, is_show=False, save_name=""
):
  """
  默認(rèn)數(shù)據(jù)存儲至 df["diff"] 字段
  :param df: 數(shù)據(jù)
  :param value_col: 數(shù)值列
  :param y_ticks_col: 標(biāo)簽列
  :return:
  """
  # 準(zhǔn)備數(shù)據(jù)
  df['colors'] = ['red' if x < 0 else 'green' for x in df[value_col]]
  df.sort_values(value_col, inplace=True)
  df.reset_index(inplace=True)
  plt.figure(figsize=(12, 12), dpi=80)

  # 設(shè)置數(shù)據(jù)
  plt.hlines(y=df.index, xmin=0, xmax=df[value_col], color=df.colors, alpha=0.4, linewidth=8)
  # 設(shè)置標(biāo)簽
  # plt.yticks(df.index, df[y_ticks_col], fontsize=12)  # y軸
  plt.yticks(df.index, df[y_ticks_col], fontsize=12)  # y軸

  # 設(shè)置圖表說明
  if x_label:
    plt.gca().set(xlabel=f'${x_label}$')
  if y_label:
    plt.gca().set(ylabel=f'${y_label}$')  # $$符號表示斜體
  plt.title(title, fontdict={'size': 12})

  if is_grid:  # 顯示網(wǎng)格
    plt.grid(linestyle='--', alpha=0.5)
  if save_name:
    plt.savefig(save_name, bbox_inches='tight', format='png', transparent=True)
  if is_show:
    plt.show()

def main():
  # df = read_excel_to_df(os.path.join(DATA_DIR, "ourbest_20230605_dockq_9_final.xls"))
  df = read_excel_to_df(os.path.join(DATA_DIR, "ourbest_20230601_tmscore_56.xls"))
  x1 = df["m0-score"]
  x2 = df["m2-score"]
  # x2 = df["m1-score"]
  df["diff"] = x2 - x1
  df.info()
  diverging_bar(
    df, value_col="diff", y_ticks_col="target",
    x_label="DockQ", y_label="Target", title="",
    is_show=True, save_name="xxx.png")


if __name__ == '__main__':
  main()