#  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.
"""Tests for TaskRunner and Experiment class."""

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

import json
import os
import tempfile
import threading
import time

from tensorflow.contrib.learn.python.learn import evaluable
from tensorflow.contrib.learn.python.learn import experiment
from tensorflow.contrib.learn.python.learn import monitors
from tensorflow.contrib.learn.python.learn import run_config
from tensorflow.contrib.learn.python.learn import trainable
from tensorflow.contrib.learn.python.learn.estimators import run_config as run_config_lib
from tensorflow.contrib.learn.python.learn.utils import saved_model_export_utils
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging
from tensorflow.python.training import saver
from tensorflow.python.training import server_lib
from tensorflow.python.training import session_run_hook
from tensorflow.python.util import compat


class SheepCounter(object):
  """To be patched in for the time module, replacing sleep() and time()."""

  def __init__(self):
    self._total_time = 0
    self._sleeptimes = []
    self._time_calls = 0

  def sleep(self, t):
    self._total_time += t
    self._sleeptimes += [t]

  def time(self):
    self._time_calls += 1
    return self._total_time

  @property
  def sleep_times(self):
    return self._sleeptimes

  @property
  def time_calls(self):
    return self._time_calls


class TestEstimator(evaluable.Evaluable, trainable.Trainable):

  def __init__(self, config=None, max_evals=5):
    self.eval_count = 0
    self.fit_count = 0
    self._max_evals = max_evals
    self.export_count = 0
    self.monitors = []
    self.eval_hooks = []
    self._config = config or run_config.RunConfig()
    self._model_dir = tempfile.mkdtemp()

  @property
  def model_dir(self):
    return self._model_dir

  @property
  def config(self):
    return self._config

  def evaluate(self, **kwargs):
    tf_logging.info('evaluate called with args: %s' % kwargs)
    if 'hooks' in kwargs:
      self.eval_hooks = kwargs['hooks']
    self.eval_count += 1
    if self.eval_count > self._max_evals:
      tf_logging.info('Ran %d evals. Done.' % self.eval_count)
      raise StopIteration()
    return [(key, kwargs[key]) for key in sorted(kwargs.keys())]

  def fake_checkpoint(self):
    save_path = os.path.join(self.model_dir, 'model.ckpt')
    with session.Session() as sess:
      var = variables.Variable(1.0, name='var0')
      save = saver.Saver({var.op.name: var})
      var.initializer.run()
      save.save(sess, save_path, global_step=0)

  def fit(self, **kwargs):
    self.fake_checkpoint()
    tf_logging.info('fit called with args: %s' % kwargs)
    self.fit_count += 1
    if 'monitors' in kwargs:
      self.monitors = kwargs['monitors']
    return [(key, kwargs[key]) for key in sorted(kwargs.keys())]

  def export_savedmodel(self, export_dir_base, serving_input_fn, **kwargs):
    tf_logging.info('export_savedmodel called with args: %s, %s, %s' %
                    (export_dir_base, serving_input_fn, kwargs))
    self.export_count += 1
    return os.path.join(
        compat.as_bytes(export_dir_base), compat.as_bytes('bogus_timestamp'))


class _NoopHook(session_run_hook.SessionRunHook):
  pass


class ExperimentTest(test.TestCase):

  def _cluster_spec(self):
    return {
        run_config_lib.TaskType.PS: ['host1:2222', 'host2:2222'],
        run_config_lib.TaskType.WORKER:
            ['host3:2222', 'host4:2222', 'host5:2222']
    }

  def test_train(self):
    est = TestEstimator()
    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        train_steps='train_steps',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics')
    fit_args = ex.train(delay_secs=0)
    self.assertEqual(1, est.fit_count)
    self.assertIn(('max_steps', 'train_steps'), fit_args)
    self.assertEqual(0, est.eval_count)

  def test_train_delay(self):
    est = TestEstimator()
    ex = experiment.Experiment(
        est, train_input_fn='train_input', eval_input_fn='eval_input')
    for delay in [0, 1, 3]:
      sheep = SheepCounter()
      with test.mock.patch.object(time, 'time', sheep.time):
        with test.mock.patch.object(time, 'sleep', sheep.sleep):
          ex.train(delay_secs=delay)
          self.assertAlmostEqual(delay, sheep.time(), delta=1e-4)

  def test_train_default_delay(self):
    for task_id in [0, 1, 3]:
      tf_config = {'task': {'index': task_id}}
      with test.mock.patch.dict('os.environ',
                                {'TF_CONFIG': json.dumps(tf_config)}):
        config = run_config.RunConfig()
      est = TestEstimator(config)
      ex = experiment.Experiment(
          est, train_input_fn='train_input', eval_input_fn='eval_input')

      sheep = SheepCounter()
      with test.mock.patch.object(time, 'time', sheep.time):
        with test.mock.patch.object(time, 'sleep', sheep.sleep):
          ex.train()
          self.assertAlmostEqual(task_id * 5, sheep.time(), delta=1e-4)

  @test.mock.patch.object(server_lib, 'Server')
  def test_train_starts_server(self, mock_server):
    # Arrange.
    tf_config = {
        'cluster': self._cluster_spec(),
        'environment': run_config_lib.Environment.CLOUD,
        'task': {
            'type': run_config_lib.TaskType.WORKER,
            'index': 1
        }
    }
    with test.mock.patch.dict('os.environ',
                              {'TF_CONFIG': json.dumps(tf_config)}):
      config = run_config_lib.RunConfig(
          master='host4:2222', num_cores=15, gpu_memory_fraction=0.314)

    est = TestEstimator(config)
    ex = experiment.Experiment(
        est, train_input_fn='train_input', eval_input_fn='eval_input')

    # Act.
    # We want to make sure we discount the time it takes to start the server
    # in our accounting of the delay, so we set a small delay here.
    sheep = SheepCounter()
    with test.mock.patch.object(time, 'time', sheep.time):
      with test.mock.patch.object(time, 'sleep', sheep.sleep):
        ex.train(delay_secs=1)
        # Ensure that the delay takes into account the time to start the server.
        self.assertAlmostEqual(1, sheep.time(), delta=1e-4)

    # Assert.
    expected_config_proto = config_pb2.ConfigProto()
    expected_config_proto.inter_op_parallelism_threads = 15
    expected_config_proto.intra_op_parallelism_threads = 15
    expected_config_proto.gpu_options.per_process_gpu_memory_fraction = 0.314
    mock_server.assert_called_with(
        config.cluster_spec,
        job_name=run_config_lib.TaskType.WORKER,
        task_index=1,
        config=expected_config_proto,
        start=False)
    mock_server.assert_has_calls([test.mock.call().start()])

  @test.mock.patch.object(server_lib, 'Server')
  def test_train_server_does_not_start_without_cluster_spec(self, mock_server):
    config = run_config_lib.RunConfig(master='host4:2222')
    ex = experiment.Experiment(
        TestEstimator(config),
        train_input_fn='train_input',
        eval_input_fn='eval_input')
    ex.train()

    # The server should not have started because there was no ClusterSpec.
    self.assertFalse(mock_server.called)

  @test.mock.patch.object(server_lib, 'Server')
  def test_train_server_does_not_start_with_empty_master(self, mock_server):
    tf_config = {'cluster': self._cluster_spec()}
    with test.mock.patch.dict('os.environ',
                              {'TF_CONFIG': json.dumps(tf_config)}):
      config = run_config_lib.RunConfig(master='')
    ex = experiment.Experiment(
        TestEstimator(config),
        train_input_fn='train_input',
        eval_input_fn='eval_input')
    ex.train()

    # The server should not have started because master was the empty string.
    self.assertFalse(mock_server.called)

  def test_train_raises_if_job_name_is_missing(self):
    tf_config = {
        'cluster': self._cluster_spec(),
        'environment': run_config_lib.Environment.CLOUD,
        'task': {
            'index': 1
        }
    }
    with test.mock.patch.dict(
        'os.environ',
        {'TF_CONFIG': json.dumps(tf_config)}), self.assertRaises(ValueError):
      config = run_config_lib.RunConfig(
          master='host3:2222'  # Normally selected by task type.
      )
      ex = experiment.Experiment(
          TestEstimator(config),
          train_input_fn='train_input',
          eval_input_fn='eval_input')
      ex.train()

  def test_evaluate(self):
    est = TestEstimator()
    est.fake_checkpoint()
    noop_hook = _NoopHook()
    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics',
        eval_hooks=[noop_hook],
        eval_steps='steps',
        eval_delay_secs=0)
    ex.evaluate()
    self.assertEqual(0, est.fit_count)
    self.assertEqual(1, est.eval_count)
    self.assertEqual([noop_hook], est.eval_hooks)

  def test_evaluate_delay(self):
    est = TestEstimator()
    est.fake_checkpoint()
    noop_hook = _NoopHook()
    ex = experiment.Experiment(
        est, train_input_fn='train_input', eval_input_fn='eval_input',
        eval_hooks=[noop_hook])

    for delay in [0, 1, 3]:
      sheep = SheepCounter()
      with test.mock.patch.object(time, 'time', sheep.time):
        with test.mock.patch.object(time, 'sleep', sheep.sleep):
          ex.evaluate(delay_secs=delay)
      self.assertAlmostEqual(delay, sheep.time(), delta=1e-4)
      self.assertEqual([noop_hook], est.eval_hooks)

  def test_continuous_eval(self):
    est = TestEstimator()
    est.fake_checkpoint()
    noop_hook = _NoopHook()
    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics',
        eval_hooks=[noop_hook],
        eval_delay_secs=0,
        continuous_eval_throttle_secs=0)
    self.assertRaises(
        StopIteration, ex.continuous_eval, evaluate_checkpoint_only_once=False)
    self.assertEqual(0, est.fit_count)
    self.assertEqual(6, est.eval_count)
    self.assertEqual([noop_hook], est.eval_hooks)

  def test_continuous_eval_throttle_delay(self):
    for delay in [0, 1, 2]:
      est = TestEstimator()
      est.fake_checkpoint()
      noop_hook = _NoopHook()
      ex = experiment.Experiment(
          est,
          train_input_fn='train_input',
          eval_input_fn='eval_input',
          eval_metrics='eval_metrics',
          eval_hooks=[noop_hook],
          continuous_eval_throttle_secs=delay,
          eval_delay_secs=0)
      sheep = SheepCounter()
      with test.mock.patch.object(time, 'time', sheep.time):
        with test.mock.patch.object(time, 'sleep', sheep.sleep):
          self.assertRaises(
              StopIteration,
              ex.continuous_eval,
              evaluate_checkpoint_only_once=False)
          self.assertAlmostEqual(5 * delay, sheep.time(), delta=1e-4)

  def test_continuous_eval_predicate_fn(self):
    est = TestEstimator()
    est.fake_checkpoint()
    noop_hook = _NoopHook()

    def _predicate_fn(unused_eval_result):
      return est.eval_count < 3

    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics',
        eval_hooks=[noop_hook],
        eval_delay_secs=0,
        continuous_eval_throttle_secs=0)
    ex.continuous_eval(evaluate_checkpoint_only_once=False,
                       continuous_eval_predicate_fn=_predicate_fn)
    self.assertEqual(0, est.fit_count)
    self.assertEqual(3, est.eval_count)
    self.assertEqual([noop_hook], est.eval_hooks)

  def test_run_local(self):
    est = TestEstimator()
    noop_hook = _NoopHook()
    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics',
        eval_hooks=[noop_hook],
        train_steps=100,
        eval_steps=100,
        local_eval_frequency=10)
    ex.local_run()
    self.assertEqual(1, est.fit_count)
    self.assertEqual(1, est.eval_count)
    self.assertEqual(1, len(est.monitors))
    self.assertEqual([noop_hook], est.eval_hooks)
    self.assertTrue(isinstance(est.monitors[0], monitors.ValidationMonitor))

  def test_train_hooks_extend_does_not_mutate_input_hooks(self):
    noop_hook = _NoopHook()
    input_hooks = [noop_hook]

    ex = experiment.Experiment(
        TestEstimator(),
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics',
        train_monitors=input_hooks)
    self.assertAllEqual([noop_hook], ex._train_monitors)

    another_noop_hook = _NoopHook()
    # Assert that the extend API mutates the hooks, but not the input hooks
    ex.extend_train_hooks([another_noop_hook])
    self.assertAllEqual([noop_hook, another_noop_hook], ex._train_monitors)
    self.assertAllEqual([noop_hook], input_hooks)

  def test_export_strategies_reset(self):
    est = TestEstimator()
    export_strategy_1 = saved_model_export_utils.make_export_strategy(
        est, 'export_input_1', exports_to_keep=None)

    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics',
        train_steps=100,
        eval_steps=100,
        export_strategies=[export_strategy_1])
    ex.train_and_evaluate()
    self.assertEqual(1, est.export_count)

    # After reset with empty list (None), the count does not change and the user
    # provided export strategy list should remain intact.
    old_es = ex.reset_export_strategies()
    ex.train_and_evaluate()
    self.assertAllEqual([export_strategy_1], old_es)
    self.assertEqual(1, est.export_count)

    # After reset with list, the count should increase with the number of items.
    export_strategy_2 = saved_model_export_utils.make_export_strategy(
        est, 'export_input_2', exports_to_keep=None)
    export_strategy_3 = saved_model_export_utils.make_export_strategy(
        est, 'export_input_3', exports_to_keep=None)

    old_es = ex.reset_export_strategies([export_strategy_2, export_strategy_3])
    ex.train_and_evaluate()
    self.assertAllEqual([], old_es)
    self.assertEqual(3, est.export_count)

  def test_train_and_evaluate(self):
    est = TestEstimator()
    noop_hook = _NoopHook()
    export_strategy = saved_model_export_utils.make_export_strategy(
        est, 'export_input', exports_to_keep=None)
    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        eval_metrics='eval_metrics',
        eval_hooks=[noop_hook],
        train_steps=100,
        eval_steps=100,
        export_strategies=export_strategy)
    ex.train_and_evaluate()
    self.assertEqual(1, est.fit_count)
    self.assertEqual(1, est.eval_count)
    self.assertEqual(1, est.export_count)
    self.assertEqual(1, len(est.monitors))
    self.assertEqual([noop_hook], est.eval_hooks)
    self.assertTrue(isinstance(est.monitors[0], monitors.ValidationMonitor))

  @test.mock.patch.object(server_lib, 'Server')
  def test_run_std_server(self, mock_server):
    # Arrange.
    tf_config = {
        'cluster': self._cluster_spec(),
        'task': {
            'type': run_config_lib.TaskType.PS,
            'index': 1
        }
    }
    with test.mock.patch.dict('os.environ',
                              {'TF_CONFIG': json.dumps(tf_config)}):
      config = run_config_lib.RunConfig(
          master='host2:2222',
          num_cores=15,
          gpu_memory_fraction=0.314,)
    est = TestEstimator(config)
    ex = experiment.Experiment(
        est, train_input_fn='train_input', eval_input_fn='eval_input')

    # Act.
    ex.run_std_server()

    # Assert.
    mock_server.assert_has_calls(
        [test.mock.call().start(), test.mock.call().join()])

  @test.mock.patch.object(server_lib, 'Server')
  def test_run_std_server_raises_without_cluster_spec(self, mock_server):
    config = run_config_lib.RunConfig(master='host4:2222')
    with self.assertRaises(ValueError):
      ex = experiment.Experiment(
          TestEstimator(config),
          train_input_fn='train_input',
          eval_input_fn='eval_input')
      ex.run_std_server()

  def test_test(self):
    est = TestEstimator()
    exp_strategy = saved_model_export_utils.make_export_strategy(
        est, 'export_input', exports_to_keep=None)
    ex = experiment.Experiment(
        est,
        train_input_fn='train_input',
        eval_input_fn='eval_input',
        export_strategies=[exp_strategy])
    ex.test()
    self.assertEqual(1, est.fit_count)
    self.assertEqual(1, est.eval_count)
    self.assertEqual(1, est.export_count)

  def test_continuous_eval_evaluates_checkpoint_once(self):
    # Temporarily disabled until we figure out the threading story on Jenkins.
    return
    # pylint: disable=unreachable

    # The TestEstimator will raise StopIteration the second time evaluate is
    # called.
    ex = experiment.Experiment(
        TestEstimator(max_evals=1),
        train_input_fn='train_input',
        eval_input_fn='eval_input')

    # This should not happen if the logic restricting evaluation of the same
    # checkpoint works. We do need some checkpoint though, otherwise Experiment
    # will never evaluate.
    ex.estimator.fake_checkpoint()

    # Start a separate thread with continuous eval
    thread = threading.Thread(
        target=lambda: ex.continuous_eval(delay_secs=0, throttle_delay_secs=0))
    thread.start()

    # The thread will die if it evaluates twice, and we should never evaluate
    # twice since we don't write another checkpoint. Since we did not enable
    # throttling, if it hasn't died after two seconds, we're good.
    thread.join(2)
    self.assertTrue(thread.is_alive())

    # But we should have evaluated once.
    count = ex.estimator.eval_count
    self.assertEqual(1, count)


if __name__ == '__main__':
  test.main()
