import torch

class NonLinearRegression(torch.nn.Module):
    def __init__(self, inputSize, output_size, hiddenSize=128, numLayers=2, dropout=0.0):
        super(NonLinearRegression, self).__init__()
        self.inputSize = inputSize
        self.output_size = output_size

        self.hiddenSize = hiddenSize
        self.numLayers = numLayers

        self.dropout = dropout

        # add linear layers with relu
        self.layers = torch.nn.ModuleList()
        self.layers.append(torch.nn.Linear(inputSize, hiddenSize))
        self.layers.append(torch.nn.Dropout(dropout))
        for _ in range(numLayers - 2):
            self.layers.append(torch.nn.Linear(hiddenSize, hiddenSize))
            self.layers.append(torch.nn.Dropout(dropout))
        self.layers.append(torch.nn.Linear(hiddenSize, output_size))

        self.relu = torch.nn.ReLU()

    def forward(self, x):
        x = torch.squeeze(x, -1)

        for layer in self.layers[:-1]:
            x = self.relu(layer(x))

        out = self.layers[-1](x)
        return out



