import torch


class FFNN(torch.nn.Module):
    def __init__(self, configuration: dict):
        super().__init__()
        self.hidden_size = configuration["hidden_size"]
        self.input_size = configuration["input_size"]
        self.output_size = 2
        self.hidden_1 = torch.nn.Linear(self.input_size, self.hidden_size)
        self.output_layer = torch.nn.Linear(self.hidden_size, 2)


    def forward(self, x: torch.Tensor)-> torch.Tensor:
        x = self.hidden_1(x)
        x = torch.relu(x)
        y = self.output_layer(x)
        return y

    def predict(self, x: torch.Tensor) -> torch.Tensor:
        outputs =  self.forward(x)
        outputs = outputs.softmax(dim=1)
        return outputs