import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import pandas as pd
from sklearn import preprocessing

# 假设数据集已经被加载为Pandas DataFrame，下面是预处理步骤
def preprocess_data(data):
    # 特征描述
    numerical_features = ['Ball Speed (mph)', 'Rally Count', 'Player Position (X, Y)']
    categorical_features = ['Serve Type', 'Match Surface', 'Player Hand']

    # 数值特征归一化
    scaler = preprocessing.StandardScaler()
    data[numerical_features] = scaler.fit_transform(data[numerical_features])

    # 类别特征独热编码
    data = pd.get_dummies(data, columns=categorical_features)

    # 将特征和标签分开
    X = data.drop('Target', axis=1)
    y = data['Target']
    
    return X, y

# 假设df是包含所有数据的DataFrame
X, y = preprocess_data(df)

# 将数据转换为PyTorch张量
X_tensor = torch.tensor(X.values).float()
y_tensor = torch.tensor(y.values).float()

# 创建数据加载器
dataset = TensorDataset(X_tensor, y_tensor)
train_loader = DataLoader(dataset, batch_size=64, shuffle=True)

# 定义CNN组件
class CNNComponent(nn.Module):
    def __init__(self):
        super(CNNComponent, self).__init__()
        self.conv1 = nn.Conv1d(in_channels=1, out_channels=64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=128, kernel_size=3, padding=1)
        self.pool = nn.MaxPool1d(2)
        self.fc1 = nn.Linear(128 * 50, 256)

    def forward(self, x):
        x = x.unsqueeze(1)  # 添加一个通道维度
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = x.view(x.size(0), -1)  # 展平
        x = F.relu(self.fc1(x))
        return x

# 定义LSTM组件
class LSTMComponent(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers):
        super(LSTMComponent, self).__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)

    def forward(self, x):
        x, _ = self.lstm(x)
        return x[:, -1, :]  # 取序列的最后一个时间步

# 定义多头注意力组件
class MultiHeadAttentionComponent(nn.Module):
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttentionComponent, self).__init__()
        self.num_heads = num_heads
        self.attention = nn.MultiheadAttention(d_model, num_heads)

    def forward(self, x):
        x = x.transpose(0, 1)  # 调整为（seq_len, batch, feature）
        attn_output, _ = self.attention(x, x, x)
        return attn_output[-1]  # 取序列的最后一个时间步

# 定义混合神经网络模型
class HybridNeuralNetwork(nn.Module):
    def __init__(self):
        super(HybridNeuralNetwork, self).__init__()
        self.cnn = CNNComponent()
        self.lstm = LSTMComponent(input_dim=256, hidden_dim=256, num_layers=1)
        self.attention = MultiHeadAttentionComponent(d_model=256, num_heads=8)
        self.fc = nn.Linear(256, 1)  # 假设我们正在预测一个连续值

    def forward(self, x):
        cnn_output = self.cnn(x)
        lstm_output = self.lstm(cnn_output.unsqueeze(0))  # 增加一个序列长度维度
        attention_output = self.attention(lstm_output.unsqueeze(0))  # 增加一个序列长度维度
        output = self.fc(attention_output)
        return output

# 初始化模型、损失函数和优化器
model = HybridNeuralNetwork()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs.squeeze(), targets)
        loss.backward()
        optimizer.step()

        if (i + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}')

# 保存模型
torch.save(model.state_dict(), 'tennis_match_model.pth')