# Copyright 2017 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.
# ==============================================================================
"""A python interface for Grappler items."""

from tensorflow.core.grappler.costs import op_performance_data_pb2
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.python.grappler import _pywrap_tf_item as tf_item


class Item(object):
  """GrapplerItem."""

  def __init__(self,
               metagraph,
               ignore_colocation=True,
               ignore_user_placement=False):
    """Creates an Item.

    Args:
      metagraph: a TensorFlow metagraph.
      ignore_colocation: if set, the tool will ignore all the colocation
        constraints generated by TensorFlow.
      ignore_user_placement: if set, all the placement annotations annotated in
        the metagraph will be ignored.
    Raises:
      ValueError: the metagraph is incomplete or invalid.
    """
    self._metagraph = metagraph
    self._item_graph = meta_graph_pb2.MetaGraphDef()
    self._item_graph.CopyFrom(metagraph)
    self._ignore_colocation = ignore_colocation
    self._ignore_user_placement = ignore_user_placement
    self._tf_item = None
    self._BuildTFItem()

  def IdentifyImportantOps(self, sort_topologically=False):
    return tf_item.TF_IdentifyImportantOps(self.tf_item, sort_topologically)

  def GetOpProperties(self):
    """Get Op properties."""
    props = tf_item.TF_GetOpProperties(self.tf_item)
    properties = {}
    for key, values in props.items():
      prop = []
      for value in values:
        # TODO(petebu): Make this conversion to a dictionary be done in the C++
        # wrapper for performance.
        prop.append(
            op_performance_data_pb2.OpInfo.TensorProperties.FromString(value))
      properties[key] = prop
    return properties

  def GetColocationGroups(self):
    """Return a list of hard colocation constraints.

    All the nodes in a colocation tuple must be placed on the same device for
    the model to work.

    Returns:
      A list of colocation tuples.
    """
    return tf_item.TF_GetColocationGroups(self.tf_item)

  @property
  def metagraph(self):
    return self._metagraph

  @property
  def tf_item(self):
    if self._item_graph != self._metagraph:
      self._BuildTFItem()
      self._item_graph.CopyFrom(self._metagraph)
    return self._tf_item

  def _BuildTFItem(self):
    self._tf_item = tf_item.TF_NewItem(self._metagraph.SerializeToString(),
                                       self._ignore_colocation,
                                       self._ignore_user_placement)