from typing import Type, Optional

import torch
from torch import nn as nn


class SimpleMlp(nn.Module):
    """
    A class for very simple multi layer perceptron
    """
    def __init__(self, in_dim=2, out_dim=1, hidden_dim=64, n_layers=2,
                 activation: Type[nn.Module] = nn.ReLU, output_activation: Optional[Type[nn.Module]] = None):
        super(SimpleMlp, self).__init__()
        layers = [nn.Linear(in_dim, hidden_dim), activation()]
        layers.extend([nn.Linear(hidden_dim, hidden_dim), activation()] * (n_layers - 2))
        layers.append(nn.Linear(hidden_dim, out_dim))
        if output_activation:
            layers.append(output_activation())
        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)