# Path: deep_architect/helpers/tensorflow_eager_support.py
# def siso_tensorflow_eager_module(name, compile_fn, name_to_hyperp, scope=None):
#     return TensorflowEagerModule(name, compile_fn, name_to_hyperp, ['in'],
#                                  ['out'], scope).get_io()


import tensorflow as tf
from deep_architect.helpers.tensorflow_eager_support import siso_tensorflow_eager_module
# mapping from a layer name to the shared layer used by the DeepArchitect modules.
class WeightSharer:

    def __init__(self):
        self.name_to_weight = {}

    def get(self, name, weight_fn):
        if name not in self.name_to_weight:
            self.name_to_weight[name] = weight_fn()
        return self.name_to_weight[name]


weight_sharer = WeightSharer()


# Now to share weights or any other object, simply use a common name for the object
# you wish to share, and pass in a function to initialize the object. The first
# time the function is called, the convolution layer will be created. Subsequent
# calls will simply retrieve the layer from the WeightSharer object.
def conv2D(filter_size, channels, name):

    def compile_fn(di, dh):
        conv_fn = lambda: tf.keras.layers.Conv2D(channels, filter_size)
        conv = weight_sharer.get(name, conv_fn)

        def fn(di, is_training=True):
            return {'out': conv(di['in'])}

        return fn


# Path: deep_architect/helpers/tensorflow_eager_support.py
# def siso_tensorflow_eager_module(name, compile_fn, name_to_hyperp, scope=None):
#     return TensorflowEagerModule(name, compile_fn, name_to_hyperp, ['in'],
#                                  ['out'], scope).get_io()


# Path: dev/tutorials/weight_sharing/weight_sharing/main.py
import tensorflow as tf
from deep_architect.helpers.tensorflow_eager_support import siso_tensorflow_eager_module

# mapping from a layer name to the shared layer used by the DeepArchitect modules.
class WeightSharer:

    def __init__(self):
        self.name_to_weight = {}

    def get(self, name, weight_fn):
        if name not in self.name_to_weight:
            self.name_to_weight[name] = weight_fn()
        return self.name_to_weight[name]


weight_sharer = WeightSharer()


# Now to share weights or any other object, simply use a common name for the object
# you wish to share, and pass in a function to initialize the object. The first
# time the function is called, the convolution layer will be created. Subsequent
# calls will simply retrieve the layer from the WeightSharer object.
def conv2D(filter_size, channels, name):

    def compile_fn(di, dh):
        conv_fn = lambda: tf.keras.layers.Conv2D(channels, filter_size)
        conv = weight_sharer.get(name, conv_fn)

        def fn(di, is_training=True):
            return {'out': conv(di['in'])}

        return fn


    return siso_tensorflow_eager_module('Conv2D', compile_fn, {})