# 和神经网络训练2.py相比，删除了汽轮发电机功率
import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt
import matplotlib

# 设置中文字体支持
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'DejaVu Sans']
matplotlib.rcParams['axes.unicode_minus'] = False


class GuaranteedMonotonicNN(nn.Module):
    """
    保证单调性的神经网络模型
    通过架构设计确保输出对指定输入特征的单调性
    """

    def __init__(self, input_size, hidden_sizes, output_size, monotonic_constraints):
        """
        初始化模型
        Args:
            input_size: 输入特征数量
            hidden_sizes: 隐藏层大小列表
            output_size: 输出大小
            monotonic_constraints: 单调性约束列表，1表示递增，-1表示递减，0表示无约束
        """
        super(GuaranteedMonotonicNN, self).__init__()
        self.monotonic_constraints = torch.tensor(monotonic_constraints, dtype=torch.float32)

        # 创建隐藏层
        self.hidden_layers = nn.ModuleList()
        prev_size = input_size

        for hidden_size in hidden_sizes:
            layer = nn.Linear(prev_size, hidden_size)
            self.hidden_layers.append(layer)
            prev_size = hidden_size

        # 输出层
        self.output_layer = nn.Linear(prev_size, output_size)

    def forward(self, x):
        """
        前向传播
        Args:
            x: 输入张量
        Returns:
            输出张量
        """
        # 通过隐藏层
        hidden = x
        for layer in self.hidden_layers:
            hidden = torch.relu(layer(hidden))

        # 输出层
        output = self.output_layer(hidden)
        return output


class MonotonicModelPredictor:
    """
    单调性神经网络模型预测器
    加载一次模型，支持多次预测
    """

    def __init__(self, model_path):
        """
        初始化预测器并加载模型
        Args:
            model_path: 模型文件路径
        """
        # 加载模型和预处理器
        checkpoint = torch.load(model_path, weights_only=False)

        # 获取特征信息
        self.input_features = checkpoint.get('input_features')
        self.output_features = checkpoint.get('output_features')

        # 重建模型
        monotonic_constraints = [1, 1, 0, -1, -1]  # 删除第一个参数后的约束
        self.model = GuaranteedMonotonicNN(
            input_size=5,  # 从6改为5
            hidden_sizes=[128, 64, 32],
            output_size=1,
            monotonic_constraints=monotonic_constraints
        )
        self.model.load_state_dict(checkpoint['model_state_dict'])

        # 获取预处理器
        self.scaler_X = checkpoint['scaler_X']
        self.scaler_y = checkpoint['scaler_y']

        # 设置为评估模式
        self.model.eval()

    def predict(self, input_data):
        """
        使用已加载的模型进行预测
        Args:
            input_data: numpy数组，形状为 (n_samples, 5)
        Returns:
            预测结果
        """
        with torch.no_grad():
            # 标准化输入数据
            input_scaled = self.scaler_X.transform(input_data)
            input_tensor = torch.FloatTensor(input_scaled)

            # 预测
            prediction_scaled = self.model(input_tensor)
            prediction = self.scaler_y.inverse_transform(prediction_scaled.numpy())

        return prediction

    def get_feature_info(self):
        """
        获取模型的特征信息
        Returns:
            输入特征列表和输出特征列表
        """
        return self.input_features, self.output_features


def preprocess_data(df):
    """
    数据预处理函数
    Args:
        df: 原始数据DataFrame
    Returns:
        标准化后的数据和特征名称
    """
    # 输入特征（删除"汽轮发电机功率"）
    input_features = ["凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数", "循泵运行方式数字编码"]
    # 输出特征
    output_feature = ["背压"]

    X = df[input_features].values
    y = df[output_feature].values

    # 标准化数据
    scaler_X = StandardScaler()
    scaler_y = StandardScaler()

    X_scaled = scaler_X.fit_transform(X)
    y_scaled = scaler_y.fit_transform(y)

    return X_scaled, y_scaled, scaler_X, scaler_y, input_features, ["背压"]


def train_model():
    """
    训练神经网络模型
    Returns:
        训练好的模型和预处理器
    """
    # 读取数据（删除"汽轮发电机功率"）
    df = pd.read_csv(r"D:\国华燃机\data.csv")
    df = df[["凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数", "循泵运行方式数字编码", "背压"]]

    # 数据预处理
    X, y, scaler_X, scaler_y, input_features, output_features = preprocess_data(df)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 转换为PyTorch张量
    X_train_tensor = torch.FloatTensor(X_train)
    y_train_tensor = torch.FloatTensor(y_train)
    X_test_tensor = torch.FloatTensor(X_test)
    y_test_tensor = torch.FloatTensor(y_test)

    # 创建数据加载器
    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)

    # 定义单调性约束（删除第一个参数的约束）
    monotonic_constraints = [1, 1, 0, -1, -1]

    # 定义模型（保证单调性）
    model = GuaranteedMonotonicNN(input_size=5, hidden_sizes=[128, 64, 32], output_size=1,
                                  monotonic_constraints=monotonic_constraints)

    # 定义优化器
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 使用普通的MSE损失函数
    criterion = nn.MSELoss()

    # 训练模型
    num_epochs = 100
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0

        for batch_X, batch_y in train_loader:
            optimizer.zero_grad()

            # 前向传播
            predictions = model(batch_X)

            # 计算损失
            loss = criterion(predictions, batch_y)

            # 反向传播
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        # 每10个epoch打印一次训练信息
        if (epoch + 1) % 10 == 0:
            model.eval()
            with torch.no_grad():
                test_predictions = model(X_test_tensor)
                test_loss = criterion(test_predictions, y_test_tensor)

            print(f'Epoch [{epoch + 1}/{num_epochs}], '
                  f'Train Loss: {total_loss / len(train_loader):.4f}, '
                  f'Test MSE: {test_loss:.4f}')

    return model, scaler_X, scaler_y, input_features, output_features, df


def predict(model, scaler_X, scaler_y, input_data):
    """
    使用训练好的模型进行预测
    Args:
        model: 训练好的模型
        scaler_X: 输入数据标准化器
        scaler_y: 输出数据标准化器
        input_data: numpy数组，形状为 (n_samples, 5)
    Returns:
        预测结果
    """
    model.eval()
    with torch.no_grad():
        # 标准化输入数据
        input_scaled = scaler_X.transform(input_data)
        input_tensor = torch.FloatTensor(input_scaled)

        # 预测
        prediction_scaled = model(input_tensor)
        prediction = scaler_y.inverse_transform(prediction_scaled.numpy())

    return prediction


def plot_monotonicity_check(predictor: MonotonicModelPredictor, base_params, param_ranges):
    """
    绘制背压随各个参数的变化曲线，检查单调性
    Args:
    """
    model = predictor.model

    # 创建子图
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    axes = axes.flatten()

    # 参数索引映射（删除"汽轮发电机功率"后的映射）
    param_indices = {
        "凝汽器热负荷": 0,
        "环境温度": 1,
        "大气湿度": 2,
        "风机运行台数": 3,
        "循泵运行方式数字编码": 4
    }

    # 定义单调性约束
    monotonic_constraints = {
        "凝汽器热负荷": 1,  # 单调递增
        "环境温度": 1,  # 单调递增
        "大气湿度": 0,  # 无约束
        "风机运行台数": -1,  # 单调递减
        "循泵运行方式数字编码": -1  # 单调递减
    }

    # 对每个参数绘制曲线
    for idx, (param_name, param_values) in enumerate(param_ranges.items()):
        ax = axes[idx]

        predictions = []
        for val in param_values:
            # 构建测试输入：使用基准参数，当前特征使用测试值
            test_input = base_params.copy()
            test_input[param_indices[param_name]] = val
            test_input = np.array([test_input])

            # 预测
            pred = predictor.predict(test_input)
            predictions.append(pred[0, 0])

        # 绘制曲线
        ax.plot(param_values, predictions, 'b-', linewidth=2, marker='o', label='模型预测')

        ax.set_xlabel(param_name)
        ax.set_ylabel('背压')
        ax.set_title(f'背压 vs {param_name}')
        ax.grid(True, alpha=0.3)
        ax.legend()

        # 添加单调性标注
        constraint = monotonic_constraints[param_name]
        if constraint == 1:
            ax.text(0.05, 0.95, '应单调递增', transform=ax.transAxes,
                    bbox=dict(boxstyle="round", facecolor='green', alpha=0.3))
        elif constraint == -1:
            ax.text(0.05, 0.95, '应单调递减', transform=ax.transAxes,
                    bbox=dict(boxstyle="round", facecolor='red', alpha=0.3))
        else:
            ax.text(0.05, 0.95, '无单调性要求', transform=ax.transAxes,
                    bbox=dict(boxstyle="round", facecolor='yellow', alpha=0.3))

    plt.tight_layout()
    plt.savefig('monotonicity_check.png', dpi=300, bbox_inches='tight')
    plt.show()

    print("单调性检查图已保存为 monotonicity_check.png")


def main():
    """
    主程序入口
    """
    # 训练模型
    model, scaler_X, scaler_y, input_features, output_features, df = train_model()

    # 保存模型和预处理器
    torch.save({
        'model_state_dict': model.state_dict(),
        'scaler_X': scaler_X,
        'scaler_y': scaler_y,
        'input_features': input_features,
        'output_features': output_features
    }, 'guaranteed_monotonic_nn_model.pth')

    print("模型已保存为 guaranteed_monotonic_nn_model.pth")

    # 示例预测（删除第一个参数）
    example_input = np.array([[1708286, 17.58, 75, 7, 1]])  # 示例值
    prediction = predict(model, scaler_X, scaler_y, example_input)
    print(f"示例预测结果: {prediction}")

    # 演示使用 MonotonicModelPredictor 进行多次预测
    print("\n=== 使用 MonotonicModelPredictor 进行多次预测 ===")
    predictor = MonotonicModelPredictor('guaranteed_monotonic_nn_model.pth')
    base_params = [1708286, 17.58, 75, 7, 1]  # 删除第一个参数
    param_ranges = {
        "凝汽器热负荷": [1708286, 1718286, 1728286, 1738286, 1748286, 1758286, 1768286,
                         1778286, 1788286, 1798286,
                         1808286],
        "环境温度": [17.58, 18.58, 19.58, 20.58, 21.58, 22.58, 23.58, 24.58, 25.58, 26.58,
                     27.58],
        "循泵运行方式数字编码": [1, 2, 3, 4, 5, 6],
        "风机运行台数": [0, 1, 2, 3, 4, 5, 6, 7, 8]
    }
    # 绘制单调性检查图
    plot_monotonicity_check(predictor, base_params, param_ranges)


if __name__ == "__main__":
    main()
    # predictor = MonotonicModelPredictor('guaranteed_monotonic_nn_model.pth')

    # plot_monotonicity_check(predictor, base_params, param_ranges)
