import torch
import torch.nn as nn
import numpy as np


class Model(nn.Module):
    """NN Model class"""

    def __init__(self, input_dim=300, hidden_dim=600, output_dim=1):
        """Initializes new instance of Model class"""

        super(Model, self).__init__()

        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim

        self.fc1 = nn.Linear(self.input_dim, self.hidden_dim)
        self.fc2 = nn.Linear(self.hidden_dim, self.hidden_dim)
        self.fc3 = nn.Linear(self.hidden_dim, self.output_dim)

        self.relu = nn.ReLU()

        self.criterion = nn.BCELoss()
        self.optimizer = torch.optim.SGD(self.parameters(), lr=0.01)

    def forward(self, x):
        """Step forward learning fn"""

        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        x = torch.sigmoid(x)
        return x

    def run_training(self, X_train, Y_train, batch_size, epochs_count):
        for _ in range(epochs_count):
            self.train()
            print(f"{Y_train.shape[0]}, {Y_train.shape[0] == self.input_dim}")
            print(f"{Y_train.shape[0]}, {Y_train.shape[0] == self.hidden_dim}")
            print(f"{Y_train.shape[0]}, {Y_train.shape[0] == self.output_dim}")
            for i in range(0, Y_train.shape[0], batch_size):
                X = X_train[i: i + batch_size]
                X = torch.tensor(X)
                y = Y_train[i: i + batch_size]
                y = torch.tensor(
                    y.astype(np.float32).to_numpy()).reshape(-1, 1)

                outputs = self(X.float())
                loss = self.criterion(outputs, y)

                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()