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
monotonic_constraints = [0, 0, 0, 0, 0]  # 删除第一个参数后的约束

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()
        self.weight_masks = []  # 用于存储权重掩码
        prev_size = input_size

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

            # 创建权重掩码，确保单调性约束的权重为正
            mask = torch.ones_like(layer.weight)
            self.weight_masks.append(mask)
            prev_size = hidden_size

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

        # 初始化权重为正数
        self._initialize_positive_weights()

    def _initialize_positive_weights(self):
        """初始化权重为正数以确保单调性"""
        with torch.no_grad():
            for layer, mask in zip(self.hidden_layers, self.weight_masks):
                # 对于有单调性约束的输入特征，确保对应权重为正
                layer.weight.data.abs_()
                layer.weight.data.clamp_(min=0)  # 确保严格为正

            # 输出层权重也初始化为正
            self.output_layer.weight.data.abs_()
            self.output_layer.weight.data.clamp_(min=0)

    def forward(self, x):
        """
        前向传播
        Args:
            x: 输入张量
        Returns:
            输出张量
        """
        # 应用单调性约束到输入
        constrained_x = x.clone()
        for i, constraint in enumerate(self.monotonic_constraints):
            if constraint == -1:  # 对于递减特征，取反以变为递增
                constrained_x[:, i] = -constrained_x[:, i]

        # 通过隐藏层
        hidden = constrained_x
        for i, (layer, mask) in enumerate(zip(self.hidden_layers, self.weight_masks)):
            # 确保权重保持为正（在前向传播中强制约束）
            with torch.no_grad():
                layer.weight.data.abs_()
                layer.weight.data.clamp_(min=0)

            hidden = torch.relu(layer(hidden))

        # 输出层（确保权重为正）
        with torch.no_grad():
            self.output_layer.weight.data.abs_()
            self.output_layer.weight.data.clamp_(min=0)

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

        # 重建模型

        self.model = GuaranteedMonotonicNN(
            input_size=5,  # 从6改为5
            hidden_sizes=[64, 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(X, y, scaler_X, scaler_y):
    """
    训练神经网络模型
    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=64, shuffle=True)

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

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

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

    # 训练模型
    num_epochs = 1000000
    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()

            # 在更新前确保权重保持为正
            with torch.no_grad():
                for layer in model.hidden_layers:
                    layer.weight.grad = layer.weight.grad * (layer.weight.data > 0).float() + \
                                        layer.weight.grad * (layer.weight.data <= 0).float() * 0.1

                # 输出层梯度处理
                model.output_layer.weight.grad = model.output_layer.weight.grad * \
                                                 (model.output_layer.weight.data > 0).float() + \
                                                 model.output_layer.weight.grad * \
                                                 (model.output_layer.weight.data <= 0).float() * 0.1

            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
                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):.6f}, '
                         f'测试集标准化MSE: {test_loss:.6f}, '
                         f'测试集原始MSE: {rmse_original:.6f}')
            if rmse_original < 0.0001:
                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_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, 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
    }

    # 定义单调性约束
    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(save_path, dpi=300, bbox_inches='tight')
    plt.show()

    logger.debug(f"单调性检查图已保存为 {save_path}")


def main():
    """
    主程序入口
    """
    # 读取数据（删除"汽轮发电机功率"）
    # df = pd.read_csv(r"D:\国华燃机\data.csv")
    df = pd.read_excel(r'C:\Users\54067\WPSDrive\217145378\WPS云盘\6科研项目\神经网络测试数据.xlsx',
                       sheet_name='训练数据')
    df.columns = ["凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数", "循泵运行方式数字编码", "背压"]
    df = df[["凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数", "循泵运行方式数字编码", "背压"]]
    # 数据预处理
    X, y, scaler_X, scaler_y, input_features, output_features = preprocess_data(df)
    logger.debug(f"输入特征：{input_features}, {scaler_X=}, {scaler_y}")
    for i in [6]:
        # 训练模型
        model_name = f'guaranteed_monotonic_nn_model{i}'
        model_path = f'{model_name}.pth'
        model = train_model(X, y, scaler_X, scaler_y)
        # 保存模型和预处理器
        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]  # 删除第一个参数
        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]
        }

        # 示例预测（删除第一个参数）
        prediction = predict(model, scaler_X, scaler_y, [base_params])
        logger.debug(f"示例预测结果: {prediction}")

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

        # 绘制单调性检查图
        plot_monotonicity_check(predictor, base_params, param_ranges, f"{model_name}.png")


if __name__ == "__main__":
    main()
