"""
文件名: train_gpu_blance_10features.py

训练部分代码

作者: 王春林
创建日期: 2023年10月18日
最后修改日期: 2023年10月20日
版本号: 1.0.0

"""

import numpy as np
import torch
from torch import nn

# 定义 MLP 网络
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(10, 32),  # 输入层
            nn.ReLU(),  # 激活函数
            nn.Linear(32, 128),  # 隐藏层
            nn.ReLU(),  # 激活函数
            nn.Linear(128, 32),  # 隐藏层
            nn.ReLU(),  # 激活函数
            nn.Linear(32, 4),  # 输出层，4个类别
        )

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

# 检查GPU是否可用
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")

def predict_with_model(data):
    # # 读取特征和标签
    # data = pd.read_excel('feature_label.xlsx')

    # 以下是你的特征名
    feature_names = ["强迫症状数字化", "人际关系敏感数字化", "抑郁数字化", "多因子症状", "母亲教养方式数字化", "父亲教养方式数字化", "自评家庭经济条件数字化", "有无心理治疗（咨询）史数字化", "学业情况数字化", "出勤情况数字化"]

    # 定义特征权重列表
    feature_weights = [0.135, 0.085, 0.08, 0.2, 0.09, 0.09, 0.06, 0.06, 0.08, 0.12]

    # 找到最大值
    max_value = max(feature_weights)

    # 缩放权重
    feature_weights_scaled = [x / max_value for x in feature_weights]

    # 打印缩放后的特征权重
    # print("Scaled Feature Weights:", feature_weights_scaled)

    # 将特征和标签分开，并做归一化处理
    X = data[feature_names].values

    # 分别乘以权重，放在归一化后
    for i in range(len(feature_names)):
        X[:, i] = X[:, i] * feature_weights_scaled[i]
    
    model = MLP().to(device)
    model.load_state_dict(torch.load('psychology.pth', map_location=device))
    model.eval()

    # 转换输入数据为 Tensor，并放到设备上
    X_tensor = torch.from_numpy(X).float().to(device)

    # 推理
    with torch.no_grad():
        outputs = model(X_tensor)

    # 获取预测结果
    _, predictions = torch.max(outputs, 1)

    # 实际类别从1开始，程序类别从0开始
    predictions += 1

    # 打印预测结果
    # print("预测结果:", predictions.cpu().numpy())
    # 返回预测结果
    return predictions.cpu().numpy().tolist()