# Path: cortex/_lib/models.py
# MODEL_PLUGINS = {}
# 
# Path: cortex/built_ins/networks/fully_connected.py
# class FullyConnectedNet(BaseNet):
# 
#     def __init__(self, dim_in, dim_out=None, dim_h=64, dim_ex=None,
#                  nonlinearity='ReLU', n_levels=None, output_nonlinearity=None,
#                  normalize_input=False, **layer_args):
#         super(FullyConnectedNet, self).__init__(
#             nonlinearity=nonlinearity, output_nonlinearity=output_nonlinearity)
# 
#         dim_h = self.get_h(dim_h, n_levels=n_levels)
# 
#         if normalize_input:
#             self.models.add_module('initial_bn', nn.BatchNorm1d(dim_in))
# 
#         dim_in = self.add_linear_layers(dim_in, dim_h, dim_ex=dim_ex,
#                                         **layer_args)
#         self.add_output_layer(dim_in, dim_out)
# 
# Path: cortex/plugins.py
# class DatasetPlugin(DatasetPluginBase):
#         class IndexingDataset(C):
# class ModelPlugin(ModelPluginBase):
#     def copy_to_local_path(self, from_path: str) -> str:
#     def add_dataset(self, mode: str, dataset: Dataset):
#     def set_dataloader_class(self, dataloader_class: type):
#     def get_path(self, source: str):
#     def set_input_names(self, input_names):
#     def set_dims(self, **kwargs):
#     def set_scale(self, scale):
#     def make_indexing(self, C):
#             def __getitem__(self, index):
#     def build(self, *args, **kwargs):
#     def routine(self, *args, **kwargs):
#     def visualize(self, *args, **kwargs):
#     def train_step(self):
#     def eval_step(self):
#     def optimizer_step(self):
#     def train_loop(self):
#     def eval_loop(self):
#     def get_dims(self, *queries):
#     def add_noise(self, key, dist=None, size=None, **kwargs):
#     def add_image(self, *args, **kwargs):
#     def add_histogram(self, *args, **kwargs):
#     def add_scatter(self, *args, **kwargs):
# def register_plugin(plugin):


from cortex._lib.models import MODEL_PLUGINS
from cortex.built_ins.networks.fully_connected import FullyConnectedNet
from cortex.plugins import ModelPlugin, register_model


def test_class(model_class, arguments):
    '''Tests simple class attributions.

    Args:
        model_class: ModulePlugin subclass.
        arguments: Arguments for the class.

    '''
    arg1 = arguments['arg1']
    arg2 = arguments['arg2']
    arg1_help = arguments['arg1_help']
    arg2_help = arguments['arg2_help']

    assert model_class._help[arg1] == arg1_help, model_class.help[arg1]
    assert model_class._help[arg2] == arg2_help, model_class.help[arg2]
    assert model_class._kwargs[arg1] == 17, model_class.kwargs[arg1]
    assert model_class._kwargs[arg2] == 19, model_class.kwargs[arg2]


def test_register(model_class):
    '''Tests registration of a model.

    Args:
        model_class: ModelPlugin subclass.

    '''

    MODEL_PLUGINS.clear()

# Path: cortex/_lib/models.py
# MODEL_PLUGINS = {}
# 
# Path: cortex/built_ins/networks/fully_connected.py
# class FullyConnectedNet(BaseNet):
# 
#     def __init__(self, dim_in, dim_out=None, dim_h=64, dim_ex=None,
#                  nonlinearity='ReLU', n_levels=None, output_nonlinearity=None,
#                  normalize_input=False, **layer_args):
#         super(FullyConnectedNet, self).__init__(
#             nonlinearity=nonlinearity, output_nonlinearity=output_nonlinearity)
# 
#         dim_h = self.get_h(dim_h, n_levels=n_levels)
# 
#         if normalize_input:
#             self.models.add_module('initial_bn', nn.BatchNorm1d(dim_in))
# 
#         dim_in = self.add_linear_layers(dim_in, dim_h, dim_ex=dim_ex,
#                                         **layer_args)
#         self.add_output_layer(dim_in, dim_out)
# 
# Path: cortex/plugins.py
# class DatasetPlugin(DatasetPluginBase):
#         class IndexingDataset(C):
# class ModelPlugin(ModelPluginBase):
#     def copy_to_local_path(self, from_path: str) -> str:
#     def add_dataset(self, mode: str, dataset: Dataset):
#     def set_dataloader_class(self, dataloader_class: type):
#     def get_path(self, source: str):
#     def set_input_names(self, input_names):
#     def set_dims(self, **kwargs):
#     def set_scale(self, scale):
#     def make_indexing(self, C):
#             def __getitem__(self, index):
#     def build(self, *args, **kwargs):
#     def routine(self, *args, **kwargs):
#     def visualize(self, *args, **kwargs):
#     def train_step(self):
#     def eval_step(self):
#     def optimizer_step(self):
#     def train_loop(self):
#     def eval_loop(self):
#     def get_dims(self, *queries):
#     def add_noise(self, key, dist=None, size=None, **kwargs):
#     def add_image(self, *args, **kwargs):
#     def add_histogram(self, *args, **kwargs):
#     def add_scatter(self, *args, **kwargs):
# def register_plugin(plugin):


# Path: tests/built_ins/models/test_build.py
from cortex._lib.models import MODEL_PLUGINS
from cortex.built_ins.networks.fully_connected import FullyConnectedNet
from cortex.plugins import ModelPlugin, register_model



def test_class(model_class, arguments):
    '''Tests simple class attributions.

    Args:
        model_class: ModulePlugin subclass.
        arguments: Arguments for the class.

    '''
    arg1 = arguments['arg1']
    arg2 = arguments['arg2']
    arg1_help = arguments['arg1_help']
    arg2_help = arguments['arg2_help']

    assert model_class._help[arg1] == arg1_help, model_class.help[arg1]
    assert model_class._help[arg2] == arg2_help, model_class.help[arg2]
    assert model_class._kwargs[arg1] == 17, model_class.kwargs[arg1]
    assert model_class._kwargs[arg2] == 19, model_class.kwargs[arg2]


def test_register(model_class):
    '''Tests registration of a model.

    Args:
        model_class: ModelPlugin subclass.

    '''

    MODEL_PLUGINS.clear()

    register_model(model_class)