from categorical_embedder.embedders.core.aux.custom_layers import get_custom_layer_class
from categorical_embedder.embedders.core.aux.loss_factory import get_loss_function


def prepare_custom_objects(custom_object_info):
    custom_objects = {}
    custom_objects.update(_prepare_custom_layers(custom_object_info["layer_info"]))
    if not custom_object_info["has_implicit_loss"]:
        custom_objects.update(_prepare_custom_loss(custom_object_info["loss_info"]))
    return custom_objects


def _prepare_custom_layers(layer_info):
    custom_layers = {}
    for layer_name in layer_info:
        custom_layers[layer_name] = get_custom_layer_class(layer_name)
    return custom_layers


def _prepare_custom_loss(loss_info):
    return {"loss": get_loss_function(loss_info)}

[((16, 36, 16, 70), 'categorical_embedder.embedders.core.aux.custom_layers.get_custom_layer_class', 'get_custom_layer_class', ({(16, 59, 16, 69): 'layer_name'}, {}), '(layer_name)', False, 'from categorical_embedder.embedders.core.aux.custom_layers import get_custom_layer_class\n'), ((21, 20, 21, 48), 'categorical_embedder.embedders.core.aux.loss_factory.get_loss_function', 'get_loss_function', ({(21, 38, 21, 47): 'loss_info'}, {}), '(loss_info)', False, 'from categorical_embedder.embedders.core.aux.loss_factory import get_loss_function\n')]