# Path: cortex/built_ins/networks/utils.py
# def finish_layer_1d(models, name, dim_out,
#                     dropout=False, layer_norm=False, batch_norm=False,
#                     nonlinearity=None):
#     if layer_norm and batch_norm:
#         logger.warning('Ignoring layer_norm because batch_norm is True')
# 
#     if dropout:
#         models.add_module(name + '_do', nn.Dropout(p=dropout))
# 
#     if layer_norm:
#         models.add_module(name + '_ln', nn.LayerNorm(dim_out))
#     elif batch_norm:
#         models.add_module(name + '_bn', nn.BatchNorm1d(dim_out))
# 
#     if nonlinearity:
#         nonlinearity = get_nonlinearity(nonlinearity)
#         models.add_module(
#             '{}_{}'.format(name, nonlinearity.__class__.__name__),
#             nonlinearity)
# 
# def get_nonlinearity(nonlinearity=None):
#     if not nonlinearity:
#         pass
#     elif callable(nonlinearity):
#         if nonlinearity == nn.LeakyReLU:
#             nonlinearity = nonlinearity(0.02, inplace=True)
#     elif hasattr(nn, nonlinearity):
#         nonlinearity = getattr(nn, nonlinearity)
#         if nonlinearity == 'LeakyReLU':
#             nonlinearity = nonlinearity(0.02, inplace=True)
#         else:
#             nonlinearity = nonlinearity()
#     elif hasattr(nn.functional, nonlinearity):
#         nonlinearity = getattr(nn.functional, nonlinearity)
#     else:
#         raise ValueError(nonlinearity)
#     return nonlinearity


from torch import nn
from torchvision import models
from .utils import finish_layer_1d, get_nonlinearity


class AlexNet(models.AlexNet):
    def __init__(self, shape, dim_out=None, fully_connected_layers=None,
                 nonlinearity='ReLU', n_steps=None,
                 **layer_args):
        super(AlexNet, self).__init__()
        fully_connected_layers = fully_connected_layers or []
        self.fc = nn.Sequential()
        dim_out_ = (256 * ((shape[0] + 4 - 10) // 32) *
                    ((shape[1] + 4 - 10) // 32))
        nonlinearity = get_nonlinearity(nonlinearity)
        for dim_h in fully_connected_layers:
            dim_in = dim_out_
            dim_out_ = dim_h
            name = 'linear_%s_%s' % (dim_in, dim_out_)
            self.fc.add_module(name, nn.Linear(dim_in, dim_out_))

# Path: cortex/built_ins/networks/utils.py
# def finish_layer_1d(models, name, dim_out,
#                     dropout=False, layer_norm=False, batch_norm=False,
#                     nonlinearity=None):
#     if layer_norm and batch_norm:
#         logger.warning('Ignoring layer_norm because batch_norm is True')
# 
#     if dropout:
#         models.add_module(name + '_do', nn.Dropout(p=dropout))
# 
#     if layer_norm:
#         models.add_module(name + '_ln', nn.LayerNorm(dim_out))
#     elif batch_norm:
#         models.add_module(name + '_bn', nn.BatchNorm1d(dim_out))
# 
#     if nonlinearity:
#         nonlinearity = get_nonlinearity(nonlinearity)
#         models.add_module(
#             '{}_{}'.format(name, nonlinearity.__class__.__name__),
#             nonlinearity)
# 
# def get_nonlinearity(nonlinearity=None):
#     if not nonlinearity:
#         pass
#     elif callable(nonlinearity):
#         if nonlinearity == nn.LeakyReLU:
#             nonlinearity = nonlinearity(0.02, inplace=True)
#     elif hasattr(nn, nonlinearity):
#         nonlinearity = getattr(nn, nonlinearity)
#         if nonlinearity == 'LeakyReLU':
#             nonlinearity = nonlinearity(0.02, inplace=True)
#         else:
#             nonlinearity = nonlinearity()
#     elif hasattr(nn.functional, nonlinearity):
#         nonlinearity = getattr(nn.functional, nonlinearity)
#     else:
#         raise ValueError(nonlinearity)
#     return nonlinearity


# Path: cortex/built_ins/networks/tv_models_wrapper.py
from torch import nn
from torchvision import models
from .utils import finish_layer_1d, get_nonlinearity



class AlexNet(models.AlexNet):
    def __init__(self, shape, dim_out=None, fully_connected_layers=None,
                 nonlinearity='ReLU', n_steps=None,
                 **layer_args):
        super(AlexNet, self).__init__()
        fully_connected_layers = fully_connected_layers or []
        self.fc = nn.Sequential()
        dim_out_ = (256 * ((shape[0] + 4 - 10) // 32) *
                    ((shape[1] + 4 - 10) // 32))
        nonlinearity = get_nonlinearity(nonlinearity)
        for dim_h in fully_connected_layers:
            dim_in = dim_out_
            dim_out_ = dim_h
            name = 'linear_%s_%s' % (dim_in, dim_out_)
            self.fc.add_module(name, nn.Linear(dim_in, dim_out_))

            finish_layer_1d(self.fc, name, dim_out_,