import torch, os
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

torch.serialization.add_safe_globals([StandardScaler])


# 定义神经网络模型
class MonotonicNeuralNetwork(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super(MonotonicNeuralNetwork, self).__init__()

        # 创建网络层
        layers = []
        prev_size = input_size
        for hidden_size in hidden_sizes:
            layers.append(nn.Linear(prev_size, hidden_size))
            layers.append(nn.ReLU())
            prev_size = hidden_size

        layers.append(nn.Linear(prev_size, output_size))
        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)


# 数据预处理函数
def preprocess_data(df):
    # 输入特征
    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


# 自定义损失函数，加入单调性约束
class MonotonicLoss(nn.Module):
    def __init__(self, monotonic_constraints):
        """
        monotonic_constraints: 列表，指定每个输入特征的单调性
        1: 单调递增, -1: 单调递减, 0: 无约束
        """
        super(MonotonicLoss, self).__init__()
        self.monotonic_constraints = torch.tensor(monotonic_constraints, dtype=torch.float32)
        self.mse = nn.MSELoss()

    def forward(self, predictions, targets, inputs, model):
        # 基本的MSE损失
        mse_loss = self.mse(predictions, targets)

        # 计算单调性约束损失
        monotonic_loss = 0
        if torch.sum(torch.abs(self.monotonic_constraints)) > 0:
            # 使用数值梯度近似计算偏导数
            inputs.requires_grad_(True)
            outputs = model(inputs)

            # 计算梯度
            gradients = torch.autograd.grad(outputs, inputs,
                                            grad_outputs=torch.ones_like(outputs),
                                            create_graph=True)[0]

            # 应用单调性约束
            constraint_violation = torch.relu(-self.monotonic_constraints * gradients)
            monotonic_loss = torch.mean(constraint_violation)

            # 重置梯度要求
            inputs.requires_grad_(False)

        # 总损失 = MSE损失 + 单调性约束损失
        total_loss = mse_loss + 0.1 * monotonic_loss  # 0.1是单调性约束的权重，可调整

        return total_loss, mse_loss, monotonic_loss


# 训练函数
def train_model():
    # 读取数据
    df = pd.read_csv(r"D:\国华燃机\data.csv")
    df = df[["凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数", "循泵运行方式数字编码", "背压"]]
    # 数据预处理
    X, y, scaler_X, scaler_y = preprocess_data(df)

    # 划分训练集和测试集
    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 = MonotonicNeuralNetwork(input_size=6, hidden_sizes=[128, 64, 32], output_size=1)

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

    # 定义单调性约束 (根据您的要求)
    # ["汽轮发电机功率", "凝汽器热负荷", "环境温度", "大气湿度", "风机运行台数", "循泵运行方式数字编码"]
    # [背压]随[汽轮发电机功率]单调递增 -> 1
    # [背压]随[凝汽器热负荷]单调递增 -> 1
    # [背压]随[环境温度]单调递增 -> 1
    # [背压]随[风机运行台数]单调递减 -> -1
    # [背压]随[循泵运行方式数字编码]单调递减 -> -1
    # [背压]与[大气湿度]关系未指定 -> 0
    monotonic_constraints = [1, 1, 1, 0, -1, -1]

    # 创建损失函数
    criterion = MonotonicLoss(monotonic_constraints)

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

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

            # 前向传播
            predictions = model(batch_X)

            # 计算损失
            loss, mse_loss, monotonic_loss = criterion(predictions, batch_y, batch_X, model)

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

            total_loss += loss.item()
            total_mse_loss += mse_loss.item()
            total_monotonic_loss += monotonic_loss.item()

        # 每10个epoch打印一次训练信息
        if (epoch + 1) % 10 == 0:
            model.eval()
            with torch.no_grad():
                test_predictions = model(X_test_tensor)
                test_loss = criterion.mse(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}, '
                  f'MSE Loss: {total_mse_loss / len(train_loader):.4f}, '
                  f'Monotonic Loss: {total_monotonic_loss / len(train_loader):.4f}')

    return model, scaler_X, scaler_y


# 预测函数
def predict(model, scaler_X, scaler_y, input_data):
    """
    使用训练好的模型进行预测
    input_data: numpy数组，形状为 (n_samples, 6)
    """
    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 load_model(model_path):
    """
    加载已训练的模型和标准化器

    Args:
        model_path: 模型文件路径

    Returns:
        model: 加载的模型
        scaler_X: 输入标准化器
        scaler_y: 输出标准化器
    """
    # 使用weights_only=False加载模型（注意：仅在信任模型来源时使用）
    checkpoint = torch.load(model_path, weights_only=False)

    # 初始化模型结构
    model = MonotonicNeuralNetwork(input_size=6, hidden_sizes=[128, 64, 32], output_size=1)

    # 加载模型权重
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()

    # 获取标准化器
    scaler_X = checkpoint['scaler_X']
    scaler_y = checkpoint['scaler_y']

    return model, scaler_X, scaler_y


def plot_monotonicity_curves(predictor, base_params: list, param_ranges: dict, save_path="default.png"):
    """
    绘制神经网络模型预测结果的曲线图，主要是输出随输入的变化曲线。
    base_params为默认神经网络的输入参数，为了绘制变化规律，分别按照param_ranges这个对象中指定的每个参数的取值列表，用神经网络模型
    计算输出，并绘制输出随每个参数取值列表中数值变化的曲线。

    绘制背压随各个输入参数变化的曲线图，验证单调性约束是否满足
    :params base_params: list
    :params param_ranges: dict
    """
    # 参数索引映射
    param_indices = {
        "汽轮发电机功率": 0,
        "凝汽器热负荷": 1,
        "环境温度": 2,
        "大气湿度": 3,
        "风机运行台数": 4,
        "循泵运行方式数字编码": 5
    }

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

    # 对每个参数绘制曲线
    for idx, (param_name, param_values) in enumerate(param_ranges.items()):
        # 存储预测结果
        predictions = []

        # 对每个参数值进行预测
        for value in param_values:
            # 复制基准参数并修改对应参数值
            params = base_params.copy()
            params[param_indices[param_name]] = value

            # 转换为numpy数组并预测
            input_data = np.array([params])
            pred = predictor.predict(input_data)
            predictions.append(pred[0][0])

        # 绘制曲线
        ax = axes[idx]
        ax.plot(param_values, predictions, 'bo-', linewidth=2, markersize=6)
        ax.set_xlabel(param_name)
        ax.set_ylabel('背压')
        ax.set_title(f'背压随{param_name}变化')
        ax.grid(True, alpha=0.3)

        # 打印单调性检查结果
        # 根据约束定义检查单调性
        expected_monotonicity = {
            "汽轮发电机功率": 1,  # 单调递增
            "凝汽器热负荷": 1,  # 单调递增
            "环境温度": 1,  # 单调递增
            "循泵运行方式数字编码": -1,  # 单调递减
            "风机运行台数": -1  # 单调递减
        }

        if param_name in expected_monotonicity:
            if expected_monotonicity[param_name] == 1:
                # 应该单调递增
                is_monotonic = all(predictions[i] <= predictions[i + 1] for i in range(len(predictions) - 1))
            else:
                # 应该单调递减
                is_monotonic = all(predictions[i] >= predictions[i + 1] for i in range(len(predictions) - 1))
        else:
            # 无约束
            is_monotonic = True

        print(f"{param_name}单调性检查: {'通过' if is_monotonic else '不通过'}")
        print(f"  数值序列: {[round(p, 2) for p in predictions]}")

    # 隐藏多余的子图
    axes[5].set_visible(False)

    plt.tight_layout()
    plt.suptitle('背压随各参数变化关系及单调性验证', fontsize=16, y=1.02)

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

    # plt.show()


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

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

        # 重建模型
        self.model = model_class(
            input_size=6,
            hidden_sizes=[128, 64, 32],
            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, 6)
        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 outer_train(i):
    model_path = f'monotonic_nn_model1{i}.pth'
    # 训练模型
    model, scaler_X, scaler_y = train_model()

    # 保存模型
    torch.save({
        'model_state_dict': model.state_dict(),
        'scaler_X': scaler_X,
        'scaler_y': scaler_y
    }, model_path)

    print("模型已保存为 monotonic_nn_model.pth")
    predictor = ModelPredictor(MonotonicNeuralNetwork, model_path)
    base_params = [280, 1708286, 17.58, 75, 7, 1]
    param_ranges = {
        "汽轮发电机功率": [280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290],
        "凝汽器热负荷": [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],
        "循泵运行方式数字编码": [0, 1, 2, 3, 4, 5, 6],
        "风机运行台数": [0, 1, 2, 3, 4, 5, 6, 7, 8]
    }
    plot_monotonicity_curves(predictor, base_params, param_ranges, f"result{i}.png")


# 主程序
if __name__ == "__main__":
    for i in range(20):
        outer_train(i)
