from yangke.pytorch.mytorch import DataFitterNet, DataSetFitting, re_train
import pandas as pd
import torch
from yangke.common.config import logger, add_file_logger

add_file_logger('log-神经网络训练.txt')
# 定义输入和输出列
input_columns = ["汽轮发电机功率", "凝汽器热负荷", "环境温度", "大气湿度", "冷却塔进水蝶阀开启数量", "风机运行台数",
                 "循泵运行方式数字编码"]
output_column = ["背压"]

# 定义模型配置
settings = {
    "dataset": {"data_file": r"D:\国华燃机\data.csv"},

    "input": input_columns,
    "output": output_column,
    "networks": {
        "train": {
            "epochs": 100,
            "batch_size": 64,
            "learning_rate": 0.001,
        },
        "loss_fn": {
            "type": "PriorLoss",
            "punish": 20000,
            "bak_info": [
                # {'isChecked': True, 'op': '任何条件', 'expression': '[背压]随[汽轮发电机功率]单调递增'},
                {'isChecked': True, 'op': '任何条件', 'expression': '[背压]随[凝汽器热负荷]单调递增'},
                {'isChecked': True, 'op': '任何条件', 'expression': '[背压]随[环境温度]单调递增'},
                {'isChecked': True, 'op': '任何条件', 'expression': '[背压]随[风机运行台数]单调递减'},
                {'isChecked': True, 'op': '任何条件', 'expression': '[背压]随[循泵运行方式数字编码]单调递减'},
            ]
        },
        "optimizer": {
            "type": "adam",
        },
        "layers": [
            {"type": "input", "cell_num": "auto"},
            {"type": "linear", "cell_num": 128, "bias": True},
            {"type": "relu"},
            {"type": "dropout", "rate": 0.2},
            {"type": "linear", "cell_num": 64, "bias": True},
            {"type": "relu"},
            {"type": "dropout", "rate": 0.2},
            {"type": "linear", "cell_num": 32, "bias": True},
            {"type": "relu"},
            {"type": "output", "cell_num": "auto"}
        ],
    },
    "save_path": "背压预测模型113.torch",
    "normalization_method": "z-score"
}


def train(idx=1):
    from yangke.pytorch.mytorch import setup_seed
    setup_seed()
    # 加载数据
    df = pd.read_csv(settings["dataset"]["data_file"])
    logger.info(f"数据形状: {df.shape}")
    logger.info(f"数据前5行: \n{df.head()}")
    # logger.info(f"相关性分析: \n{df.corr()}")

    # 检查凝汽器热负荷和背压的相关性
    correlation = df["凝汽器热负荷"].corr(df["背压"])
    logger.info(f"凝汽器热负荷与背压相关性: {correlation}")

    # 创建数据集
    dataset = DataSetFitting(df, input_columns, output_column)
    dataset.set_standard(normal_method="z-score")  # 设置数据标准化方法

    # 创建模型
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = DataFitterNet(settings)
    model.to(device)

    # 训练模型
    model.start_train(dataset)

    # 保存模型
    model.save_yk(f"背压预测模型{idx}.torch")
    plot_backpressure_curves(model, f"yk_torch{idx}.png")


def load_trained_model(model_path):
    """
    加载已训练的神经网络模型

    :param model_path: 模型文件路径
    :return: 加载的模型对象
    """
    # 使用DataFitterNet的静态方法加载模型
    # noinspection all
    torch.serialization.add_safe_globals([pd.Series])
    loaded_model = DataFitterNet.load_yk(model_path)
    return loaded_model


def plot_backpressure_curves(model, save_path='default.png'):
    """
    绘制背压随不同参数变化的曲线图
    包括：背压随循泵运行方式、汽轮发电机功率、凝汽器热负荷、环境温度的变化
    """
    import matplotlib.pyplot as plt
    import numpy as np
    import matplotlib
    # 设置中文字体支持
    matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
    matplotlib.rcParams['axes.unicode_minus'] = False

    # 创建图形布局
    fig, axes = plt.subplots(2, 3, figsize=(15, 12))
    fig.suptitle('背压随不同参数的变化关系', fontsize=16)

    # 1. 背压随循泵运行方式变化（固定其他参数）
    base_params_1 = [173.723, 518001, 8.614, 13, 3, 2, 2]
    pump_modes = list(range(1, 7))
    bp_pump = []
    for mode in pump_modes:
        params = base_params_1.copy()
        params[6] = mode
        pred = model.prediction(params)
        bp_pump.append(pred[0] if isinstance(pred, (list, np.ndarray)) else pred)

    axes[0, 0].plot(pump_modes, bp_pump, 'bo-', linewidth=2, markersize=6)
    axes[0, 0].set_xlabel('循泵运行方式数字编码')
    axes[0, 0].set_ylabel('背压')
    axes[0, 0].set_title('背压随循泵运行方式变化')
    axes[0, 0].grid(True, alpha=0.3)

    # 2. 背压随风机台数变化
    fans = [0, 1, 2, 3, 4, 5, 6, 7, 8]
    bp_fan = []
    for fan in fans:
        params = base_params_1.copy()
        params[5] = fan
        pred = model.prediction(params)
        bp_fan.append(pred[0] if isinstance(pred, (list, np.ndarray)) else pred)

    axes[0, 1].plot(fans, bp_fan, 'go-', linewidth=2, markersize=6)
    axes[0, 1].set_xlabel('风机运行台数')
    axes[0, 1].set_ylabel('背压')
    axes[0, 1].set_title('背压随风机运行台数变化')
    axes[0, 1].grid(True, alpha=0.3)

    # 3. 背压随凝汽器热负荷变化
    heat_loads = [1708286, 1718286, 1728286, 1738286, 1748286, 1758286, 1768286, 1778286, 1788286, 1798286, 1808286]
    bp_heat = []
    for heat in heat_loads:
        params = base_params_1.copy()
        params[1] = heat
        pred = model.prediction(params)
        bp_heat.append(pred[0] if isinstance(pred, (list, np.ndarray)) else pred)

    axes[1, 0].plot(heat_loads, bp_heat, 'go-', linewidth=2, markersize=6)
    axes[1, 0].set_xlabel('凝汽器热负荷')
    axes[1, 0].set_ylabel('背压')
    axes[1, 0].set_title('背压随凝汽器热负荷变化')
    axes[1, 0].grid(True, alpha=0.3)

    # 4. 背压随环境温度变化
    powers = [17.58, 18.58, 19.58, 20.58, 21.58, 22.58, 23.58, 24.58, 25.58, 26.58, 27.58]
    bp_temp = []
    for temp in powers:
        params = base_params_1.copy()
        params[2] = temp
        pred = model.prediction(params)
        bp_temp.append(pred[0] if isinstance(pred, (list, np.ndarray)) else pred)

    axes[1, 1].plot(powers, bp_temp, 'mo-', linewidth=2, markersize=6)
    axes[1, 1].set_xlabel('环境温度')
    axes[1, 1].set_ylabel('背压')
    axes[1, 1].set_title('背压随环境温度变化')
    axes[1, 1].grid(True, alpha=0.3)

    plt.tight_layout()
    # plt.show()
    # 保存图片为PNG文件
    plt.savefig(save_path, dpi=300, bbox_inches='tight', format='png')
    print(f"图片已保存为 {save_path}")

    # 记录数据到日志
    logger.info("背压随各参数变化数据:")
    logger.info(f"循泵运行方式: {list(zip(pump_modes, bp_pump))}")
    logger.info(f"凝汽器热负荷: {list(zip(heat_loads, bp_heat))}")
    logger.info(f"环境温度: {list(zip(powers, bp_temp))}")

    # 5. 背压随功率变化
    powers = [170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180]
    bp_temp = []
    for temp in powers:
        params = base_params_1.copy()
        params[0] = temp
        pred = model.prediction(params)
        bp_temp.append(pred[0] if isinstance(pred, (list, np.ndarray)) else pred)

    axes[0, 2].plot(powers, bp_temp, 'mo-', linewidth=2, markersize=6)
    axes[0, 2].set_xlabel('汽轮发电机功率')
    axes[0, 2].set_ylabel('背压')
    axes[0, 2].set_title('背压随汽轮发电机功率变化')
    axes[0, 2].grid(True, alpha=0.3)

    plt.tight_layout()
    # plt.show()
    # 保存图片为PNG文件
    plt.savefig(save_path, dpi=300, bbox_inches='tight', format='png')
    print(f"图片已保存为 {save_path}")

    # 记录数据到日志
    logger.info("背压随各参数变化数据:")
    logger.info(f"循泵运行方式: {list(zip(pump_modes, bp_pump))}")
    logger.info(f"凝汽器热负荷: {list(zip(heat_loads, bp_heat))}")
    logger.info(f"环境温度: {list(zip(powers, bp_temp))}")


if __name__ == "__main__":
    logger.info(f"{settings=}")
    train()
    # input_columns = ["汽轮发电机功率", "凝汽器热负荷", "环境温度", "大气湿度", "冷却塔进水蝶阀开启数量", "风机运行台数",
    #                  "循泵运行方式数字编码"]
    # re_train(data_file=r"D:\国华燃机\data.csv", settings1=settings)
    #
    model = load_trained_model(f"背压预测模型113.torch")
    plot_backpressure_curves(model, f"yk_torch{115}.png")
    print(f"{model.prediction([173.723, 518001, 8.614, 13, 3, 2, 2])=}") # 3.5
    # 平均相对误差为：0.025340，最大相对误差为0.864250