Wraps an arbitrary method so it does variable sharing.

  This decorator creates variables the first time it calls `method`, and reuses
  them for subsequent calls. The object that calls `method` provides a
  `tf.VariableScope`, either as a `variable_scope` attribute or as the return
  value of an `_enter_variable_scope()` method.

  The first time the wrapped method is invoked, it enters the caller's
  `tf.VariableScope` with `reuse=False`. On all subsequent calls it enters the
  same variable scope with `reuse=True`.

  Variables are created in the context of the `tf.VariableScope` provided by the
  caller object. Ops are created with an additional `tf.name_scope()`, which
  adds a scope for the wrapped method name. For example:

  ```python
  class MyClass(object):

    def __init__(self, name):
      with tf.variable_scope(None, default_name=name) as variable_scope:
        self.variable_scope = variable_scope

    @snt.reuse_variables
    def add_x(self, tensor):
      x = tf.get_variable("x", shape=tensor.get_shape())
      return tensor + x

  module = MyClass("my_module_name")
  input_tensor = tf.zeros(shape=(5,))

  # This creates the variable "my_module_name/x"
  # and op "my_module_name/add_x/add"
  output = module.add_x(input_tensor)
  ```

  For performance when executing eagerly it may be desirable to additionally
  annotate these methods using `defun`, such that they are encapsulated as
  graph functions. This is not recommended if your method returns a variable
  since the output of `defun` would be an op that returned the variable's value
  when evaluated (rather than the variable instance).

  ```python
  class FooModule(snt.AbstractModule):
    def _build(self, inputs):
      return complex_math(inputs)

    @tfe.defun
    @snt.reuse_variables
    def more_complex_stuff(self, inputs):
      return more_complex_math(inputs)
  ```

  Args:
    method: The method to wrap.

  Returns:
    The wrapped method.

def reuse_variables(method):
  """Wraps an arbitrary method so it does variable sharing.

  This decorator creates variables the first time it calls `method`, and reuses
  them for subsequent calls. The object that calls `method` provides a
  `tf.VariableScope`, either as a `variable_scope` attribute or as the return
  value of an `_enter_variable_scope()` method.

  The first time the wrapped method is invoked, it enters the caller's
  `tf.VariableScope` with `reuse=False`. On all subsequent calls it enters the
  same variable scope with `reuse=True`.

  Variables are created in the context of the `tf.VariableScope` provided by the
  caller object. Ops are created with an additional `tf.name_scope()`, which
  adds a scope for the wrapped method name. For example:

  ```python
  class MyClass(object):

    def __init__(self, name):
      with tf.variable_scope(None, default_name=name) as variable_scope:
        self.variable_scope = variable_scope

    @snt.reuse_variables
    def add_x(self, tensor):
      x = tf.get_variable("x", shape=tensor.get_shape())
      return tensor + x

  module = MyClass("my_module_name")
  input_tensor = tf.zeros(shape=(5,))

  # This creates the variable "my_module_name/x"
  # and op "my_module_name/add_x/add"
  output = module.add_x(input_tensor)
  ```

  For performance when executing eagerly it may be desirable to additionally
  annotate these methods using `defun`, such that they are encapsulated as
  graph functions. This is not recommended if your method returns a variable
  since the output of `defun` would be an op that returned the variable's value
  when evaluated (rather than the variable instance).

  ```python
  class FooModule(snt.AbstractModule):
    def _build(self, inputs):
      return complex_math(inputs)

    @tfe.defun
    @snt.reuse_variables
    def more_complex_stuff(self, inputs):
      return more_complex_math(inputs)
  ```

  Args:
    method: The method to wrap.

  Returns:
    The wrapped method.
  """
  initialized_variable_scopes_eager = set()
  initialized_variable_scopes_graph = weakref.WeakKeyDictionary()

  # Ensure that the argument passed in is really a method by checking that the
  # first positional argument to it is "self".
  arg_spec = inspect.getargspec(method)
  is_method = arg_spec.args and arg_spec.args[0] == "self"

  if not is_method:
    raise TypeError("reuse_variables can only be used with methods.")

  @wrapt.decorator
  def eager_test(method, obj, args, kwargs):
    """Validates runtime state in eager mode."""
    # If @reuse_variables is combined with @property, obj is passed in args
    # and method is still unbound at this stage.
    if obj is None:
      obj = args[0]

    if tf.executing_eagerly() and not hasattr(obj, "_template"):
      raise ValueError(
          "reuse_variables is not supported in eager mode except in Sonnet "
          "modules.")

    return method(*args, **kwargs)

  @wrapt.decorator
  def call_method(method, obj, args, kwargs):
    """Calls `method` with a variable scope whose reuse flag is set correctly.

    The first time the wrapper is called it creates a
    `(tf.Graph, tf.VariableScope)` key and checks it for membership in
    `initialized_variable_scopes`. The check is `False` if and only if this is
    the first time the wrapper has been called with the key, otherwise the
    check is `True`. The result of this check is used as the `reuse` flag for
    entering the provided variable scope before calling `method`.

    Here are two examples of how to use the reuse_variables decorator.

    1. Decorate an arbitrary instance method with a `variable_scope` attribute:

      ```python
      class Reusable(object):

        def __init__(self, name):
          with tf.variable_scope(None, default_name=name) as vs:
            self.variable_scope = vs

        @snt.reuse_variables
        def add_a(self, input_tensor):
          a = tf.get_variable("a", shape=input_tensor.get_shape())
          return a + input_tensor

      obj = Reusable("reusable")
      x = tf.constant(5.0)
      out1 = obj.add_a(x)
      out2 = obj.add_a(x)
      # out1 == out2
      ```

    2. Decorating a snt.AbstractModule instance method:

      ```python
      class ReusableModule(snt.AbstractModule):

        @snt.reuse_variables
        def add_a(self, input_tensor):
          a = tf.get_variable("a", shape=input_tensor.get_shape())
          return a + input_tensor

        # We don't need @snt.reuse_variables here because build is
        wrapped by # `tf.make_template` inside `snt.AbstractModule`.
        def _build(self, input_tensor):
          b = tf.get_variable("b", shape=input_tensor.get_shape())
          return b + self.add_a(input_tensor)

      obj = Reusable("reusable")
      x = tf.constant(5.0)
      out1 = obj(x)
      out2 = obj(x)
      # out1 == out2
      ```

    Args:
      method: The method to wrap.
      obj: The object instance passed to the wrapped method.
      args: The positional arguments (Tensors) passed to the wrapped method.
      kwargs: The keyword arguments passed to the wrapped method.

    Returns:
      Output of the wrapped method.

    Raises:
      ValueError: If no variable scope is provided or if `method` is a method
                  and a variable_scope keyword argument is also provided.
    """

    # If @reuse_variables is combined with @property, obj is passed in args
    # and method is still unbound at this stage.
    if obj is None:
      obj = args[0]

    def default_context_manager(reuse=None):
      variable_scope = obj.variable_scope
      return tf.variable_scope(variable_scope, reuse=reuse)

    variable_scope_context_manager = getattr(obj, "_enter_variable_scope",
                                             default_context_manager)

    with tf.init_scope():
      # We need `init_scope` incase we're running inside a defun. In that case
      # what we want is information about where the function will be called not
      # where the function is being built.
      graph = tf.get_default_graph()
      will_call_in_eager_context = tf.executing_eagerly()

    if will_call_in_eager_context:
      initialized_variable_scopes = initialized_variable_scopes_eager
    else:
      if graph not in initialized_variable_scopes_graph:
        initialized_variable_scopes_graph[graph] = set()
      initialized_variable_scopes = initialized_variable_scopes_graph[graph]

    # Temporarily enter the variable scope to capture it
    with variable_scope_context_manager() as tmp_variable_scope:
      variable_scope = tmp_variable_scope

    reuse = variable_scope.name in initialized_variable_scopes

    # Enter the pure variable scope with reuse correctly set
    with variable_scope_ops._pure_variable_scope(  # pylint:disable=protected-access
        variable_scope, reuse=reuse) as pure_variable_scope:
      current_name_scope = tf.get_default_graph().get_name_scope()
      # Force tf.name_scope to treat current_name_scope as an "absolute" scope
      # so we can re-enter it.
      if current_name_scope and current_name_scope[-1] != "/":
        current_name_scope += "/"
      with tf.name_scope(current_name_scope):
        module_name = pure_variable_scope.name
        method_name = to_snake_case(method.__name__)
        method_name_scope = "{}/{}".format(module_name, method_name)
        with tf.name_scope(method_name_scope) as scope:
          if hasattr(obj, "_capture_variables"):
            with obj._capture_variables():  # pylint: disable=protected-access
              out_ops = method(*args, **kwargs)
          else:
            out_ops = method(*args, **kwargs)
      initialized_variable_scopes.add(pure_variable_scope.name)
      try:
        # If `obj` is a Sonnet module, let it know it's been connected
        # to the TF graph.
        obj._is_connected = True  # pylint: disable=protected-access
        if not tf.executing_eagerly():
          obj._add_connected_subgraph(  # pylint: disable=protected-access
              method, out_ops, scope, args, kwargs)
      except AttributeError:
        pass
    return out_ops

  return eager_test(call_method(method))