import torch
import numpy as np
n_features = 1000
batch_size = 5

class NeuralNetworkModel(torch.nn.Module):

    def __init__(self):
        super(NeuralNetworkModel, self).__init__()
        self.fc1 = torch.nn.Linear(n_features, 500)
        self.fc2 = torch.nn.Linear(500, 1)

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

def get_loss_acc(model, X_dataset, Y_dataset, criterion, optimizer,):
    loss_score = 0
    acc_score = 0
    items_total = 0
    model.eval()
    for i in range(0, Y_dataset.shape[0], batch_size):
        X = X_dataset[i:i+batch_size]
        X = torch.tensor(X.astype(np.float32))
        Y = Y_dataset[i:i+batch_size]
        Y = torch.tensor(Y.astype(np.float32)).reshape(-1,1)
        Y_predictions = model(X)
        acc_score += torch.sum((Y_predictions > 0.5) == Y).item()
        items_total += Y.shape[0]
        optimizer.zero_grad()
        loss = criterion(Y_predictions, Y)
        loss.backward()
        optimizer.step()
        loss_score += loss.item() * Y.shape[0]
    return (loss_score / items_total), (acc_score / items_total)