# features= input data/ predictors = LSTAT/RM/DIS
# target = target variable/ dependent variable = MEDV

# convert into tensors
inputs = torch.tensor(X_train.values, dtype=torch.float32)
labels = torch.tensor(y_train.values, dtype=torch.float32).reshape(-1, 1)

# Division of the data set into training and test data
train_size = int(0.8 * len(inputs))
test_size = len(inputs) - train_size


# The NN_ReLU class inherits properties from nn.Module
class Net(nn.Module):
    def __init__(self, input_size, output_size):
        #super()-methode executes the heritage
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_size, 300)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(300, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x
    

# Hyperparameter
input_size = 3 #Number of features/dependent variables, LSTAT/RM/DIS
output_size = 1 #one neuron for the result MEDV
learning_rate = 0.0001
num_epochs = 100

#Initializing the neural network
model = Net(input_size, output_size)

# loss function and optimization
lossFunc = nn.MSELoss()  # Mean Squared Error 
optimizer = optim.Adam(model.parameters(), lr= learning_rate)

# Training of the neural network
for epoch in range(num_epochs):
    for idx, element in enumerate(inputs):

        # forward propagation
        outputs = model(element)
        loss = lossFunc(outputs, labels[idx])

        # Backpropagation and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if epoch % 10 ==0:
            loss_value = loss.item()
            current = idx * len(inputs)
            print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {loss_value}, Batch: {current}/{len(inputs)}")

# Evaluation of performance
predicted_labels = [output.item() for output in model(inputs)]
mse = mean_squared_error(labels, predicted_labels)
print(f"Mean Squared Error: {mse}")

r2 = r2_score(labels, predicted_labels)
print("R²:", r2)