"""
指定两组数据标签，绘制相关系数热力图，用于定向分析
"""
import sys
sys.path.append("..")
from temp import config_analysis
from temp.utils_analysis import create_custom_heatmap, set_chinese_font, cross_corr
from Database.database_ import Database, DatabaseParam
import json
import matplotlib.pyplot as plt
from datetime import datetime
time_now = datetime.now().strftime('%H:%M:%S')
#series_a_name = ['管廊间温度1','加药间温度1','膜车间温度1','泵房温度1']   # 需要修改, 列标签
#series_b_name = ['RO1 1段酸洗记次数','RO1 1段J洗记次数', 'RO1 1段FY洗记次数', 'RO1 2段FY洗记次数', '1#超滤反洗水泵运行时间', '段间泵A温度', '反渗透高压泵A温度', '反渗透冲洗水泵A温度','反渗透外供水泵A温度','超滤反洗水泵A温度','超滤供水泵A温度','超滤总进水浊度','超滤反洗泵A 累计电量'] # 需要修改， 行标签
#series_a_name = ['UF1跨膜压差','UF2跨膜压差', 'UF3跨膜压差', 'UF4跨膜压差','超滤总产水PH', 'RO1产水电导','RO2产水电导','RO3产水电导','RO4产水电导','反渗透总产水电导','RO1一段浓水压力','RO2一段浓水压力','RO3一段浓水压力','RO4一段浓水压力','RO1一段进水压力','RO2一段进水压力','RO3一段进水压力','RO4一段进水压力']
#series_b_name = ['超滤总进水压力','超滤总产水压力','超滤总进水浊度', '反渗透总进水温度','反渗透总进水电导','超滤总产水PH','反渗透总进水PH','外供水-PH','1#阻垢剂加药泵(RO)运行频率','2#阻垢剂加药泵(RO)运行频率','3#阻垢剂加药泵(RO)运行频率', '4#阻垢剂加药泵(RO)运行频率','1#反渗透阻垢剂流量','2#反渗透阻垢剂流量','3#反渗透阻垢剂流量','4#反渗透阻垢剂流量','还原剂流量','还原剂药箱1液位','还原剂药箱2液位','盐酸药箱液位','RO1一段浓水压力','RO2一段浓水压力','RO3一段浓水压力','RO4一段浓水压力','RO1一段进水压力','RO2一段进水压力','RO3一段进水压力','RO4一段进水压力']
series_a_name = ['超滤总进水压力','RO1产水电导','RO2产水电导','RO3产水电导','RO4产水电导','超滤总产水PH','反渗透总产水电导']
series_b_name = ['RO1一段浓水压力','RO2一段浓水压力','RO3一段浓水压力','RO4一段浓水压力','RO1一段进水压力','RO2一段进水压力','RO3一段进水压力','RO4一段进水压力','超滤总进水压力','反渗透总进水温度']
#series_a_name = ['RO1产水电导','RO2产水电导','RO3产水电导','RO4产水电导']
#series_b_name = ['1#阻垢剂加药泵(RO)运行频率','2#阻垢剂加药泵(RO)运行频率','3#阻垢剂加药泵(RO)运行频率', '4#阻垢剂加药泵(RO)运行频率','1#反渗透阻垢剂流量','2#反渗透阻垢剂流量','3#反渗透阻垢剂流量','4#反渗透阻垢剂流量']
db_param = DatabaseParam(
        db_host='192.168.50.4',
        db_user='root',
        db_password='*B-@p2b+97D5xAF1e6',
        db_name='ws_data',
        db_port=4000)
# 检查字段
with open('../GetItem/92_all_items_name_code_transfer.json', 'r', encoding='utf-8') as f:  # 总字段加载文件
    name_code_transfer = json.load(f)
    print(f'加载name与code映射字典，共{name_code_transfer.get('len')}条')
name_2_code = name_code_transfer.get('name_2_code')
code_2_name = name_code_transfer.get('code_2_name')
del name_code_transfer

for name in series_a_name:
    if name not in name_2_code.keys():
        raise IOError(f'指定字段{name}不存在', '输入A序列：',series_a_name)
for name in series_b_name:
    if name not in name_2_code.keys():
        raise IOError(f'指定字段{name}不存在', '输入B序列：',series_b_name)
# 设置中文字体
set_chinese_font()
# 按组计算
with Database(db_param) as db:
    series_a_code = [name_2_code.get(i) for i in series_a_name]
    series_b_code = [name_2_code.get(i) for i in series_b_name]
    group_df = db.query_sql_time_series_group2data_frame(
        code_name_dict=code_2_name,
        project_id=config_analysis.PROJECT_ID,
        sheet_name=config_analysis.DB_SHEET_NAME,
        data_codes=series_a_code + series_b_code,
        start_year=config_analysis.CHECK_YEAR_START,
        end_year=config_analysis.CHECK_YEAR_END,
        start_month=config_analysis.CHECK_MONTH_START,
        end_month=config_analysis.CHECK_MONTH_END,
        start_day=config_analysis.CHECK_DAY_START,
        end_day=config_analysis.CHECK_DAY_END)
    # 剔除不存在的字段
    series_a_code = [i for i in series_a_code if i in group_df.columns.tolist()]
    series_b_code = [i for i in series_b_code if i in group_df.columns.tolist()]

    corr = cross_corr(series_a_code, series_b_code, group_df, code_2_name)
    create_custom_heatmap(corr_matrix=corr, title=f'中荷水厂数据相关系数热力图 ' + time_now)

    # 绘制曲线
    for code in series_a_code+series_b_code:

        plt.figure(figsize=(12, 6))
        # 原始数据
        plt.plot(group_df['time'], group_df[code], 'b-', alpha=0.3, label='原始数据')
        plt.title(f'{code_2_name.get(code)} 时间序列')
        plt.xlabel('时间')
        plt.ylabel(code_2_name.get(code))
        plt.legend()
        plt.grid(True)
        plt.xticks(rotation=45)  # 旋转x轴标签以便阅读
        plt.savefig(f'{code_2_name.get(code)}.png', bbox_inches='tight', dpi=300)
        plt.close()
        print(f"已保存 {code_2_name.get(code)} 的曲线图到 {code_2_name.get(code)}.png")