import sys
import pandas as pd
import matplotlib.pyplot as plt
from config import FILE_NAME, PAINT_RANK_SIZE

def load_data():
    try:
        dataframe = pd.read_csv(FILE_NAME, names=["country", "person", "death"], encoding="utf-8")
    except BaseException:
        print("错误编码格式")
        sys.exit(1)
    
    #删除第一行,得到前PAINT_RANK_SIZE名的国家
    return dataframe.drop(0).head(PAINT_RANK_SIZE)

def paint_bar():
    dataframe = load_data()
    
    # 设置中文显示
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['figure.figsize'] = (PAINT_RANK_SIZE * 0.8, 5)  # 设置figure_size尺寸
    
    # x轴坐标
    x = dataframe["country"].values
    # y轴坐标
    y = dataframe["death"].values
    # 绘制柱状图
    plt.bar(x, y)
    # 设置x轴名称
    plt.xlabel("国家",fontsize=14)
    # 设置y轴名称
    plt.ylabel("死亡人数",fontsize=14)
    plt.savefig('image_bar.jpg', dpi=800)
    plt.show()

def paint_pie():
    dataframe = load_data()
    
    labels = dataframe["country"].values
    counts = dataframe["death"].values
    dist = [x / 5e6 for x in counts]#分离距离与counts相关
    
    plt.pie(counts, dist, labels, autopct='%.1f%%')
    plt.savefig('image_pie.jpg', dpi=800)
    plt.show()