import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch.optim as optim

# 读取CSV文件
train_csv_path = 'mouse_data.csv'  # 替换为你的CSV文件路径
train_csv = pd.read_csv(train_csv_path, header=None)
test_csv_path = 'mouse_data_test.csv'  # 替换为你的CSV文件路径
test_csv = pd.read_csv(test_csv_path, header=None)

# 提取dx, dy和标签
dx_dy_train = train_csv.iloc[:, 0].apply(lambda x: list(map(int, x.split(','))))
dx_dy_labels_train = train_csv.apply(lambda row: [list(map(int, row[i].split(','))) for i in range(1,10)], axis=1)
dx_dy_test = test_csv.iloc[:, 0].apply(lambda x: list(map(int, x.split(','))))
dx_dy_labels_test = test_csv.apply(lambda row: [list(map(int, row[i].split(','))) for i in range(1,10)], axis=1)

# 转换为PyTorch Tensor
dx_dy_train_tensor = torch.Tensor(dx_dy_train.tolist())
dx_dy_train_tensor = dx_dy_train_tensor.unsqueeze(1)
labels_train_tensor = torch.Tensor(dx_dy_labels_train.tolist())
dx_dy_test_tensor = torch.Tensor(dx_dy_test.tolist())
dx_dy_test_tensor = dx_dy_test_tensor.unsqueeze(1)
labels_test_tensor = torch.Tensor(dx_dy_labels_test.tolist())

# 创建自定义Dataset
class CustomDataset(Dataset):
    def __init__(self, dx_dy, labels):
        self.dx_dy = dx_dy
        self.labels = labels

    def __len__(self):
        return len(self.dx_dy)

    def __getitem__(self, idx):
        return self.dx_dy[idx], self.labels[idx]

# 创建Dataset和DataLoader
train_dataset = CustomDataset(dx_dy_train_tensor, labels_train_tensor)
test_dataset = CustomDataset(dx_dy_test_tensor, labels_test_tensor)
batch_size = 64
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)

# 定义神经网络模型
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(2, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 18)

    def forward(self, input_data):
        x = torch.flatten(input_data, start_dim=1)
        x = self.fc1(x)
        x = nn.ReLU()(x)
        x = self.fc2(x)
        x = nn.ReLU()(x)
        x = self.fc3(x)
        return x.view(-1, 9, 2)

# 初始化模型、损失函数和优化器
model = SimpleNet()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.9)

# 训练模型
epochs = 1000
for epoch in range(epochs):
    for batch_dx_dy, batch_labels in train_dataloader:
        optimizer.zero_grad()
        output = model(batch_dx_dy)
        loss = criterion(output, batch_labels)
        print(loss)
        loss.backward()
        optimizer.step()
        scheduler.step()
print('test')
for idx, data in enumerate(test_dataloader):
        batch_dx_dy, batch_labels = data
        output = model(batch_dx_dy)
        loss = criterion(output, batch_labels)
        print(loss)
        if idx == len(test_dataloader)-1:
            print(batch_dx_dy[0])
            print(output[0])