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

# Load to Pandas Dataframe
train_path = 'training.dat'
test_path = 'testing.dat'

train_df = pd.read_csv(train_path, delim_whitespace=True, header=None)
test_df = pd.read_csv(test_path, delim_whitespace=True, header=None)

# Assuming that all columns are features and classes are based on row indices
X_train = torch.tensor(train_df.values).float()  # Convert feature data to tensor
X_test = torch.tensor(test_df.values).float()    # Convert feature data to tensor

# Create labels based on row indices
y_train = torch.tensor([(i // 25) for i in range(len(train_df))]).long()
y_test = torch.tensor([(i // 25) for i in range(len(test_df))]).long()

# Create data loaders
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=5, shuffle=True)
test_loader = DataLoader(TensorDataset(X_test, y_test), batch_size=5)

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.hidden1 = nn.Linear(4, 8)
        self.hidden2 = nn.Linear(8, 8)
        self.hidden3 = nn.Linear(8, 8)
        self.hidden4 = nn.Linear(8, 8)
        self.output = nn.Linear(8, 3)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.hidden1(x))
        x = self.relu(self.hidden2(x))
        x = self.relu(self.hidden3(x))
        x = self.relu(self.hidden4(x))
        x = self.output(x)
        return x

# Instantiate the model, loss function, and optimizer
model = NeuralNetwork()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
for epoch in range(50):  # number of epochs
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels) # 오차(오류) 계산
        loss.backward() # 오차(오류) 역전파
        optimizer.step() # 모델 가중치값 업데이트 -> Adam optimizer 기반

    print(f'Epoch {epoch+1}, Loss: {loss.item()}')

# Testing the model
correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in test_loader:
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {100 * correct / total}%')