# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
# pylint: disable=unidiomatic-typecheck
"""Prototype decorator for defining legacy-graph-mode functions."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import weakref

from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.core.protobuf import struct_pb2
from tensorflow.python.eager import context
from tensorflow.python.eager import function
from tensorflow.python.eager import lift_to_graph
from tensorflow.python.framework import composite_tensor
from tensorflow.python.framework import func_graph
from tensorflow.python.framework import importer
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.training.tracking import data_structures
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import tf_export


class VariableHolder(object):
  """Holds variables for a python function."""

  def __init__(self, fn=None, share_variables=False):
    self._fn = fn

    self._share_variables = share_variables
    self._variables_by_name = data_structures.Mapping()

  @property
  def variables(self):
    return self._variables_by_name

  def variable_creator_scope(self, next_creator, **kwargs):
    """Creates variables & adds them to collections to match legacy code."""
    collections = kwargs.pop("collections", None)
    v = None

    # Get expected variable name.
    with ops.name_scope(
        kwargs.get("name", None), "Variable", skip_on_eager=False) as name:
      variable_name = ops.name_from_scope_name(name)
      kwargs["name"] = name

    if self._share_variables:
      v = self._variables_by_name.get(variable_name, None)

    if v is None:
      v = next_creator(**kwargs)
      self._variables_by_name[variable_name] = v

    if collections is None:
      collections = [ops.GraphKeys.GLOBAL_VARIABLES]
    if v.trainable and ops.GraphKeys.TRAINABLE_VARIABLES not in collections:
      collections = list(collections) + [ops.GraphKeys.TRAINABLE_VARIABLES]

    ops.add_to_collections(collections, v)

    return v

  def __call__(self, *args, **kwargs):
    return self.call_with_variable_creator_scope(self._fn)(*args, **kwargs)

  def call_with_variable_creator_scope(self, fn):

    def wrapped(*args, **kwargs):
      with variable_scope.variable_creator_scope(self.variable_creator_scope):
        return fn(*args, **kwargs)

    return wrapped


def _get_element_from_tensor_info(tensor_info, graph):
  """Simplified copy of the deprecated `get_tensor_from_tensor_info`."""
  encoding = tensor_info.WhichOneof("encoding")
  if encoding == "name":
    # We may get operations here in some cases. TensorInfo is a bit of a
    # misnomer if so.
    return graph.as_graph_element(tensor_info.name)
  elif encoding == "coo_sparse":
    return sparse_tensor.SparseTensor(
        graph.get_tensor_by_name(tensor_info.coo_sparse.indices_tensor_name),
        graph.get_tensor_by_name(tensor_info.coo_sparse.values_tensor_name),
        graph.get_tensor_by_name(
            tensor_info.coo_sparse.dense_shape_tensor_name))
  elif encoding == "composite_tensor":
    struct_coder = nested_structure_coder.StructureCoder()
    spec_proto = struct_pb2.StructuredValue(
        type_spec_value=tensor_info.composite_tensor.type_spec)
    spec = struct_coder.decode_proto(spec_proto)
    components = [graph.get_tensor_by_name(component.name) for component in
                  tensor_info.composite_tensor.components]
    return spec._from_components(components)  # pylint: disable=protected-access
  else:
    raise ValueError(f"Invalid TensorInfo.encoding: {encoding}. Valid "
                     "encodings are 'name', 'coo_sparse', and "
                     "'composite_tensor'.")


def _lift_single_variable(old_variable, graph, variable_holder):
  """Lifts `old_variable` out of the `FuncGraph` `graph`."""
  new_variable = resource_variable_ops.UninitializedVariable(
      shape=old_variable.shape,
      dtype=old_variable.dtype,
      name=old_variable.op.name,
      trainable=old_variable.trainable,
      extra_handle_data=old_variable.handle)
  new_variable._initializer_op = old_variable._initializer_op  # pylint: disable=protected-access
  graph.add_capture(new_variable.handle, old_variable.handle)
  # Now that we've added the new variable to graph.captures,
  # graph.capture will use that cached value and do some post-processing
  # on the capture like recording it on the tape.
  graph.capture(new_variable.handle)
  # pylint: disable=protected-access
  variable_name = new_variable.name.split(":")[0]
  variable_holder._variables_by_name[variable_name] = new_variable
  graph._weak_variables.append(weakref.ref(new_variable))
  # pylint: enable=protected-access
  graph.watch_variable(new_variable)
  return new_variable