# Copyright 2016 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.
# ==============================================================================
"""Monitors tests."""

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

import collections
import shutil
import tempfile
import time

from six.moves import xrange  # pylint: disable=redefined-builtin

from tensorflow.contrib import testing
from tensorflow.contrib.framework.python.framework import checkpoint_utils
from tensorflow.contrib.framework.python.ops import variables as variables_lib
from tensorflow.contrib.learn.python import learn
from tensorflow.contrib.learn.python.learn import estimators
from tensorflow.python.client import session as session_lib
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.summary import summary
from tensorflow.python.training import gradient_descent
from tensorflow.python.training import monitored_session
from tensorflow.python.training import saver


class _MyEveryN(learn.monitors.EveryN):

  def __init__(self, every_n_steps=100, first_n_steps=1):
    super(_MyEveryN, self).__init__(
        every_n_steps=every_n_steps, first_n_steps=first_n_steps)
    self._steps_begun = []
    self._steps_ended = []
    self._post_steps = []

  @property
  def steps_begun(self):
    return self._steps_begun

  @property
  def steps_ended(self):
    return self._steps_ended

  @property
  def post_steps(self):
    return self._post_steps

  def every_n_step_begin(self, step):
    super(_MyEveryN, self).every_n_step_begin(step)
    self._steps_begun.append(step)
    return []

  def every_n_step_end(self, step, outputs):
    super(_MyEveryN, self).every_n_step_end(step, outputs)
    self._steps_ended.append(step)
    return False

  def every_n_post_step(self, step, session):
    super(_MyEveryN, self).every_n_post_step(step, session)
    self._post_steps.append(step)
    return False


class MonitorsTest(test.TestCase):
  """Monitors tests."""

  def setUp(self):
    # Mock out logging calls so we can verify whether correct tensors are being
    # monitored.
    self._actual_log = logging.info

    def mockLog(*args, **kwargs):  # pylint: disable=invalid-name
      self.logged_message = args
      self._actual_log(*args, **kwargs)

    logging.info = mockLog

  def tearDown(self):
    logging.info = self._actual_log

  def _run_monitor(self,
                   monitor,
                   num_epochs=3,
                   num_steps_per_epoch=10,
                   pass_max_steps=True):
    if pass_max_steps:
      max_steps = num_epochs * num_steps_per_epoch - 1
    else:
      max_steps = None
    monitor.begin(max_steps=max_steps)
    for epoch in xrange(num_epochs):
      monitor.epoch_begin(epoch)
      should_stop = False
      step = epoch * num_steps_per_epoch
      next_epoch_step = step + num_steps_per_epoch
      while (not should_stop) and (step < next_epoch_step):
        tensors = monitor.step_begin(step)
        output = ops.get_default_session().run(tensors) if tensors else {}
        output = dict(
            zip([t.name if isinstance(t, ops.Tensor) else t for t in tensors],
                output))
        should_stop = monitor.step_end(step=step, output=output)
        monitor.post_step(step=step, session=None)
        step += 1
      monitor.epoch_end(epoch)
    monitor.end()

  def test_base_monitor(self):
    with ops.Graph().as_default() as g, self.test_session(g):
      self._run_monitor(learn.monitors.BaseMonitor())

  def test_every_0(self):
    monitor = _MyEveryN(every_n_steps=0, first_n_steps=-1)
    with ops.Graph().as_default() as g, self.test_session(g):
      self._run_monitor(monitor, num_epochs=3, num_steps_per_epoch=10)
      expected_steps = list(range(30))
      self.assertAllEqual(expected_steps, monitor.steps_begun)
      self.assertAllEqual(expected_steps, monitor.steps_ended)