# 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.
# ==============================================================================

"""Utilities for V2 control flow."""

from tensorflow.core.framework import attr_value_pb2
from tensorflow.python.eager import context
from tensorflow.python.eager.polymorphic_function import atomic_function
from tensorflow.python.eager.polymorphic_function import concrete_function
from tensorflow.python.eager.polymorphic_function import tracing_compilation
from tensorflow.python.eager.polymorphic_function import transform
from tensorflow.python.framework import function_def_to_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework.func_graph import FuncGraph
from tensorflow.python.ops import control_flow_util
from tensorflow.python.ops import control_flow_v2_func_graphs
from tensorflow.python.ops import gradients_util
from tensorflow.python.util import keras_deps
from tensorflow.python.util import tf_contextlib
from tensorflow.python.util.tf_export import tf_export

_EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = None
_DISABLE_LOWER_USING_SWITCH_MERGE = False


CondBranchFuncGraph = control_flow_v2_func_graphs.CondBranchFuncGraph
WhileCondFuncGraph = control_flow_v2_func_graphs.WhileCondFuncGraph
WhileBodyFuncGraph = control_flow_v2_func_graphs.WhileBodyFuncGraph


def in_defun():
  """Returns if the current graph is, or is nested in, a defun."""
  if context.executing_eagerly(): return False

  graph = ops.get_default_graph()
  while (isinstance(graph, CondBranchFuncGraph) or
         isinstance(graph, WhileBodyFuncGraph) or
         isinstance(graph, WhileCondFuncGraph)):
    graph = graph.outer_graph
  return isinstance(graph, FuncGraph)


def in_while_loop_defun(graph):
  """Returns if the graph is a while loop FuncGraph."""
  if context.executing_eagerly(): return False
  return (isinstance(graph, WhileCondFuncGraph) or
          isinstance(graph, WhileBodyFuncGraph))


def create_new_tf_function(func_graph):
  """Converts func_graph to a TF_Function and adds it to the current graph.

  Args:
    func_graph: FuncGraph

  Returns:
    The name of the new TF_Function.
  """
  transform.apply_func_graph_transforms(func_graph)
  func = atomic_function.from_func_graph(func_graph.name, func_graph, {})

  func_graph.outer_graph._add_function_recursive(func)  # pylint: disable=protected-access
  return func_graph.name


def unique_fn_name(scope, name):
  """Returns a unique name to use for a control flow function.

  Args:
    scope: A name scope string.
    name: An identifier for this function (e.g. "true", "body").

  Returns:
    A string, the name to use for the function.
  """
  return ("%s%s_%s" % (scope, name, ops.uid())).replace("/", "_")


def unique_grad_fn_name(forward_name):
  return "%s_grad_%s" % (forward_name, ops.uid())


def maybe_set_lowering_attr(op, lower_using_switch_merge=None):
  """Sets the flag to enable lowering on `op` if necessary.

  Lowering allows cond_v2 and while_v2 to avoid some of the limitations of
  Functions, allowing users to specify devices & colocation inside of cond_v2
  and while_v2 input functions, and enabling non-strict evaluation & partial
  pruning. This brings v2 control flow closer to feature parity with v1 control
  flow.

  However, we do not lower in the following cases:
    - When the `If` or `While` ops are in the XLA context. Because it is easier
      for XLA to apply its own optimizations when dealing with un-lowered
      control flow operators than with low-level control flow primitives.
    - When the eager execution context specifies the executor of functions to
      be the single threaded executor (see context.function_executor_type()).
      Because the single threaded executor does not support v1 control flow ops.
    - When 'lower_using_switch_merge' is explicitly set to False.

  Args:
    op: An `If` or `While` Operation.
    lower_using_switch_merge: Explicit value to lower or not (optional).
  """
  if lower_using_switch_merge is not None:
    # pylint: disable=protected-access
    op._set_attr("_lower_using_switch_merge",
                 attr_value_pb2.AttrValue(b=lower_using_switch_merge))
    # pylint: enable=protected-access
  elif (not _DISABLE_LOWER_USING_SWITCH_MERGE and
        not control_flow_util.GraphOrParentsInXlaContext(op.graph) and
        context.context().function_call_options.executor_type !=
        "SINGLE_THREADED_EXECUTOR"):
    # pylint: disable=protected-access
    op._set_attr("_lower_using_switch_merge", attr_value_pb2.AttrValue(b=True))
    # pylint: enable=protected-access


def maybe_propagate_compile_time_consts_in_xla(op):
  """Tells XLA whether to propagate compile-time consts in the loop body.

  This is needed to make compile time constants available to ops, for example
  `max_num_elements` in `EmptyTensorList`, inside the loop body. Ideally this
  would always be turned on, but that doesn't work with legacy functionalized
  while_loops.

  Args:
    op: A `While` Operation.
  """
  if control_flow_util.GraphOrParentsInXlaContext(op.graph):
    # pylint: disable=protected-access
    op._set_attr("_xla_propagate_compile_time_consts",
                 attr_value_pb2.AttrValue(b=True))
    # pylint: enable=protected-access


def resource_input_index(tensor_name, input_names, node_defs, functions):
  """Returns the index of the input corresponding to `tensor_name`.

  This method is used to find the corresponding index of an arbitrary resource
  tensor in a function (the function could be a loop body). We assume that
  resource handles are never created in functions, so that every resource
  tensor can be traced back to a function input.

  The awkward signature of this method is to make it work with both FuncGraphs
  and FunctionDefs. This is so we can recurse on function call ops without
  building the corresponding FuncGraph (note that even if a FuncGraph for a
  FunctionDef already exists, the input/output/node names may have been
  changed when the FuncGraph was serialized to the FunctionDef, which makes it
  unusable with this algorithm).

  Args:
    tensor_name: the name of the resource tensor to be resolved to an input.
    input_names: a list of the names of all inputs to the function.
    node_defs: a dict mapping op name -> NodeDef for every op in the function.
    functions: a dict mapping function name -> AtomicFunction.

  Returns:
    The index into input_names corresponding to `tensor_name`.
  """
  while tensor_name not in input_names:
    # FunctionDefs and graphs use different tensor naming conventions.
    parts = tensor_name.split(":")
    if len(parts) == 3:
      op_name, _, output_idx = parts
    elif len(parts) == 2:
      op_name, output_idx = parts
    else:
      assert len(parts) == 1
      op_name = parts[0]
      output_idx = 0
      tensor_name = "%s:%d" % (tensor_name, output_idx)
      # Check again for cases where the tensor suffix (":0") is stripped out.
      if tensor_name in input_names:
        break
    output_idx = int(output_idx)
    node_def = node_defs[op_name]

    def _extract_input_index(function_attribute_name):
      func_name = node_def.attr[function_attribute_name].func.name
      fdef = functions[func_name].cached_definition
      output_arg_name = fdef.signature.output_arg[output_idx].name
      output_tensor_name = fdef.ret[output_arg_name]
      return resource_input_index(
          output_tensor_name, [arg.name for arg in fdef.signature.input_arg],
          {ndef.name: ndef for ndef in fdef.node_def}, functions)

    if node_def.op in ("Identity", "While"):
      # Captured resources occur at the same index in the lists of inputs and
      # outputs of a while or identity op. So we lookup the input of `tensor.op`
      # at the same index as the index of `tensor` in the `tensor.op.outputs`.
      tensor_name = node_def.input[output_idx]
    elif node_def.op in ("PartitionedCall", "StatefulPartitionedCall"):
      # Functions output any captured resource tensors used by their
      # gradients.  `tensor_name` is one of these outputs from a nested
      # function call, so recursively find the corresponding input in the
      # nested FunctionDef.
      tensor_name = node_def.input[_extract_input_index("f")]
    elif node_def.op in ("If", "StatelessIf"):
      input_index = _extract_input_index("then_branch")
      if input_index != _extract_input_index("else_branch"):
        raise AssertionError(
            ("Expected cond branches ({} op) to each have the same "
             "input->output mapping of resources.").format(node_def.op))
      tensor_name = node_def.input[
          # Ignore the `cond` input; the function inputs come after.
          input_index + 1]
    else:
      # We assume there are no other ops types that will "forward" resource
      # handles like this, so all other handles must have been created by the
      # op. (Note that cond_v2 wraps resource handle outputs in optionals,
      # which we'll end up accumulating).
      raise ValueError("Taking gradient of a while loop which creates "
                       "a resource in its body is not supported: %s (%s)"
                       % (op_name, node_def.op))

  return input_names.index(tensor_name)


@tf_contextlib.contextmanager
def clear_control_inputs():
  """Clears the control inputs but preserves the ControlFlowContext.

  This is needed to preserve the XLAControlFlowControl when clearing
  control inputs for the gradient accumulators in while_v2.
  `ops.control_dependencies` does not allow that.

  Yields:
    A context manager in which the ops created will not have any control inputs
    by default but the control flow context is the same.
  """
  # pylint: disable=protected-access
  control_flow_context = ops.get_default_graph()._get_control_flow_context()
  with ops.control_dependencies(None):
    ops.get_default_graph()._set_control_flow_context(control_flow_context)
    yield
  # pylint: enable=protected-access


def _is_tpu_strategy(strategy):
  return (strategy is not None and
          strategy.__class__.__name__.startswith("TPUStrategy"))


def _is_building_keras_layer():
  # TODO(srbs): Remove this function when we no long support session with Keras.
  keras_call_context_function = keras_deps.get_call_context_function()
  if keras_call_context_function:
    return keras_call_context_function().layer is not None
  else:
    return False


def output_all_intermediates():
  """Whether to output all intermediates of a functional control flow op.

  The default behavior is to output intermediates only when building a Keras
  Layer in graph mode and that too when certain other conditions are met:
  1. We do not output intermediates if the functional control flow op
     is being built inside a FuncGraph which is not a If/While graph. This
     guards against outputting intermediates in eager mode since keras adds
     tensors to a FuncGraph named "keras_graph" in that case. Also because we
     do not output intermediates of tf.function (since this feature is only for
     backwards compatibility) outputting intermediates of functional control
     flow ops built inside tf.function is of no value.
  2. We do not output intermediates when the compilation is using XLA or for a
     TPU.
  3. We do not output intermediates when a single threaded executor is used
     since that does not perform inlining and pruning.

  Returns:
    A bool telling whether to output all intermediates.
  """
  if _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE is not None:
    return _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE
  if in_defun():
    return False
  if control_flow_util.GraphOrParentsInXlaContext(ops.get_default_graph()):
    return False
  if (context.context().function_call_options.executor_type ==
      "SINGLE_THREADED_EXECUTOR"):
    return False
  return _is_building_keras_layer()


def get_func_graph(op, input_shapes, func_name):
  """Generates and returns a FuncGraph for the given op and input_shapes."""
  fdef = None
  graph = op.graph
  # Recursively search the func in graphs.
  while graph is not None:
    func = graph._get_function(func_name)  # pylint: disable=protected-access
    if func is not None:
      fdef = func.cached_definition
      break
    if hasattr(graph, "outer_graph"):
      graph = graph.outer_graph
    else:
      break

  if fdef is None:
    raise KeyError("%s cannot be found in the graph" % func_name)

  # `op.graph` may not be the same as `ops.get_default_graph()` e.g.
  # in the case of nested if ops or when the gradient is being computed
  # from inside a Defun. We build the `func_graph` with `op.graph` as its
  # `outer_graph`. This resembles how the `FuncGraph` was built in the
  # forward pass. We need this so that we can resolve references to tensors
  # in `func_graph` from its gradient graph in `_resolve_grad_inputs`.
  with op.graph.as_default():
    func_graph = function_def_to_graph.function_def_to_graph(
        fdef, input_shapes=input_shapes)

  # TODO(xjun): Ideally we want to retrieve the gradient functions instead of
  # re-create them. But the lifetime of gradient functions of PartitionedCall
  # ops is attached to ParitionedCall ops in the original func_graph and
  # when we are inside this function we don't have access to the original
  # func_graph or PartitionedCall ops. See cl/499362867 and cl/273858076 for
  # more context.
  for operation in func_graph.get_operations():
    if operation.type in ["PartitionedCall", "StatefulPartitionedCall"]:
      f = graph._get_function(operation.get_attr("f").name)  # pylint: disable=protected-access
      try:
        cf = concrete_function.ConcreteFunction.from_func_graph(
            f.graph,
            f.function_type,
            attrs=f.cached_definition.attr,
        )
      except AttributeError:
        # f is not found or f is a _DefinedFunction that doesn't have a graph.
        continue
      operation._gradient_function = cf._get_gradient_function()  # pylint: disable=protected-access

  return func_graph


def get_op_and_outputs(op_or_outputs):
  if isinstance(op_or_outputs, ops.Operation):
    return op_or_outputs, []
  elif not op_or_outputs:  # Empty list.
    return None, []
  else:
    return op_or_outputs[0].op, op_or_outputs


def graph_wrapped_for_higher_order_tape_gradients(graph):
  """Check if `graph` is wrapped by `run_as_function_for_tape_gradients`."""
  while graph is not None:
    if "cflow_gradient_wrapper" in getattr(graph, "name", ""):
      return True
    graph = getattr(graph, "outer_graph", None)
  return False


def run_as_function_for_tape_gradients(make_op, inputs):
  """Fix higher-order tape gradients by wrapping `make_op` in a function.

  Args:
    make_op: A function that takes a list of inputs and returns a list of output
      tensors. This function should set any handle data relevant to its outputs
      before returning.
    inputs: A list of tensors to check for tape gradients and pass to
      `make_op`. These should include all tensors used in `make_op`.

  Returns:
    Tensors corresponding to `make_op`'s output.
  """
  # GradientTapes created inside a function currently don't work well with
  # un-wrapped control flow ops in that same function. Wrapping in an extra
  # layer of intermediate function means we run extra logic in the function
  # gradient code to record the correct intermediates on the tape.
  #
  # The function attribute inputs to control flow ops are not hashable, so we
  # pass everything as a capture to bypass defun's caching.
  if (gradients_util.PossibleTapeGradientTypes(inputs)
      == gradients_util.POSSIBLE_GRADIENT_TYPES_HIGHER_ORDER
      # We only need one function between the tape and the op; if we've already
      # wrapped once, we stop wrapping to avoid infinite recursion.
      and not (ops.get_default_graph().building_function
               and "cflow_gradient_wrapper" in ops.get_default_graph().name)):
    results = tracing_compilation.call_function(
        (inputs,),
        tracing_options=tracing_compilation.TracingOptions(
            make_op, "cflow_gradient_wrapper", autograph=False
        ),
    )
    return results
  else:
    return make_op(inputs)


@tf_export(v1=["experimental.output_all_intermediates"])
def set_output_all_intermediates(state):  # pylint: disable=invalid-name
  """Whether to output all intermediates from functional control flow ops.

  The "default" behavior to is to output all intermediates when using v2 control
  flow inside Keras models in graph mode. This is needed to support taking
  gradients of v2 control flow. In graph mode, Keras can sometimes freeze the
  forward graph before the gradient computation which does not work for v2
  control flow since it requires updating the forward ops to output the needed
  intermediates. We work around this by proactively outputting the needed
  intermediates when building the forward pass itself. Ideally any such extra
  tensors should be pruned out at runtime. However, if for any reason this
  doesn't work for you or if you have an inference-only model you can turn this
  behavior off using
  `tf.compat.v1.experimental.output_all_intermediates(False)`.

  If with the default behavior you are still seeing errors of the form
  "Connecting to invalid output X of source node Y which has Z outputs" try
  setting `tf.compat.v1.experimental.output_all_intermediates(True)` and
  please file an issue at https://github.com/tensorflow/tensorflow/issues.

  Args:
    state: True, False or None. None restores the default behavior.
  """
  global _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE
  _EXPERIMENTAL_OUTPUT_ALL_INTERMEDIATES_OVERRIDE = state  # pylint: disable=protected-access