from torch import Tensor, nn


def safe_log(x: Tensor) -> Tensor:
    """Compute natural logarithm in safe manner.

    :param root: x
    :returns: log(x)
    """
    return x.clamp(min=1e-10).log()


def make_nn(n_layers: int, d_input: int, d_hidden: int, d_output: int) -> nn.Module:
    """Creates a fully connected neural network with ReLU activations an batch normalization.

    :param n_layers: Number of layers.
    :param input_channels: Input dimension.
    :param hidden_channels: Hidden dimension.
    :param output_channels: Output dimension.
    :return: Neural network nn.Module.
    """
    layers = []
    layers.append(nn.Linear(d_input, d_hidden))
    layers.append(nn.BatchNorm1d(d_hidden))
    layers.append(nn.LeakyReLU())
    for _ in range(n_layers - 2):
        layers.append(nn.Linear(d_hidden, d_hidden))
        layers.append(nn.BatchNorm1d(d_hidden))
        layers.append(nn.LeakyReLU())
    layers.append(nn.Linear(d_hidden, d_output))
    return nn.Sequential(*layers)
