from torch import nn


class Net(nn.Module):
  def __init__(self, layer_sizes, activ):
    """
    In the constructor we instantiate nn modules and assign them as
    member variables.
    """
    super(Net, self).__init__()
    if activ == 'sigmoid':
        activation = nn.Sigmoid()
    elif activ == 'tanh':
        activation = nn.Tanh()
    else:
        activation = nn.ReLU()
    self.layer1 = nn.Sequential(nn.Linear(layer_sizes[0], layer_sizes[1]), activation)
    self.layer2 = nn.Sequential(nn.Linear(layer_sizes[1], layer_sizes[2]), activation)
    self.layer3 = nn.Sequential(nn.Linear(layer_sizes[2], layer_sizes[3]), activation)
    self.layer4 = nn.Linear(layer_sizes[3], layer_sizes[4], bias=False) # No biases in the final linear layer


  def forward(self, x):
    """
    In the forward function we accept a Variable of input data and we must return
    a Variable of output data. We can use Modules defined in the constructor as
    well as arbitrary operators on Variables.
    """
    output = self.layer1(x)
    output = output.sub(output.mean(0).expand(output.size())) # Removing mean from tensor
    output = self.layer2(output)
    output = output.sub(output.mean(0).expand(output.size())) # Removing mean from tensor
    output = self.layer3(output)
    output = output.sub(output.mean(0).expand(output.size())) # Removing mean from tensor
    output = self.layer4(output)
    output = output.sub(output.mean(0).expand(output.size())) # Removing mean from tensor
    return output