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

# Define MLP model
class MLP(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x):
        out = self.fc1(x)
        out = self.relu1(out)
        out = self.fc2(out)
        out = self.sigmoid(out)
        return out

# Define custom dataset
class PsychologyDataset(Dataset):
    def __init__(self, data_file):
        self.data = pd.read_excel(data_file)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        features = self.data.iloc[idx, 36:43].values.astype(np.float32)
        str = self.data.iloc[idx, -1]
        #print(idx,str,self.data.iloc[0, 0])
        label = -1
        if(str=="是"):
            label = 1
        else:
            label = 0
        #print(features)
        label = np.float32(label)
        #return torch.tensor(features, dtype=torch.float), label
        return features, label

# Set hyperparameters
input_size = 7
hidden_size = 16
output_size = 1
lr = 0.01
num_epochs = 100

# Load data
dataset = PsychologyDataset("data/data_src.xlsx")
dataloader = DataLoader(dataset, batch_size=1, shuffle=False)

# Instantiate model, loss function, and optimizer
model = MLP(input_size, hidden_size, output_size)
criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=lr)

# Train model
for epoch in range(num_epochs):
    running_loss = 0.0
    train_corrects = 0
    #print(type(dataloader))
    for i, data in enumerate(dataloader):
        #print("数据序号：", i, data)
        #continue
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels.unsqueeze(1))
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        _, preds = torch.max(outputs, 1)
        train_corrects += torch.sum(preds == labels.data)
    print('Epoch [%d/%d], Loss: %.4f' % (epoch+1, num_epochs, running_loss / len(dataloader)))
    train_acc = train_corrects.double() / len(dataloader)
    print('Epoch [%d/%d], ACC: %.4f' % (epoch+1, num_epochs, train_acc))

# Save trained model
torch.save(model.state_dict(), 'psychology_model.pth')