# E:\repos\trae\lib4python\yangke\sis\BeiJingShenHua\神经网络训练4.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
from yangke.common.config import logger

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


class SimpleNN(nn.Module):
    """
    普通神经网络模型（不考虑单调性约束）
    """

    def __init__(self, input_size, hidden_sizes, output_size):
        """
        初始化模型
        Args:
            input_size: 输入特征数量
            hidden_sizes: 隐藏层大小列表
            output_size: 输出大小
        """
        super(SimpleNN, self).__init__()

        # 创建隐藏层
        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)

        # 初始化权重
        self._initialize_weights()

    def _initialize_weights(self):
        """初始化权重"""
        for layer in self.hidden_layers:
            nn.init.xavier_uniform_(layer.weight)
            layer.bias.data.fill_(0.01)

        # 输出层权重初始化
        nn.init.xavier_uniform_(self.output_layer.weight)
        self.output_layer.bias.data.fill_(0.01)

    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 SimpleModelPredictor:
    """
    普通神经网络模型预测器
    加载一次模型，支持多次预测
    """

    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')

        # 重建模型
        self.model = SimpleNN(
            input_size=9,  # 包括原始特征和构造特征
            hidden_sizes=[128, 256, 128, 64],
            output_size=1
        )
        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_df = pd.DataFrame(input_data, columns=["凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数",
                                                         "循泵运行方式数字编码"])
            input_df['x1_sq'] = input_df['凝汽器热负荷'] ** 2
            input_df['x1_x2'] = input_df['凝汽器热负荷'] * input_df['环境温度']
            input_df['x4_x5'] = input_df['风机运行台数'] * input_df['循泵运行方式数字编码']
            input_df['x5_sqrt'] = np.sqrt(input_df['循泵运行方式数字编码'])  # 确保 x5 >= 0

            # 标准化输入数据
            input_scaled = self.scaler_X.transform(input_df.values)
            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:
        标准化后的数据和特征名称
    """
    # 确保列名正确
    df.columns = ["凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数", "循泵运行方式数字编码", "背压"]

    # 构造非线性特征
    df['x1_sq'] = df['凝汽器热负荷'] ** 2
    df['x1_x2'] = df['凝汽器热负荷'] * df['环境温度']
    df['x4_x5'] = df['风机运行台数'] * df['循泵运行方式数字编码']
    df['x5_sqrt'] = np.sqrt(df['循泵运行方式数字编码'])  # 注意：x5 >= 0

    # 输入特征（包括原始和构造特征）
    input_features = [
        "凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数", "循泵运行方式数字编码",
        "x1_sq", "x1_x2", "x4_x5", "x5_sqrt"
    ]
    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, output_feature


def train_model(X, y, scaler_X, scaler_y, input_features, output_features):
    """
    训练神经网络模型
    Returns:
        训练好的模型和预处理器
    """
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, 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=32, shuffle=True)

    # 定义模型（普通神经网络）
    model = SimpleNN(input_size=9, hidden_sizes=[128, 256, 128, 64], output_size=1)

    # 定义优化器（降低学习率）
    optimizer = optim.Adam(model.parameters(), lr=0.0001)

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

    # 训练模型
    num_epochs = 50000
    best_loss = float('inf')
    patience = 1000
    patience_counter = 0

    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():
                predictions_scaled = model(X_test_tensor)  # 预测的标准化y值
                predictions = scaler_y.inverse_transform(predictions_scaled.numpy())  # 预测的原始y值
                test_loss = criterion(predictions_scaled, y_test_tensor)  # 标准化MSE
                mse_original = np.mean((predictions - y_test) ** 2)
                rmse_original = np.sqrt(mse_original)

            logger.debug(f'Epoch [{epoch + 1}/{num_epochs}], '
                         f'训练集标准化MSE: {total_loss / len(train_loader):.8f}, '
                         f'测试集标准化MSE: {test_loss:.8f}, '
                         f'测试集原始MSE: {rmse_original:.8f}')

            # 早停机制
            if mse_original < best_loss:
                best_loss = mse_original
                patience_counter = 0
                # 保存最佳模型
                torch.save({
                    'model_state_dict': model.state_dict(),
                    'scaler_X': scaler_X,
                    'scaler_y': scaler_y,
                    'input_features': input_features,
                    'output_features': output_features
                }, 'best_model.pth')
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    logger.debug("早停触发，提前结束训练")
                    break

    return model


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_df = pd.DataFrame(input_data, columns=["凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数",
                                                     "循泵运行方式数字编码"])
        input_df['x1_sq'] = input_df['凝汽器热负荷'] ** 2
        input_df['x1_x2'] = input_df['凝汽器热负荷'] * input_df['环境温度']
        input_df['x4_x5'] = input_df['风机运行台数'] * input_df['循泵运行方式数字编码']
        input_df['x5_sqrt'] = np.sqrt(input_df['循泵运行方式数字编码'])

        # 标准化输入数据
        input_scaled = scaler_X.transform(input_df.values)
        input_tensor = torch.FloatTensor(input_scaled)

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

    return prediction


def plot_prediction_check(predictor: SimpleModelPredictor, base_params, param_ranges, save_path):
    """
    绘制背压随各个参数的变化曲线
    Args:
    """
    model = predictor.model

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

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

    # 对每个参数绘制曲线
    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()

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


def main():
    """
    主程序入口
    """
    # 读取数据
    df = pd.read_excel(r'C:\Users\54067\WPSDrive\217145378\WPS云盘\6科研项目\神经网络测试数据.xlsx',
                       sheet_name='训练数据')

    # 数据预处理（包含特征工程）
    X, y, scaler_X, scaler_y, input_features, output_features = preprocess_data(df)
    logger.debug(f"输入特征：{input_features}")

    # 训练模型
    model_name = 'simple_nn_model'
    model_path = f'{model_name}.pth'
    model = train_model(X, y, scaler_X, scaler_y, input_features, output_features)

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

    logger.debug(f"模型已保存为 {model_path}")

    # 示例预测
    base_params = [1, 10.1, 2, 3, 5]
    prediction = predict(model, scaler_X, scaler_y, [base_params])
    logger.debug(f"示例预测结果: {prediction}")

    # 演示使用 SimpleModelPredictor 进行多次预测
    logger.debug("\n=== 使用 SimpleModelPredictor 进行多次预测 ===")
    predictor = SimpleModelPredictor(model_path)

    # 绘制预测检查图
    param_ranges = {
        "凝汽器热负荷": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        "环境温度": [10.1, 11, 11.9, 12.8, 13.7, 14.6, 15.5, 16.4, 17.3, 18.2],
        "大气湿度": [2, 3, 4, 5, 6, 7, 8, 9, 10],
        "循泵运行方式数字编码": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12],
        "风机运行台数": [5, 6, 7, 8, 9, 10, 11, 12, 13]
    }

    plot_prediction_check(predictor, base_params, param_ranges, f"{model_name}.png")


if __name__ == "__main__":
    main()
