# Copyright 2015 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.
# ==============================================================================
"""## Functions for copying elements from one graph to another.

These functions allow for recursive copying of elements (ops and variables)
from one graph to another. The copied elements are initialized inside a
user-specified scope in the other graph. There are separate functions to
copy ops and variables.
There is also a function to retrieve the copied version of an op from the
first graph inside a scope in the second graph.

@@copy_op_to_graph
@@copy_variable_to_graph
@@get_copied_op
"""

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

from copy import deepcopy
from tensorflow.python.ops.variables import Variable
from tensorflow.python.client.session import Session
from tensorflow.python.framework import ops

__all__ = ['copy_op_to_graph', 'copy_variable_to_graph', 'get_copied_op']


def copy_variable_to_graph(org_instance, to_graph, scope=''):
  """Given a `Variable` instance from one `Graph`, initializes and returns
  a copy of it from another `Graph`, under the specified scope
  (default `""`).

  Args:
    org_instance: A `Variable` from some `Graph`.
    to_graph: The `Graph` to copy the `Variable` to.
    scope: A scope for the new `Variable` (default `""`).

  Returns:
    The copied `Variable` from `to_graph`.

  Raises:
    TypeError: If `org_instance` is not a `Variable`.
  """

  if not isinstance(org_instance, Variable):
    raise TypeError(str(org_instance) + ' is not a Variable')

  #The name of the new variable
  if scope != '':
    new_name = (scope + '/' + org_instance.name[:org_instance.name.index(':')])
  else:
    new_name = org_instance.name[:org_instance.name.index(':')]

  #Get the collections that the new instance needs to be added to.
  #The new collections will also be a part of the given scope,
  #except the special ones required for variable initialization and
  #training.
  collections = []
  for name, collection in org_instance.graph._collections.items():
    if org_instance in collection:
      if (name == ops.GraphKeys.GLOBAL_VARIABLES or
          name == ops.GraphKeys.TRAINABLE_VARIABLES or scope == ''):
        collections.append(name)
      else:
        collections.append(scope + '/' + name)

  #See if it's trainable.
  trainable = (
      org_instance in org_instance.graph.get_collection(
          ops.GraphKeys.TRAINABLE_VARIABLES))
  #Get the initial value
  with org_instance.graph.as_default():
    temp_session = Session()
    init_value = temp_session.run(org_instance.initialized_value())

  #Initialize the new variable
  with to_graph.as_default():
    new_var = Variable(
        init_value,
        trainable,
        name=new_name,
        collections=collections,
        validate_shape=False)

  return new_var


def copy_op_to_graph(org_instance, to_graph, variables, scope=''):
  """Returns a copy of an operation from another Graph under a specified scope.

  Given an `Operation` `org_instance` from one `Graph`,
  initializes and returns a copy of it from another `Graph`,
  under the specified scope (default `""`).

  The copying is done recursively, so any `Operation` whose output
  is required to evaluate the `org_instance`, is also copied (unless
  already done).

  Since `Variable` instances are copied separately, those required
  to evaluate `org_instance` must be provided as input.

  Args:
    org_instance: An `Operation` from some `Graph`. Could be a
      `Placeholder` as well.
    to_graph: The `Graph` to copy `org_instance` to.
    variables: An iterable of `Variable` instances to copy `org_instance` to.
    scope: A scope for the new `Variable` (default `""`).

  Returns:
    The copied `Operation` from `to_graph`.

  Raises:
    TypeError: If `org_instance` is not an `Operation` or `Tensor`.
  """

  #The name of the new instance
  if scope != '':
    new_name = scope + '/' + org_instance.name
  else:
    new_name = org_instance.name

  #Extract names of variables
  copied_variables = dict((x.name, x) for x in variables)

  #If a variable by the new name already exists, return the
  #correspondng tensor that will act as an input
  if new_name in copied_variables:
    return to_graph.get_tensor_by_name(copied_variables[new_name].name)

  #If an instance of the same name exists, return appropriately
  try: