# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/nasnet_space.py
# class SSF_NasnetA(mo.SearchSpaceFactory):
# 
#     def __init__(self):
#         mo.SearchSpaceFactory.__init__(
#             self, lambda: generate_search_space(5, 18, 2, 36, 3.0))
# 
# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/genetic_space.py
# class SSF_Genetic(mo.SearchSpaceFactory):
# 
#     def __init__(self):
#         mo.SearchSpaceFactory.__init__(
#             self, lambda: generate_search_space([3, 4, 5], [64, 128, 256],
#                                                 [5, 5, 5]))
# 
# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/nasbench_space.py
# class SSF_Nasbench(mo.SearchSpaceFactory):
# 
#     def __init__(self):
#         mo.SearchSpaceFactory.__init__(
#             self, lambda: generate_search_space(3, 3, 5, 128))
# 
# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/hierarchical_space.py
# def flat_search_space(num_classes):
#     c0 = 64
#     kernel_size = 3
#     motif_info = {2: {'num_nodes': 11, 'num_motifs': 1}}
# 
#     motif_hyperp = create_motif_hyperp(motif_info)
#     motif_fn = lambda c: create_motif(motif_hyperp[2][0], motif_hyperp,
#                                       motif_info, 2, c)
#     return mo.siso_sequential([
#         conv2d_cell(c0, kernel_size),
#         motif_fn(c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(c0),
#         separable_conv2d(2 * c0, kernel_size, stride=2),
#         motif_fn(2 * c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(2 * c0),
#         separable_conv2d(4 * c0, kernel_size, stride=2),
#         motif_fn(4 * c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(4 * c0),
#         separable_conv2d(8 * c0, kernel_size, stride=2),
#         global_average_pooling(),
#         dense(num_classes)
#     ])
# 
# def hierarchical_search_space(num_classes):
# 
#     c0 = 64  # initial number of channel C
#     kernel_size = 3
#     motif_info = {
#         2: {
#             'num_nodes': 4,
#             'num_motifs': 6
#         },
#         3: {
#             'num_nodes': 5,
#             'num_motifs': 1
#         }
#     }
#     motif_hyperp = create_motif_hyperp(motif_info)
#     motif_fn = lambda c: create_motif(motif_hyperp[3][0], motif_hyperp,
#                                       motif_info, 3, c)
#     return mo.siso_sequential([
#         conv2d_cell(c0, kernel_size),
#         motif_fn(c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(c0),
#         separable_conv2d(2 * c0, kernel_size, stride=2),
#         motif_fn(2 * c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(2 * c0),
#         separable_conv2d(4 * c0, kernel_size, stride=2),
#         motif_fn(4 * c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(4 * c0),
#         separable_conv2d(8 * c0, kernel_size, stride=2),
#         global_average_pooling(),
#         dense(num_classes)
#     ])
# 
# Path: deep_architect/contrib/misc/search_spaces/tensorflow/dnn.py
# def constant_initializer(c):
#     def init_fn(shape):
# def truncated_normal_initializer(stddev):
#     def init_fn(shape):
# def xavier_initializer_affine(gain=1.0):
#     def init_fn(shape):
# def relu():
# def affine(h_num_hidden, h_W_init_fn, h_b_init_fn):
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def dropout(h_keep_prob):
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def batch_normalization():
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def affine_simplified(h_num_hidden):
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def nonlinearity(h_nonlin_name):
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def dnn_cell(h_num_hidden, h_nonlin_name, h_swap, h_opt_drop, h_opt_bn,
#              h_drop_keep_prob):
# def dnn_net(num_classes):
#         W = tf.Variable(dh['W_init_fn']([n, m]))


from deep_architect.contrib.misc.search_spaces.tensorflow_eager.nasnet_space import SSF_NasnetA
from deep_architect.contrib.misc.search_spaces.tensorflow_eager.genetic_space import SSF_Genetic
from deep_architect.contrib.misc.search_spaces.tensorflow_eager.nasbench_space import SSF_Nasbench
from deep_architect.contrib.misc.search_spaces.tensorflow_eager.hierarchical_space import flat_search_space, hierarchical_search_space
from deep_architect.contrib.misc.search_spaces.tensorflow import dnn
import deep_architect.modules as mo

name_to_search_space_factory_fn = {
    'dnn':
    lambda num_classes: mo.SearchSpaceFactory(lambda: dnn.dnn_net(num_classes)),
    'nasnet':
    lambda num_classes: SSF_NasnetA(),
    'hierarchical':
    lambda num_classes: mo.SearchSpaceFactory(lambda: hierarchical_search_space(
        num_classes)),

# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/nasnet_space.py
# class SSF_NasnetA(mo.SearchSpaceFactory):
# 
#     def __init__(self):
#         mo.SearchSpaceFactory.__init__(
#             self, lambda: generate_search_space(5, 18, 2, 36, 3.0))
# 
# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/genetic_space.py
# class SSF_Genetic(mo.SearchSpaceFactory):
# 
#     def __init__(self):
#         mo.SearchSpaceFactory.__init__(
#             self, lambda: generate_search_space([3, 4, 5], [64, 128, 256],
#                                                 [5, 5, 5]))
# 
# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/nasbench_space.py
# class SSF_Nasbench(mo.SearchSpaceFactory):
# 
#     def __init__(self):
#         mo.SearchSpaceFactory.__init__(
#             self, lambda: generate_search_space(3, 3, 5, 128))
# 
# Path: deep_architect/contrib/misc/search_spaces/tensorflow_eager/hierarchical_space.py
# def flat_search_space(num_classes):
#     c0 = 64
#     kernel_size = 3
#     motif_info = {2: {'num_nodes': 11, 'num_motifs': 1}}
# 
#     motif_hyperp = create_motif_hyperp(motif_info)
#     motif_fn = lambda c: create_motif(motif_hyperp[2][0], motif_hyperp,
#                                       motif_info, 2, c)
#     return mo.siso_sequential([
#         conv2d_cell(c0, kernel_size),
#         motif_fn(c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(c0),
#         separable_conv2d(2 * c0, kernel_size, stride=2),
#         motif_fn(2 * c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(2 * c0),
#         separable_conv2d(4 * c0, kernel_size, stride=2),
#         motif_fn(4 * c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(4 * c0),
#         separable_conv2d(8 * c0, kernel_size, stride=2),
#         global_average_pooling(),
#         dense(num_classes)
#     ])
# 
# def hierarchical_search_space(num_classes):
# 
#     c0 = 64  # initial number of channel C
#     kernel_size = 3
#     motif_info = {
#         2: {
#             'num_nodes': 4,
#             'num_motifs': 6
#         },
#         3: {
#             'num_nodes': 5,
#             'num_motifs': 1
#         }
#     }
#     motif_hyperp = create_motif_hyperp(motif_info)
#     motif_fn = lambda c: create_motif(motif_hyperp[3][0], motif_hyperp,
#                                       motif_info, 3, c)
#     return mo.siso_sequential([
#         conv2d_cell(c0, kernel_size),
#         motif_fn(c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(c0),
#         separable_conv2d(2 * c0, kernel_size, stride=2),
#         motif_fn(2 * c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(2 * c0),
#         separable_conv2d(4 * c0, kernel_size, stride=2),
#         motif_fn(4 * c0),
#         separable_conv2d(c0, kernel_size),
#         motif_fn(4 * c0),
#         separable_conv2d(8 * c0, kernel_size, stride=2),
#         global_average_pooling(),
#         dense(num_classes)
#     ])
# 
# Path: deep_architect/contrib/misc/search_spaces/tensorflow/dnn.py
# def constant_initializer(c):
#     def init_fn(shape):
# def truncated_normal_initializer(stddev):
#     def init_fn(shape):
# def xavier_initializer_affine(gain=1.0):
#     def init_fn(shape):
# def relu():
# def affine(h_num_hidden, h_W_init_fn, h_b_init_fn):
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def dropout(h_keep_prob):
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def batch_normalization():
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def affine_simplified(h_num_hidden):
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def nonlinearity(h_nonlin_name):
#     def compile_fn(di, dh):
#         def forward_fn(di):
# def dnn_cell(h_num_hidden, h_nonlin_name, h_swap, h_opt_drop, h_opt_bn,
#              h_drop_keep_prob):
# def dnn_net(num_classes):
#         W = tf.Variable(dh['W_init_fn']([n, m]))


# Path: examples/contrib/kubernetes/search_space_factory.py
from deep_architect.contrib.misc.search_spaces.tensorflow_eager.nasnet_space import SSF_NasnetA
from deep_architect.contrib.misc.search_spaces.tensorflow_eager.genetic_space import SSF_Genetic
from deep_architect.contrib.misc.search_spaces.tensorflow_eager.nasbench_space import SSF_Nasbench
from deep_architect.contrib.misc.search_spaces.tensorflow_eager.hierarchical_space import flat_search_space, hierarchical_search_space
from deep_architect.contrib.misc.search_spaces.tensorflow import dnn
import deep_architect.modules as mo


name_to_search_space_factory_fn = {
    'dnn':
    lambda num_classes: mo.SearchSpaceFactory(lambda: dnn.dnn_net(num_classes)),
    'nasnet':
    lambda num_classes: SSF_NasnetA(),
    'hierarchical':
    lambda num_classes: mo.SearchSpaceFactory(lambda: hierarchical_search_space(
        num_classes)),

    'flat':