# Path: cortex/built_ins/networks/utils.py
# def apply_nonlinearity(x, nonlinearity, **nonlinearity_args):
#     if nonlinearity:
#         if isinstance(nonlinearity, str):
#             nonlinearity = get_nonlinearity(nonlinearity)
#         if callable(nonlinearity):
#             if isinstance(nonlinearity, nn.PReLU):
#                 nonlinearity.to(x.device)
#             try:
#                 x = nonlinearity(x, **nonlinearity_args)
#             except RuntimeError:
#                 nonlinearity.to('cpu')
#                 x = nonlinearity(x, **nonlinearity_args)
#         else:
#             raise ValueError(nonlinearity, type(nonlinearity))
#     return x
# 
# 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
# 
# 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)


import torch.nn as nn
import torch
from .utils import apply_nonlinearity, get_nonlinearity, finish_layer_1d

class BaseNet(nn.Module):
    '''Basic convenience network for cortex.

    Attributes:
        models: A sequence of

    '''

    def __init__(self, nonlinearity='ReLU', output_nonlinearity=None):
        super(BaseNet, self).__init__()

        self.models = nn.Sequential()

        self.output_nonlinearity = output_nonlinearity
        self.layer_nonlinearity = get_nonlinearity(nonlinearity)

    def forward(self,
                x: torch.Tensor,
                nonlinearity: str = None,
                **nonlinearity_args: dict) -> torch.Tensor:
        self.states = []
        if nonlinearity is None:
            nonlinearity = self.output_nonlinearity
        elif not nonlinearity:
            nonlinearity = None

        for model in self.models:
            x = model(x)
            self.states.append(x)

# Path: cortex/built_ins/networks/utils.py
# def apply_nonlinearity(x, nonlinearity, **nonlinearity_args):
#     if nonlinearity:
#         if isinstance(nonlinearity, str):
#             nonlinearity = get_nonlinearity(nonlinearity)
#         if callable(nonlinearity):
#             if isinstance(nonlinearity, nn.PReLU):
#                 nonlinearity.to(x.device)
#             try:
#                 x = nonlinearity(x, **nonlinearity_args)
#             except RuntimeError:
#                 nonlinearity.to('cpu')
#                 x = nonlinearity(x, **nonlinearity_args)
#         else:
#             raise ValueError(nonlinearity, type(nonlinearity))
#     return x
# 
# 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
# 
# 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)


# Path: cortex/built_ins/networks/base_network.py
import torch.nn as nn
import torch
from .utils import apply_nonlinearity, get_nonlinearity, finish_layer_1d


class BaseNet(nn.Module):
    '''Basic convenience network for cortex.

    Attributes:
        models: A sequence of

    '''

    def __init__(self, nonlinearity='ReLU', output_nonlinearity=None):
        super(BaseNet, self).__init__()

        self.models = nn.Sequential()

        self.output_nonlinearity = output_nonlinearity
        self.layer_nonlinearity = get_nonlinearity(nonlinearity)

    def forward(self,
                x: torch.Tensor,
                nonlinearity: str = None,
                **nonlinearity_args: dict) -> torch.Tensor:
        self.states = []
        if nonlinearity is None:
            nonlinearity = self.output_nonlinearity
        elif not nonlinearity:
            nonlinearity = None

        for model in self.models:
            x = model(x)
            self.states.append(x)

        x = apply_nonlinearity(x, nonlinearity, **nonlinearity_args)