from .mnist_LeNet import MNIST_LeNet, MNIST_LeNet_Autoencoder
from .cifar10_LeNet import CIFAR10_LeNet, CIFAR10_LeNet_Autoencoder
from .cifar10_LeNet_elu import CIFAR10_LeNet_ELU, CIFAR10_LeNet_ELU_Autoencoder
from .tiny_imagenet_LeNet import TinyImagenet_LeNet, TinyImagenet_LeNet_Autoencoder
from .gtsrb_LeNet import GTSRB_LeNet, GTSRB_LeNet_Autoencoder


def build_network(net_name):
    """Builds the neural network."""

    implemented_networks = ('mnist_LeNet', 'cifar10_LeNet', 'cifar10_LeNet_ELU', 'tiny_imagenet_LeNet', 'gtsrb_LeNet')
    assert net_name in implemented_networks

    net = None

    if net_name == 'mnist_LeNet':
        net = MNIST_LeNet()

    if net_name == 'cifar10_LeNet':
        net = CIFAR10_LeNet()

    if net_name == 'cifar10_LeNet_ELU':
        net = CIFAR10_LeNet_ELU()

    if net_name == 'tiny_imagenet_LeNet':
        net = TinyImagenet_LeNet()

    if net_name == 'gtsrb_LeNet':
        net = GTSRB_LeNet()

    return net


def build_autoencoder(net_name):
    """Builds the corresponding autoencoder network."""

    implemented_networks = ('mnist_LeNet', 'cifar10_LeNet', 'cifar10_LeNet_ELU', 'tiny_imagenet_LeNet', 'gtsrb_LeNet')
    assert net_name in implemented_networks

    ae_net = None

    if net_name == 'mnist_LeNet':
        ae_net = MNIST_LeNet_Autoencoder()

    if net_name == 'cifar10_LeNet':
        ae_net = CIFAR10_LeNet_Autoencoder()

    if net_name == 'cifar10_LeNet_ELU':
        ae_net = CIFAR10_LeNet_ELU_Autoencoder()

    if net_name == 'tiny_imagenet_LeNet':
        ae_net = TinyImagenet_LeNet_Autoencoder()

    if net_name == 'gtsrb_LeNet':
        ae_net = GTSRB_LeNet_Autoencoder()

    return ae_net
