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

"""Operations to emit summaries."""

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

import os

from tensorflow.contrib.summary import gen_summary_ops
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.layers import utils
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import summary_op_util
from tensorflow.python.training import training_util
from tensorflow.python.util import tf_contextlib

# Name for a collection which is expected to have at most a single boolean
# Tensor. If this tensor is True the summary ops will record summaries.
_SHOULD_RECORD_SUMMARIES_NAME = "ShouldRecordSummaries"

_SUMMARY_COLLECTION_NAME = "_SUMMARY_V2"
_SUMMARY_WRITER_INIT_COLLECTION_NAME = "_SUMMARY_WRITER_V2"


def should_record_summaries():
  """Returns boolean Tensor which is true if summaries should be recorded."""
  should_record_collection = ops.get_collection(_SHOULD_RECORD_SUMMARIES_NAME)
  if not should_record_collection:
    return False
  if len(should_record_collection) != 1:
    raise ValueError(
        "More than one tensor specified for whether summaries "
        "should be recorded: %s" % should_record_collection)
  return should_record_collection[0]


# TODO(apassos) consider how to handle local step here.
@tf_contextlib.contextmanager
def record_summaries_every_n_global_steps(n):
  """Sets the should_record_summaries Tensor to true if global_step % n == 0."""
  collection_ref = ops.get_collection_ref(_SHOULD_RECORD_SUMMARIES_NAME)
  old = collection_ref[:]
  with ops.device("cpu:0"):
    collection_ref[:] = [math_ops.equal(training_util.get_global_step() % n, 0)]
  yield
  collection_ref[:] = old


@tf_contextlib.contextmanager
def always_record_summaries():
  """Sets the should_record_summaries Tensor to always true."""
  collection_ref = ops.get_collection_ref(_SHOULD_RECORD_SUMMARIES_NAME)
  old = collection_ref[:]
  collection_ref[:] = [True]
  yield
  collection_ref[:] = old


@tf_contextlib.contextmanager
def never_record_summaries():
  """Sets the should_record_summaries Tensor to always false."""
  collection_ref = ops.get_collection_ref(_SHOULD_RECORD_SUMMARIES_NAME)
  old = collection_ref[:]
  collection_ref[:] = [False]
  yield
  collection_ref[:] = old


class SummaryWriter(object):
  """Encapsulates a summary writer."""

  def __init__(self, resource):
    self._resource = resource
    if context.in_eager_mode():
      self._resource_deleter = resource_variable_ops.EagerResourceDeleter(
          handle=self._resource, handle_device="cpu:0")

  def set_as_default(self):
    context.context().summary_writer_resource = self._resource

  @tf_contextlib.contextmanager
  def as_default(self):
    if self._resource is None:
      yield
    else:
      old = context.context().summary_writer_resource
      context.context().summary_writer_resource = self._resource
      yield
      # Flushes the summary writer in eager mode or in graph functions, but not
      # in legacy graph mode (you're on your own there).
      with ops.device("cpu:0"):
        gen_summary_ops.flush_summary_writer(self._resource)
      context.context().summary_writer_resource = old


def create_summary_file_writer(logdir,
                               max_queue=None,
                               flush_secs=None,
                               filename_suffix=None,
                               name=None):
  """Creates a summary file writer in the current context.

  Args: