import os
import datetime
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
from pyecharts import options as opts
from pyecharts.charts import Timeline, Map
from pyecharts.globals import ThemeType

plt.rcParams['font.sans-serif'] = ['simhei']  # 中文字体支持
plt.rcParams['font.serif'] = ['simhei']  # 中文字体支持
plt.rcParams['figure.figsize'] = (9, 7)  # 输出图像尺寸
sns.set_style(
    "darkgrid", {"font.sans-serif": ['simhei', 'Droid Sans Fallback']})  # 中文字体支持
pd.set_option('display.unicode.east_asian_width', True)  # 中文输出对齐
data = pd.read_csv('data/DXYArea.csv')  # 读取数据文件

# 数据概览
print('\033[32m数据概览：\033[0m')
print(data.head())
print()

print('\033[32m查看列名：\033[0m')
print(list(data.columns))
print()

# 删除冗余列
print('\033[32m删除冗余列：\033[0m')
data.drop(columns=['continentName', 'continentEnglishName', 'countryName', 'countryEnglishName', 'provinceEnglishName',
          'province_zipCode', 'cityEnglishName', 'city_zipCode', 'province_suspectedCount', 'city_suspectedCount'], inplace=True)
print(data.head())
print()

print('\033[32m查看现有列名：\033[0m')
print(list(data.columns))
print()

# 筛选出广东省的数据
print('\033[32m筛选出广东省的数据：\033[0m')
data = data[data['provinceName'] == '广东省']
print(data.head())
print(data.tail())
print()

# 筛选出2022年的数据
print('\033[32m筛选出2022年的数据：\033[0m')
data = data[data['updateTime'] >= '2022-01-01']
print(data.head())
print(data.tail())
print()

# 修改列数据类型
print('修改列数据类型：')
print('\033[31m修改前列数据类型：\033[0m')
print(data.info())
data = data.astype({'city_confirmedCount': 'int64',
                   'city_curedCount': 'int64', 'city_deadCount': 'int64'})
print('\033[32m修改后列数据类型：\033[0m')
print(data.info())
print()

# 修改列名
print('\033[32m修改列名、重建索引：\033[0m')
data.rename(columns={'provinceName': '省名', 'province_confirmedCount': '省累计确诊', 'province_curedCount': '省累计治愈', 'province_deadCount': '省累计死亡',
            'cityName': '市名', 'city_confirmedCount': '市累计确诊', 'city_curedCount': '市累计治愈', 'city_deadCount': '市累计死亡', 'updateTime': '更新时间'}, inplace=True)
# 新增现有确诊列
data['省现有确诊'] = data['省累计确诊'] - data['省累计治愈'] - data['省累计死亡']
data['市现有确诊'] = data['市累计确诊'] - data['市累计治愈'] - data['市累计死亡']
# 调整列顺序
data = data[['省名', '省现有确诊', '省累计确诊', '省累计治愈', '省累计死亡',
             '市名', '市现有确诊', '市累计确诊', '市累计治愈', '市累计死亡', '更新时间']]
print('\033[32m修改后的列名：\033[0m')
print(list(data.columns))
# 重建索引
data.reset_index(inplace=True, drop=True)
print('\033[32m修改后的数据：\033[0m')
print(data.head())
print()
print(data.tail())
print()

# 按市名分组
print('\033[32m按市名分组：\033[0m')
cites = set(data['市名'])
cites -= {'待明确地区'}
print(cites)
cites_data = dict()
for city in cites:
    temp_date = '2022-05-20'
    cites_data[city] = data[data['市名'] == city].copy()
    # print(f'查看{city}更新时间：\n',cites_data[city]['更新时间'].values)
    # 发现更新日期不连续，且有的日期重复，需要进行处理
    # 首先处理重复日期数据，保留更新时间最新的数据
    for i, row in cites_data[city].iterrows():
        if row['更新时间'][:10] == temp_date:
            cites_data[city].drop(i, inplace=True)
        else:
            temp_date = row['更新时间'][:10]
    cites_data[city].loc[:, '更新时间'] = pd.to_datetime(
        cites_data[city]['更新时间'].apply(lambda x: x[:10])).values
    # print(cites_data[city].head())
    cites_data[city].reset_index(inplace=True, drop=True)
    # 然后处理不连续日期数据，根据前后最近数据进行填充
    # 新开一个DataFrame用于存放填充的数据
    new_data = pd.DataFrame(columns=cites_data[city].columns)
    current_date = datetime.datetime(2022, 5, 20)
    current_index = 0
    for i, row in cites_data[city].iterrows():
        next_date = row['更新时间']
        days = (current_date-next_date).days
        if days > 1:
            for j in range(1, days):
                new_row = row.copy()
                new_row['省累计确诊'] = round(
                    (cites_data[city].iloc[current_index]['省累计确诊'] - row['省累计确诊']) / days * j + row['省累计确诊'])
                new_row['省累计治愈'] = round(
                    (cites_data[city].iloc[current_index]['省累计治愈'] - row['省累计治愈']) / days * j + row['省累计治愈'])
                new_row['省累计死亡'] = round(
                    (cites_data[city].iloc[current_index]['省累计死亡'] - row['省累计死亡']) / days * j + row['省累计死亡'])
                new_row['省现有确诊'] = new_row['省累计确诊'] - \
                    new_row['省累计治愈'] - new_row['省累计死亡']
                new_row['市累计确诊'] = round(
                    (cites_data[city].iloc[current_index]['市累计确诊'] - row['市累计确诊']) / days * j + row['市累计确诊'])
                new_row['市累计治愈'] = round(
                    (cites_data[city].iloc[current_index]['市累计治愈'] - row['市累计治愈']) / days * j + row['市累计治愈'])
                new_row['市累计死亡'] = round(
                    (cites_data[city].iloc[current_index]['市累计死亡'] - row['市累计死亡']) / days * j + row['市累计死亡'])
                new_row['市现有确诊'] = new_row['市累计确诊'] - \
                    new_row['市累计治愈'] - new_row['市累计死亡']
                new_row['更新时间'] = next_date + datetime.timedelta(days=j)
                new_row = new_row.to_frame().T
                new_row['更新时间'] = pd.to_datetime(
                    new_row['更新时间'].apply(lambda x: str(x)[:10])).values
                # 填充的数据放到新的DataFrame
                new_data = pd.concat([new_data, new_row], ignore_index=True)
            current_index = i
            current_date = next_date
        else:
            current_index = i
            current_date = next_date
    # 将新填充的数据添加到原数据中
    cites_data[city] = pd.concat(
        [cites_data[city], new_data], ignore_index=True)
    # 按日期顺序排序
    cites_data[city].sort_values(by='更新时间', inplace=True)
    cites_data[city].loc[:, '更新时间'] = pd.to_datetime(
        cites_data[city]['更新时间'].apply(lambda x: str(x)[:10])).values
    # 重建索引
    cites_data[city].reset_index(inplace=True, drop=True)
    # 增加省、市每日新增确诊数
    for i, row in cites_data[city].iterrows():
        if i == 0:
            cites_data[city]['省新增确诊'] = ''
            cites_data[city]['市新增确诊'] = ''
            cites_data[city].loc[0, '省新增确诊'] = 0
            cites_data[city].loc[0, '市新增确诊'] = 0
        else:
            cites_data[city].loc[i, '市新增确诊'] = row['市累计确诊']-cites_data[city].iloc[i -
                                                                                  1]['市累计确诊'] if row['市累计确诊']-cites_data[city].iloc[i-1]['市累计确诊'] > 0 else 0
            cites_data[city].loc[i, '省新增确诊'] = row['省累计确诊']-cites_data[city].iloc[i -
                                                                                  1]['省累计确诊'] if row['省累计确诊']-cites_data[city].iloc[i-1]['省累计确诊'] > 0 else 0
    # 调整列顺序
    cites_data[city] = cites_data[city][['省名', '省新增确诊', '省现有确诊', '省累计确诊', '省累计治愈', '省累计死亡',
                                         '市名', '市新增确诊', '市现有确诊', '市累计确诊', '市累计治愈', '市累计死亡', '更新时间']]
    # print(cites_data[city].head())
    # print(cites_data[city].tail())
    # assert 0

# 数据分析与可视化
if not os.path.exists('./output/'):
    os.makedirs('./output/')
    data = cites_data['广州']
    # 绘制并保存省每日新增确诊数、现有确诊数、累计确诊数折线图
    sns.lineplot(data=data, x='更新时间', y='省新增确诊')
    plt.savefig('output/省新增确诊数.jpg', dpi=1024)
    plt.clf()
    sns.lineplot(data=data, x='更新时间', y='省现有确诊')
    plt.savefig('output/省现有确诊数.jpg', dpi=1024)
    plt.clf()
    sns.lineplot(data=data, x='更新时间', y='省累计确诊')
    plt.savefig('output/省累计确诊数.jpg', dpi=1024)
    plt.clf()

    data = None
    for city in cites:
        data = pd.concat([data, cites_data[city]], ignore_index=True)
    # 绘制并保存各市每日新增确诊数、现有确诊数、累计确诊数折线图
    sns.lineplot(data=data, x='更新时间', y='市新增确诊', hue='市名', style='市名')
    plt.savefig('output/各市新增确诊数.jpg', dpi=1024)
    plt.clf()
    sns.lineplot(data=data, x='更新时间', y='市现有确诊', hue='市名', style='市名')
    plt.savefig('output/各市现有确诊数.jpg', dpi=1024)
    plt.clf()
    sns.lineplot(data=data, x='更新时间', y='市累计确诊', hue='市名', style='市名')
    plt.savefig('output/各市累计确诊数.jpg', dpi=1024)
    plt.clf()

# 现有确诊动态地图


def current_confirmed_timeline_map() -> Timeline:
    global cites
    tl = Timeline(init_opts=opts.InitOpts(page_title="2022年广东省每日现有确诊动态地图",
                                          theme=ThemeType.DARK,
                                          width="1300px", height="666px"),
                  )
    index = 0
    cites = list(cites)
    for date in pd.date_range(start='2022-01-01', end='2022-05-19'):
        date = str(date)[:10]
        confirmedCount = []
        # assert 0
        for city in cites:
            confirmedCount.append(cites_data[city].iloc[index]['市现有确诊'])
        zipped = zip([x + '市' for x in cites], confirmedCount)
        f_map = (
            Map(init_opts=opts.InitOpts(width="1000px",
                                        height="600px",
                                        page_title="2022年广东省每日现有确诊动态地图",
                                        bg_color=None))
            .add(series_name="现有确诊数量",
                             data_pair=[list(z) for z in zipped],
                             maptype="广东",
                             is_map_symbol_show=False)
            .set_global_opts(
                title_opts=opts.TitleOpts(title="2022年广东省每日现有确诊动态地图",
                                          subtitle=f"{date[5:7]}月{date[8:10]}日数据\n",
                                          pos_left="center", ),
                legend_opts=opts.LegendOpts(
                    is_show=True, pos_top="40px", pos_right="30px"),
                visualmap_opts=opts.VisualMapOpts(
                    is_piecewise=True, range_text=['', ''], pieces=[
                        {"min": 200, "color": "#751d0d"},
                        {"min": 100, "max": 199, "color": "#ae2a23"},
                        {"min": 50, "max": 99, "color": "#d6564c"},
                        {"min": 20, "max": 49, "color": "#f19178"},
                        {"min": 10, "max": 19, "color": "#f7d3a6"},
                        {"min": 1, "max": 9, "color": "#fdf2d3"},
                        {"min": 0, "max": 0, "color": "#AAFFAA"}
                    ]),
            )
            .set_series_opts(label_opts=opts.LabelOpts(is_show=True),
                             markpoint_opts=opts.MarkPointOpts(
                                 symbol_size=[90, 90], symbol='circle'),
                             effect_opts=opts.EffectOpts(is_show='True', )
                             )
        )
        tl.add(f_map, f"{date[5:7]}月{date[8:10]}日")
        tl.add_schema(is_timeline_show=True,  # 是否显示
                      play_interval=500,  # 播放间隔
                      symbol=None,  # 图标
                      is_loop_play=True  # 循环播放
                      )
        index += 1
    return tl

# 新增确诊动态地图


def new_confirmed_timeline_map() -> Timeline:
    global cites
    tl = Timeline(init_opts=opts.InitOpts(page_title="2022年广东省每日新增确诊动态地图",
                                          theme=ThemeType.DARK,
                                          width="1300px", height="666px"),
                  )
    index = 0
    cites = list(cites)
    for date in pd.date_range(start='2022-01-01', end='2022-05-19'):
        date = str(date)[:10]
        confirmedCount = []
        # assert 0
        for city in cites:
            confirmedCount.append(cites_data[city].iloc[index]['市新增确诊'])
        zipped = zip([x + '市' for x in cites], confirmedCount)
        f_map = (
            Map(init_opts=opts.InitOpts(width="1000px",
                                        height="600px",
                                        page_title="2022年广东省每日新增确诊动态地图",
                                        bg_color=None))
            .add(series_name="新增确诊数量",
                             data_pair=[list(z) for z in zipped],
                             maptype="广东",
                             is_map_symbol_show=False)
            .set_global_opts(
                title_opts=opts.TitleOpts(title="2022年广东省每日新增确诊动态地图",
                                          subtitle=f"{date[5:7]}月{date[8:10]}日数据\n",
                                          pos_left="center", ),
                legend_opts=opts.LegendOpts(
                    is_show=True, pos_top="40px", pos_right="30px"),
                visualmap_opts=opts.VisualMapOpts(
                    is_piecewise=True, range_text=['', ''], pieces=[
                        {"min": 70, "color": "#751d0d"},
                        {"min": 40, "max": 69, "color": "#ae2a23"},
                        {"min": 20, "max": 39, "color": "#d6564c"},
                        {"min": 10, "max": 19, "color": "#f19178"},
                        {"min": 5, "max": 9, "color": "#f7d3a6"},
                        {"min": 1, "max": 4, "color": "#fdf2d3"},
                        {"min": 0, "max": 0, "color": "#AAFFAA"}
                    ]),
            )
            .set_series_opts(label_opts=opts.LabelOpts(is_show=True),
                             markpoint_opts=opts.MarkPointOpts(
                                 symbol_size=[90, 90], symbol='circle'),
                             effect_opts=opts.EffectOpts(is_show='True', )
                             )
        )
        tl.add(f_map, f"{date[5:7]}月{date[8:10]}日")
        tl.add_schema(is_timeline_show=True,  # 是否显示
                      play_interval=500,  # 播放间隔
                      symbol=None,  # 图标
                      is_loop_play=True  # 循环播放
                      )
        index += 1
    return tl


# 生成动态地图的HTML文件
current_confirmed_timeline_map().render('./output/2022年广东省每日现有确诊动态地图.html')
new_confirmed_timeline_map().render('./output/2022年广东省每日新增确诊动态地图.html')