from torch import nn

from .utils import get_act
from .norm import Norm1D, Norm2D


class Dense(nn.Module):

    def __init__(self, input_dim, output_dim, bias=True, norm1d='none', norm2d='none', act='none'):
        super(Dense, self).__init__()
        assert norm1d == 'none' or norm2d == 'none', "one of [norm1d, norm2d] must be none"

        if norm1d != 'none':
            self.nn_norm = Norm1D(input_dim, norm1d)
        elif norm2d != 'none':
            self.nn_norm = Norm2D(input_dim, norm2d)
        else:
            self.nn_norm = None

        self.nn_act = get_act(act)
        self.nn_linear = nn.Linear(input_dim, output_dim, bias)

    def weight(self):
        return self.nn_linear.weight

    def forward(self, x):
        if self.nn_norm is not None:
            x = self.nn_norm(x)
        x = self.nn_act(x)
        x = self.nn_linear(x)
        return x