import torch.nn as nn
import torch.nn.functional as F


class SimpleNet(nn.Module):
    def __init__(self, input_dim, params):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(input_dim, params["hidden_size1"])
        self.fc2 = nn.Linear(params["hidden_size1"], params["hidden_size2"])
        self.fc3 = nn.Linear(params["hidden_size2"], 1)
        self.drop1 = nn.Dropout(p=params["dropout_p1"])
        self.drop2 = nn.Dropout(p=params["dropout_p2"])
        self.bn1 = nn.BatchNorm1d(params["hidden_size1"])
        self.bn2 = nn.BatchNorm1d(params["hidden_size2"])

    def forward(self, x):
        x = self.fc1(x)
        x = self.drop1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = self.drop2(x)
        x = self.bn2(x)
        x = F.relu(x)
        x = self.fc3(x)
        x = F.sigmoid(x)
        return x
