# 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
#
#     https://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.
# ==============================================================================
"""SubGraphView: a subgraph view on an existing tf.Graph.
"""

from __future__ import annotations

import copy

import six
from six import iteritems
from six import StringIO

from . import select
from . import util
from tensorflow.python.framework import ops as tf_ops

__all__ = [
    "SubGraphView",
    "make_view",
    "make_view_from_scope",
]


def _finalize_index(index_or_t, ts):
    """Returns index as is or return index of tensor in `ts`."""
    if isinstance(index_or_t, six.integer_types):
        return index_or_t
    else:
        return ts.index(index_or_t)


def _finalize_indices(list_of_index_or_t, ts):
    """Returns index in `indices` as is or replace with tensor's index."""
    return [_finalize_index(index_or_t, ts) for index_or_t in list_of_index_or_t]


def _check_within_range(mapping, n, repetition):
    """Check is the mapping is valid.

    Args:
      mapping: an iterable of integer.
      n: define the input domain as [0, n-1]. Note that the mapping can be
        under-complete, that is, it can only contain a subset of the integers on
        [0, n-1].
      repetition: if True repetition are allowed (the function is surjective)
        otherwise repetition are not allowed (the function is injective).
    Raises:
      ValueError: if the mapping is out of range ot if repetition is False and
        the mapping has some repetition.
    """
    for i in mapping:
        if not 0 <= i < n:
            raise ValueError("Out of [0, {}[ range: {}".format(n, i))
    if not repetition and len(set(mapping)) != len(mapping):
        raise ValueError("Found repetition in mapping: {}".format(mapping))


class SubGraphView(object):
    """A subgraph view on an existing `tf.Graph`.

    An instance of this class is a subgraph view on an existing `tf.Graph`.
    "subgraph" means that it can represent part of the whole `tf.Graph`.
    "view" means that it only provides a passive observation and do not to act
    on the `tf.Graph`. Note that in this documentation, the term "subgraph" is
    often used as substitute to "subgraph view".

    A subgraph contains:

    * a list of input tensors, accessible via the `inputs` property.
    * a list of output tensors, accessible via the `outputs` property.
    * and the operations in between, accessible via the "ops" property.

    An subgraph can be seen as a function F(i0, i1, ...) -> o0, o1, ... It is a
    function which takes as input some input tensors and returns as output some
    output tensors. The computation that the function performs is encoded in the
    operations of the subgraph.

    The tensors (input or output) can be of two kinds:

    - connected: a connected tensor connects to at least one operation contained
    in the subgraph. One example is a subgraph representing a single operation
    and its inputs and outputs: all the input and output tensors of the op
    are "connected".
    - passthrough: a passthrough tensor does not connect to any operation
    contained in the subgraph. One example is a subgraph representing a
    single tensor: this tensor is passthrough. By default a passthrough tensor is
    present both in the input and output tensors of the subgraph. It can however
    be remapped to only appear as an input (or output) only.

    The input and output tensors can be remapped. For instance, some input tensor
    can be omitted. For instance, a subgraph representing an operation with two
    inputs can be remapped to only take one input. Note that this does not change
    at all the underlying `tf.Graph` (remember, it is a view). It means that
    the other input is being ignored, or is being treated as "given".
    The analogy with functions can be extended like this: F(x,y) is the original
    function. Remapping the inputs from [x, y] to just [x] means that the subgraph
    now represent the function F_y(x) (y is "given").

    The output tensors can also be remapped. For instance, some output tensor can
    be omitted. Other output tensor can be duplicated as well. As mentioned
    before, this does not change at all the underlying `tf.Graph`.
    The analogy with functions can be extended like this: F(...)->x,y is the
    original function. Remapping the outputs from [x, y] to just [y,y] means that
    the subgraph now represent the function M(F(...)) where M is the function
    M(a,b)->b,b.

    It is useful to describe three other kind of tensors:

    * internal: an internal tensor is a tensor connecting operations contained
      in the subgraph. One example in the subgraph representing the two
      operations A and B connected sequentially: -> A -> B ->. The middle arrow
      is an internal tensor.
    * actual input: an input tensor of the subgraph, regardless of whether it is
      listed in "inputs" or not (masked-out).
    * actual output: an output tensor of the subgraph, regardless of whether it is
      listed in "outputs" or not (masked-out).
    * hidden input: an actual input which has been masked-out using an
      input remapping. In other word, a hidden input is a non-internal tensor
      not listed as a input tensor and one of whose consumers belongs to
      the subgraph.
    * hidden output: a actual output which has been masked-out using an output
      remapping. In other word, a hidden output is a non-internal tensor
      not listed as an output and one of whose generating operations belongs to
      the subgraph.

    Here are some useful guarantees about an instance of a SubGraphView:

    * the input (or output) tensors are not internal.
    * the input (or output) tensors are either "connected" or "passthrough".
    * the passthrough tensors are not connected to any of the operation of
    the subgraph.

    Note that there is no guarantee that an operation in a subgraph contributes
    at all to its inputs or outputs. For instance, remapping both the inputs and
    outputs to empty lists will produce a subgraph which still contains all the
    original operations. However, the remove_unused_ops function can be used to
    make a new subgraph view whose operations are connected to at least one of
    the input or output tensors.

    An instance of this class is meant to be a lightweight object which is not
    modified in-place by the user. Rather, the user can create new modified
    instances of a given subgraph. In that sense, the class SubGraphView is meant
    to be used like an immutable python object.

    A common problem when using views is that they can get out-of-sync with the
    data they observe (in this case, a `tf.Graph`). This is up to the user to
    ensure that this doesn't happen. To keep on the safe side, it is recommended
    that the life time of subgraph views are kept very short. One way to achieve
    this is to use subgraphs within a "with make_sgv(...) as sgv:" Python context.

    To alleviate the out-of-sync problem, some functions are granted the right to
    modified subgraph in place. This is typically the case of graph manipulation
    functions which, given some subgraphs as arguments, can modify the underlying
    `tf.Graph`. Since this modification is likely to render the subgraph view
    invalid, those functions can modify the argument in place to reflect the
    change. For instance, calling the function swap_inputs(svg0, svg1) will modify
    svg0 and svg1 in place to reflect the fact that their inputs have now being
    swapped.
    """

    def __init__(self, inside_ops=(), passthrough_ts=()):
        """Create a subgraph containing the given ops and the "passthrough" tensors.

        Args:
          inside_ops: an object convertible to a list of `tf.Operation`. This list
            defines all the operations in the subgraph.
          passthrough_ts: an object convertible to a list of `tf.Tensor`. This list
            define all the "passthrough" tensors. A passthrough tensor is a tensor
            which goes directly from the input of the subgraph to it output, without
            any intermediate operations. All the non passthrough tensors are
            silently ignored.
        Raises:
          TypeError: if inside_ops cannot be converted to a list of `tf.Operation`
            or if `passthrough_ts` cannot be converted to a list of `tf.Tensor`.
        """

        inside_ops = util.make_list_of_op(inside_ops)
        passthrough_ts = util.make_list_of_t(passthrough_ts)
        ops_and_ts = inside_ops + passthrough_ts
        if ops_and_ts:
            self._graph = util.get_unique_graph(ops_and_ts)
            self._ops = inside_ops

            # Compute inside and outside tensor
            inputs, outputs, insides = select.compute_boundary_ts(inside_ops)

            # Compute passthrough tensors, silently ignoring the non-passthrough ones.
            all_tensors = frozenset(inputs + outputs + list(insides))
            self._passthrough_ts = [t for t in passthrough_ts if t not in all_tensors]

            # Set inputs and outputs.
            self._input_ts = inputs + self._passthrough_ts
            self._output_ts = outputs + self._passthrough_ts
        else:
            self._graph = None
            self._passthrough_ts = []
            self._input_ts = []
            self._output_ts = []
            self._ops = []

    def __copy__(self):
        """Create a copy of this subgraph.

        Note that this class is a "view", copying it only create another view and
        does not copy the underlying part of the `tf.Graph`.

        Returns:
          A new identical instance of the original subgraph view.
        """
        cls = self.__class__
        result = cls.__new__(cls)
        for k, v in iteritems(self.__dict__):
            if k == "_graph":
                setattr(result, k, v)
            else:
                setattr(result, k, list(v))  # copy the list
        return result

    def _assign_from(self, other):
        """Assign other to itself.

        Args:
          other: another subgraph-view.
        Returns:
          A new instance identical to the original one.
        Raises:
          TypeError: if other is not an SubGraphView.
        """
        if not isinstance(other, SubGraphView):
            raise TypeError("Expected SubGraphView, got: {}".format(type(other)))
        # pylint: disable=protected-access
        self._graph = other._graph
        self._ops = list(other._ops)
        self._passthrough_ts = list(other._passthrough_ts)
        self._input_ts = list(other._input_ts)
        self._output_ts = list(other._output_ts)
        # pylint: enable=protected-access

    def copy(self):
        """Return a copy of itself.

        Note that this class is a "view", copying it only create another view and
        does not copy the underlying part of the tf.Graph.

        Returns:
          A new instance identical to the original one.
        """
        return copy.copy(self)

    def _remap_default(self, remove_input_map=True, remove_output_map=True):
        """Remap in the place the inputs and/or outputs to the default mapping.

        Args:
          remove_input_map: if True the input map is reset to the default one.
          remove_output_map: if True the output map is reset to the default one.
        """
        if not remove_input_map and not remove_output_map:
            return

        # Compute inside and outside tensor
        inputs, outputs, _ = select.compute_boundary_ts(self._ops)
        if remove_input_map:
            self._input_ts = list(inputs) + self._passthrough_ts
        if remove_output_map:
            self._output_ts = list(outputs) + self._passthrough_ts

    def remap_default(self, remove_input_map=True, remove_output_map=True):
        """Remap the inputs and/or outputs to the default mapping.

        Args:
          remove_input_map: if True the input map is reset to the default one.
          remove_output_map: if True the output map is reset to the default one.
        Returns:
          A new modified instance of the original subgraph view with its
            input and/or output mapping reset to the default one.
        """
        res = self.copy()
        res._remap_default(remove_input_map, remove_output_map)  # pylint: disable=protected-access
        return res

    def _remap_inputs(self, new_input_indices):
        """Remap the inputs of the subgraph in-place."""
        new_input_indices = _finalize_indices(new_input_indices, self._input_ts)
        _check_within_range(new_input_indices, len(self._input_ts), repetition=False)
        self._input_ts = [self._input_ts[i] for i in new_input_indices]

    def _remap_outputs(self, new_output_indices):
        """Remap the outputs of the subgraph in-place."""
        new_output_indices = _finalize_indices(new_output_indices, self._output_ts)
        _check_within_range(new_output_indices, len(self._output_ts), repetition=True)
        self._output_ts = [self._output_ts[i] for i in new_output_indices]

    def _remap_outputs_make_unique(self):
        """Remap the outputs in place so that all the tensors appears only once."""
        output_ts = list(self._output_ts)
        self._output_ts = []
        util.concatenate_unique(self._output_ts, output_ts)

    def _remap_outputs_to_consumers(self):
        """Remap the outputs in place to match the number of consumers."""
        self._remap_outputs_make_unique()
        output_ts = list(self._output_ts)
        self._output_ts = []
        for t in output_ts:
            self._output_ts += [t] * len(t.consumers())

    def remap_outputs_make_unique(self):
        """Remap the outputs so that all the tensors appears only once."""
        res = copy.copy(self)
        res._remap_outputs_make_unique()  # pylint: disable=protected-access
        return res

    def remap_outputs_to_consumers(self):
        """Remap the outputs to match the number of consumers."""
        res = copy.copy(self)
        res._remap_outputs_to_consumers()  # pylint: disable=protected-access
        return res

    def _remove_unused_ops(self, control_inputs=True):
        """Remove unused ops in place.

        Args:
          control_inputs: if True, control inputs are used to detect used ops.
        Returns:
          A new subgraph view which only contains used operations.
        """
        ops = select.get_walks_union_ops(
            self.connected_inputs, self.connected_outputs, within_ops=self._ops, control_inputs=control_inputs
        )
        self._ops = [op for op in self._ops if op in ops]

    def remove_unused_ops(self, control_inputs=True):
        """Remove unused ops.

        Args:
          control_inputs: if True, control inputs are used to detect used ops.
        Returns:
          A new subgraph view which only contains used operations.
        """
        res = copy.copy(self)
        res._remove_unused_ops(control_inputs)  # pylint: disable=protected-access
        return res

    def remap_inputs(self, new_input_indices):
        """Remap the inputs of the subgraph.

        If the inputs of the original subgraph are [t0, t1, t2], remapping to [2,0]
        will create a new instance whose inputs is [t2, t0].

        Note that this is only modifying the view: the underlying `tf.Graph` is not
        affected.

        Args:
          new_input_indices: an iterable of integers or tf.Tensors
            representing a mapping between the old inputs and the new ones.
            Integers must be positive and smaller than the number of old inputs.
            tf.Tensors must belong to the old list of inputs.
            This mapping can be under-complete and must be without repetitions.
        Returns:
          A new modified instance of the original subgraph view with remapped
            inputs.
        """
        res = self.copy()
        res._remap_inputs(new_input_indices)  # pylint: disable=protected-access
        return res

    def remap_outputs(self, new_output_indices):
        """Remap the output of the subgraph.

        If the output of the original subgraph are [t0, t1, t2], remapping to
        [1,1,0] will create a new instance whose outputs is [t1, t1, t0].

        Note that this is only modifying the view: the underlying tf.Graph is not
        affected.

        Args:
          new_output_indices: an iterable of integers or tf.Tensors
            representing a mapping between the old outputs and the new ones.
            Integers must be positive and smaller than the number of old outputs.
            tf.Tensors must belong to the old list of outputs.
            This mapping can be under-complete and can have repetitions.
        Returns:
          A new modified instance of the original subgraph view with remapped
            outputs.
        """
        res = copy.copy(self)
        res._remap_outputs(new_output_indices)  # pylint: disable=protected-access
        return res

    def remap(self, new_input_indices=None, new_output_indices=None):
        """Remap the inputs and outputs of the subgraph.

        Note that this is only modifying the view: the underlying tf.Graph is not
        affected.

        Args:
          new_input_indices: an iterable of integers or tf.Tensors
            representing a mapping between the old inputs and the new ones.
            Integers must be positive and smaller than the number of old inputs.
            tf.Tensors must belong to the old list of inputs.
            This mapping can be under-complete and must be without repetitions.
          new_output_indices: an iterable of integers or tf.Tensors
            representing a mapping between the old outputs and the new ones.
            Integers must be positive and smaller than the number of old outputs.
            tf.Tensors must belong to the old list of outputs.
            This mapping can be under-complete and can have repetitions.
        Returns:
          A new modified instance of the original subgraph view with remapped
            inputs and outputs.
        """
        res = copy.copy(self)
        if new_input_indices is not None:
            res._remap_inputs(new_input_indices)  # pylint: disable=protected-access
        if new_output_indices is not None:
            res._remap_outputs(new_output_indices)  # pylint: disable=protected-access
        return res

    def find_op_by_name(self, op_name):
        """Return the op named op_name.

        Args:
          op_name: the name to search for
        Returns:
          The op named op_name.
        Raises:
          ValueError: if the op_name could not be found.
          AssertionError: if the name was found multiple time.
        """
        res = [op for op in self._ops if op.name == op_name]
        if not res:
            raise ValueError("{} not in subgraph.".format(op_name))
        if len(res) > 1:
            raise AssertionError("More than 1 op named: {}!".format(op_name))
        return res[0]

    def __str__(self):
        if not self:
            return "SubGraphView: empty"

        def op_name(op):
            return op.name

        def tensor_name(t):
            if t in self._passthrough_ts:
                return "{} *".format(t.name)
            else:
                return t.name

        def print_list(name, iterable, get_name):
            if iterable:
                print("** {}[{}]:".format(name, len(iterable)), file=res)
                print("\n".join(["  {}".format(get_name(elem)) for elem in iterable]), file=res)
            else:
                print("** {}: empty".format(name), file=res)

        res = StringIO()
        print("SubGraphView (graphid={}):".format(id(self.graph)), file=res)
        print_list("ops", self._ops, op_name)
        print_list("inputs", self._input_ts, tensor_name)
        print_list("outputs", self._output_ts, tensor_name)
        return res.getvalue()

    @property
    def graph(self):
        """The underlying `tf.Graph`."""
        return self._graph

    @property
    def ops(self):
        """The operations in this subgraph view."""
        return self._ops

    @property
    def inputs(self):
        """The input tensors of this subgraph view."""
        return util.ListView(self._input_ts)

    @property
    def connected_inputs(self):
        """The connected input tensors of this subgraph view."""
        return [t for t in self._input_ts if t not in self._passthrough_ts]

    @property
    def outputs(self):
        """The output tensors of this subgraph view."""
        return util.ListView(self._output_ts)

    @property
    def connected_outputs(self):
        """The connected output tensors of this subgraph view."""
        return [t for t in self._output_ts if t not in self._passthrough_ts]

    @property
    def passthroughs(self):
        """The passthrough tensors, going straight from input to output."""
        return util.ListView(self._passthrough_ts)

    def __bool__(self):
        """Allows for implicit boolean conversion."""
        return self._graph is not None

    # Python 3 wants __bool__, Python 2.7 wants __nonzero__
    __nonzero__ = __bool__

    def op(self, op_id):
        """Get an op by its index."""
        return self._ops[op_id]

    def is_passthrough(self, t):
        """Check whether a tensor is passthrough."""
        return t in self._passthrough_ts

    def __enter__(self):
        """Allow Python context to minimize the life time of a subgraph view.

        A subgraph view is meant to be a lightweight and transient object. A short
        lifetime will alleviate the "out-of-sync" issue mentioned earlier. For that
        reason, a SubGraphView instance can be used within a Python context. For
        example:

        from tensorflow.contrib import graph_editor as ge
        with ge.make_sgv(...) as sgv:
          print(sgv)

        Returns:
          Itself.
        """
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass

    def input_index(self, t):
        """Find the input index corresponding to the given input tensor t.

        Args:
          t: the input tensor of this subgraph view.
        Returns:
          The index in the self.inputs list.
        Raises:
          Error: if t in not an input tensor.
        """
        try:
            subgraph_id = self._input_ts.index(t)
        except ValueError:
            raise ValueError("Can't find {} in inputs of subgraph {}.".format(t.name, self.name))
        return subgraph_id

    def output_index(self, t):
        """Find the output index corresponding to given output tensor t.

        Args:
          t: the output tensor of this subgraph view.
        Returns:
          The index in the self.outputs list.
        Raises:
          Error: if t in not an output tensor.
        """
        try:
            subgraph_id = self._output_ts.index(t)
        except ValueError:
            raise ValueError("Can't find {} in outputs of subgraph {}.".format(t.name, self.name))
        return subgraph_id

    def consumers(self):
        """Return a Python set of all the consumers of this subgraph view.

        A consumer of a subgraph view is a tf.Operation which is a consumer
        of one of the output tensors and is not in the subgraph.

        Returns:
          A list of `tf.Operation` which are the consumers of this subgraph view.
        """
        ops_set = frozenset(self._ops)
        res = []
        for output in self._output_ts:
            consumers = [op for op in output.consumers() if op not in ops_set]
            util.concatenate_unique(res, consumers)
        return res


def _check_graph(sgv, graph):
    """Check if sgv belongs to the given graph.

    Args:
      sgv: a SubGraphView.
      graph: a graph or None.
    Returns:
      The SubGraphView sgv.
    Raises:
      TypeError: if sgv is not a SubGraphView or if graph is not None and not
        a tf.Graph.
      ValueError: if the graph of sgv and the given graph are not None and
        different.
    """
    if not isinstance(sgv, SubGraphView):
        raise TypeError("Expected a SubGraphView, got: {}".format(type(graph)))
    if graph is None or not sgv.graph:
        return sgv
    if not isinstance(graph, tf_ops.Graph):
        raise TypeError("Expected a tf.Graph, got: {}".format(type(graph)))
    if sgv.graph is not graph:
        raise ValueError("Graph mismatch.")
    return sgv


def make_view(*args, **kwargs):
    """Create a SubGraphView from selected operations and passthrough tensors.

    Args:
      *args: list of 1) regular expressions (compiled or not) or 2) (array of)
        `tf.Operation` 3) (array of) `tf.Tensor`. Those objects will be converted
        into a list of operations and a list of candidate for passthrough tensors.
      **kwargs: keyword graph is used 1) to check that the ops and ts are from
        the correct graph 2) for regular expression query
    Returns:
      A subgraph view.
    Raises:
      TypeError: if the optional keyword argument graph is not a `tf.Graph`
        or if an argument in args is not an (array of) `tf.Tensor`
        or an (array of) `tf.Operation` or a string or a regular expression.
      ValueError: if one of the keyword arguments is unexpected.
    """
    # get keywords arguments
    graph = kwargs["graph"] if "graph" in kwargs else None

    # already a view?
    if len(args) == 1 and isinstance(args[0], SubGraphView):
        return _check_graph(args[0], graph)

    ops, ts = select.select_ops_and_ts(*args, **kwargs)
    sgv = SubGraphView(ops, ts)
    return _check_graph(sgv, graph)


def make_view_from_scope(scope, graph):
    """Make a subgraph from a name scope.

    Args:
      scope: the name of the scope.
      graph: the `tf.Graph`.
    Returns:
      A subgraph view representing the given scope.
    """
    ops = select.get_name_scope_ops(graph, scope)
    return SubGraphView(ops)
