import torch
import torch.nn as nn
import torch.optim as optim
from ucimlrepo import fetch_ucirepo
from sklearn.model_selection import KFold
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler
import numpy as np

# 定义BP神经网络类
class BPNeuralNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(BPNeuralNetwork, self).__init__()
        # 初始化输入层、隐藏层和输出层的大小
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size

        # 初始化权重和偏置
        self.fc1 = torch.nn.Linear(self.input_size, self.hidden_size)  # 输入层到隐藏层的全连接层
        self.fc2 = torch.nn.Linear(self.hidden_size, self.output_size)  # 隐藏层到输出层的全连接层

    def forward(self, x):
        # 前向传播
        x = torch.sigmoid(self.fc1(x))  # 隐藏层的输出
        x = torch.sigmoid(self.fc2(x))  # 输出层的输出
        return x

# 将标签转换为one-hot编码
def one_hot_encode(y):
    # 确保 y 中的每个元素是字符串
    y = np.array([str(label) for label in y])
    # 创建一个标签到整数的映射
    label_to_int = {label: i for i, label in enumerate(np.unique(y))}
    # 将标签转换为整数
    y_int = np.array([label_to_int[label] for label in y])
    n_values = np.max(y_int) + 1
    return torch.eye(n_values)[y_int]  # 返回one-hot编码

# fetch dataset 
wine_quality = fetch_ucirepo(id=186) 
  
# data (as pandas dataframes) 
X = wine_quality.data.features.values  # 特征数据
y = wine_quality.data.targets.values  # 标签数据

# 特征缩放
scaler = StandardScaler()
X = scaler.fit_transform(X)  # 标准化特征数据

# 对标签进行one-hot编码
y_encoded = one_hot_encode(y)

# 将数据转换为PyTorch张量并移至GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
X = torch.tensor(X, dtype=torch.float32).to(device)
y_encoded = y_encoded.to(device)

# 十折交叉验证
kf = KFold(n_splits=10, shuffle=True, random_state=42)
accuracies = []

for train_index, test_index in kf.split(X):
    X_train, X_test = X[train_index], X[test_index]  # 训练集和测试集的特征数据
    y_train, y_test = y_encoded[train_index], y_encoded[test_index]  # 训练集和测试集的标签数据

    # 创建并训练BP神经网络
    nn = BPNeuralNetwork(input_size=X_train.shape[1], hidden_size=20, output_size=y_train.shape[1]).to(device)  # 修改隐藏层大小
    criterion = torch.nn.MSELoss()  # 使用均方误差损失函数
    optimizer = optim.SGD(nn.parameters(), lr=0.0001)  # 使用随机梯度下降优化器

    for epoch in range(50000):  # 增加训练轮数
        optimizer.zero_grad()  # 清零梯度
        output = nn(X_train)  # 前向传播
        loss = criterion(output, y_train)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新权重和偏置

        if epoch % 1000 == 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}')  # 打印损失

    # 预测并计算准确率
    with torch.no_grad():
        predictions = nn(X_test)
        accuracy = accuracy_score(torch.argmax(y_test, dim=1).cpu().numpy(), torch.argmax(predictions, dim=1).cpu().numpy())  # 计算准确率
    accuracies.append(accuracy)  # 存储每次交叉验证的准确率

print(f'Average Accuracy: {np.mean(accuracies)}')  # 打印平均准确率