# 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.
# ==============================================================================
"""Utility functions for summary creation."""

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

import functools
import re

import numpy as np

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import standard_ops

__all__ = ['assert_summary_tag_unique', 'is_summary_tag_unique',
           'summarize_tensor', 'summarize_activation', 'summarize_tensors',
           'summarize_collection', 'summarize_variables', 'summarize_weights',
           'summarize_biases', 'summarize_activations',]

# TODO(wicke): add more unit tests for summarization functions.


def assert_summary_tag_unique(tag):
  if not is_summary_tag_unique(tag):
    raise ValueError('Conflict with summary tag: %s already exists' % tag)


def _add_scalar_summary(tensor, tag=None):
  """Add a scalar summary operation for the tensor.

  Args:
    tensor: The tensor to summarize.
    tag: The tag to use, if None then use tensor's op's name.

  Returns:
    The created histogram summary.

  Raises:
    ValueError: If the tag is already in use or the rank is not 0.
  """
  tensor.get_shape().assert_has_rank(0)
  tag = tag or tensor.op.name
  assert_summary_tag_unique(tag)
  return standard_ops.scalar_summary(tag, tensor, name='%s_summary' % tag)


def _add_histogram_summary(tensor, tag=None):
  """Add a summary operation for the histogram of a tensor.

  Args:
    tensor: The tensor to summarize.
    tag: The tag to use, if None then use tensor's op's name.

  Returns:
    The created histogram summary.

  Raises:
    ValueError: If the tag is already in use.
  """
  tag = tag or tensor.op.name
  assert_summary_tag_unique(tag)
  return standard_ops.histogram_summary(tag, tensor, name='%s_summary' % tag)


def is_summary_tag_unique(tag):
  """Checks if a summary tag is unique.

  Args:
    tag: The tag to use

  Returns:
    True if the summary tag is unique.
  """
  existing_tags = [tensor_util.constant_value(summary.op.inputs[0])
                   for summary in ops.get_collection(ops.GraphKeys.SUMMARIES)]
  existing_tags = [name.tolist() if isinstance(name, np.ndarray) else name
                   for name in existing_tags]
  return tag.encode() not in existing_tags


def summarize_activation(op):
  """Summarize an activation.

  This applies the given activation and adds useful summaries specific to the
  activation.

  Args:
    op: The tensor to summarize (assumed to be a layer activation).
  Returns:
    The summary op created to summarize `op`.
  """
  if op.op.type in ('Relu', 'Softplus', 'Relu6'):
    # Using inputs to avoid floating point equality and/or epsilons.
    _add_scalar_summary(
        standard_ops.reduce_mean(standard_ops.to_float(standard_ops.less(
            op.op.inputs[0], standard_ops.cast(0.0, op.op.inputs[0].dtype)))),
        '%s/zeros' % op.op.name)
  if op.op.type == 'Relu6':
    _add_scalar_summary(
        standard_ops.reduce_mean(standard_ops.to_float(standard_ops.greater(
            op.op.inputs[0], standard_ops.cast(6.0, op.op.inputs[0].dtype)))),
        '%s/sixes' % op.op.name)
  return _add_histogram_summary(op, '%s/activation' % op.op.name)


def summarize_tensor(tensor, tag=None):
  """Summarize a tensor using a suitable summary type.

  This function adds a summary op for `tensor`. The type of summary depends on
  the shape of `tensor`. For scalars, a `scalar_summary` is created, for all
  other tensors, `histogram_summary` is used.

  Args:
    tensor: The tensor to summarize
    tag: The tag to use, if None then use tensor's op's name.

  Returns:
    The summary op created or None for string tensors.
  """
  # Skips string tensors and boolean tensors (not handled by the summaries).
  if (tensor.dtype.is_compatible_with(dtypes.string) or
      tensor.dtype.base_dtype == dtypes.bool):
    return None

  if tensor.get_shape().ndims == 0:
    # For scalars, use a scalar summary.
    return _add_scalar_summary(tensor, tag)
  else:
    # We may land in here if the rank is still unknown. The histogram won't
    # hurt if this ends up being a scalar.
    return _add_histogram_summary(tensor, tag)


def summarize_tensors(tensors, summarizer=summarize_tensor):
  """Summarize a set of tensors."""
  return [summarizer(tensor) for tensor in tensors]


def summarize_collection(collection, name_filter=None,
                         summarizer=summarize_tensor):
  """Summarize a graph collection of tensors, possibly filtered by name."""
  tensors = []
  for op in ops.get_collection(collection):
    if name_filter is None or re.match(name_filter, op.op.name):
      tensors.append(op)
  return summarize_tensors(tensors, summarizer)


# Utility functions for commonly used collections
summarize_variables = functools.partial(summarize_collection,
                                        ops.GraphKeys.VARIABLES)


summarize_weights = functools.partial(summarize_collection,
                                      ops.GraphKeys.WEIGHTS)


summarize_biases = functools.partial(summarize_collection,
                                     ops.GraphKeys.BIASES)


def summarize_activations(name_filter=None, summarizer=summarize_activation):
  """Summarize activations, using `summarize_activation` to summarize."""
  return summarize_collection(ops.GraphKeys.ACTIVATIONS, name_filter,
                              summarizer)
