"""Tests for AdaNet AutoEnsembleEstimator in TF 1.

Copyright 2018 The AdaNet 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

    https://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.
"""

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

import contextlib
import os
import shutil
import sys

from absl import flags
from absl import logging
from absl.testing import parameterized
from adanet import tf_compat
from adanet.autoensemble.common import AutoEnsembleSubestimator
from adanet.autoensemble.estimator import _GeneratorFromCandidatePool
from adanet.autoensemble.estimator import AutoEnsembleEstimator
import tensorflow.compat.v1 as tf

# pylint: disable=g-direct-tensorflow-import
from tensorflow.python.estimator.export import export
# pylint: enable=g-direct-tensorflow-import

logging.set_verbosity(logging.INFO)


# Ensures "local_init_op" is called.
class CheckLocalInitOpEstimator(tf.estimator.Estimator):

  def __init__(self):
    super(CheckLocalInitOpEstimator,
          self).__init__(model_fn=self._get_model_fn())

  def _get_model_fn(self):

    def _model_fn(features, labels, mode, params):

      del labels
      del params

      flag = tf.Variable(initial_value=False, collections=[])
      set_flag = tf.assign(flag, True)
      test_flag = tf.debugging.Assert(tf.equal(flag, True), [flag])

      scaffold = tf.train.Scaffold(
          local_init_op=tf.group(tf.train.Scaffold.default_local_init_op(),
                                 set_flag))

      # Note: Not consuming the feature stales the test input_fn.
      feature = next(iter(features.values()))
      with tf.control_dependencies([feature, test_flag]):
        batch_size = tf.shape(feature)[0]
        predictions = tf.zeros([batch_size, 1])

      return tf.estimator.EstimatorSpec(
          mode,
          loss=tf.constant(-1.0),
          train_op=test_flag,
          scaffold=scaffold,
          predictions=predictions)

    return _model_fn


# Creates an a file in the model directory during training, and check the
# presence of this file during eval and inference.
class CheckAssetEstimator(tf.estimator.Estimator):

  def __init__(self, config):
    super(CheckAssetEstimator, self).__init__(
        config=config, model_fn=self._get_model_fn())

  def _get_model_fn(self):

    def _model_fn(features, labels, mode, params):

      del labels
      del params

      scaffold = None
      training_hooks = None

      if mode == tf.estimator.ModeKeys.TRAIN:

        class CheckpointSaverListener(tf.train.CheckpointSaverListener):

          def after_save(subself, session, global_step_value):  # pylint: disable=no-self-argument
            tf.logging.info("Creating data in %s", self._model_dir)
            with tf.gfile.Open(os.path.join(self._model_dir, "data"), "w") as f:
              f.write("Hello")
            assert tf.io.gfile.exists(os.path.join(self._model_dir, "data"))

        listener = CheckpointSaverListener()
        saver_hook = tf.estimator.CheckpointSaverHook(
            self._model_dir, listeners=[listener], save_steps=1000000)
        training_hooks = [saver_hook]

      else:

        def check_asset():
          tf.logging.info("Checking data %s", self._model_dir)
          assert tf.io.gfile.exists(os.path.join(self._model_dir, "data"))
          return 1

        check_asset_op = tf.compat.v1.py_func(check_asset, [], tf.int64)
        scaffold = tf.train.Scaffold(
            local_init_op=tf.group(tf.train.Scaffold.default_local_init_op(),
                                   check_asset_op))

      feature = next(iter(features.values()))
      with tf.control_dependencies([feature]):
        batch_size = tf.shape(feature)[0]
        predictions = tf.zeros([batch_size, 1])
        train_op = tf.no_op()
        loss = tf.constant(-1.0)

      return tf.estimator.EstimatorSpec(
          mode,
          loss=loss,
          train_op=train_op,
          scaffold=scaffold,
          predictions=predictions,
          training_hooks=training_hooks)

    return _model_fn


class AutoEnsembleEstimatorTest(parameterized.TestCase, tf.test.TestCase):

  def setUp(self):
    super(AutoEnsembleEstimatorTest, self).setUp()
    # Setup and cleanup test directory.
    # Flags are not automatically parsed at this point.
    flags.FLAGS(sys.argv)
    self.test_subdirectory = os.path.join(flags.FLAGS.test_tmpdir, self.id())
    shutil.rmtree(self.test_subdirectory, ignore_errors=True)
    os.makedirs(self.test_subdirectory)

  def tearDown(self):
    super(AutoEnsembleEstimatorTest, self).tearDown()
    shutil.rmtree(self.test_subdirectory, ignore_errors=True)

  # pylint: disable=g-long-lambda
  @parameterized.named_parameters(
      {
          "testcase_name":
              "dict_candidate_pool",
          "candidate_pool":
              lambda head, feature_columns, optimizer: {
                  "dnn":
                      tf.estimator.DNNEstimator(
                          head=head,
                          feature_columns=feature_columns,
                          optimizer=optimizer,
                          hidden_units=[3]),
                  "linear":
                      tf.estimator.LinearEstimator(
                          head=head,
                          feature_columns=feature_columns,
                          optimizer=optimizer),
              },
          "want_loss":
              .209,
      },
      {
          "testcase_name":
              "list_candidate_pool",
          "candidate_pool":
              lambda head, feature_columns, optimizer: [
                  tf.estimator.DNNEstimator(
                      head=head,
                      feature_columns=feature_columns,
                      optimizer=optimizer,
                      hidden_units=[3]),
                  tf.estimator.LinearEstimator(
                      head=head,
                      feature_columns=feature_columns,
                      optimizer=optimizer),
              ],
          "want_loss":
              .209,
      },
      {
          "testcase_name":
              "candidate_pool_lambda",
          "candidate_pool":
              lambda head, feature_columns, optimizer: lambda config: {
                  "dnn":
                      tf.estimator.DNNEstimator(
                          head=head,
                          feature_columns=feature_columns,
                          optimizer=optimizer,
                          hidden_units=[3],
                          config=config),
                  "linear":
                      tf.estimator.LinearEstimator(
                          head=head,
                          feature_columns=feature_columns,
                          optimizer=optimizer,
                          config=config),
              },
          "want_loss":
              .209,
      },
      {
          "testcase_name":
              "bagging",
          "candidate_pool":
              lambda head, feature_columns, optimizer: {
                  "same_train_data":
                      AutoEnsembleSubestimator(
                          tf.estimator.LinearEstimator(
                              head=head,
                              feature_columns=feature_columns,
                              optimizer=optimizer)),
                  "different_train_data":
                      AutoEnsembleSubestimator(
                          tf.estimator.DNNEstimator(
                              head=head,
                              feature_columns=feature_columns,
                              optimizer=optimizer,
                              hidden_units=[3]),
                          train_input_fn=lambda: ({
                              "input_1": tf.constant([[0., 1.]])
                          }, tf.constant([[1.]]))),
              },
          # TODO: Figure out why this test's loss changes with every
          # change to the TensorFlow graph.
          "want_loss":
              0.2,
      },
      {
          "testcase_name":
              "bagging_out_of_range_error",
          "max_train_steps":
              15,
          "candidate_pool":
              lambda head, feature_columns, optimizer: {
                  "same_train_data":
                      AutoEnsembleSubestimator(
                          tf.estimator.LinearEstimator(
                              head=head,
                              feature_columns=feature_columns,
                              optimizer=optimizer)),
                  "different_train_data":
                      AutoEnsembleSubestimator(
                          tf.estimator.DNNEstimator(
                              head=head,
                              feature_columns=feature_columns,
                              optimizer=optimizer,
                              hidden_units=[3]),
                          # TODO: Dataset must have at least 2 batches,
                          # otherwise all of training terminates.
                          train_input_fn=lambda: tf.data.Dataset.
                          from_tensor_slices(({
                              "input_1": [[0., 1.], [0., 1.]]
                          }, [[1.], [1.]])).batch(1),
                      ),
              },
          # TODO: Figure out why this test's loss changes with every
          # change to the TensorFlow graph.
          "want_loss":
              0.3,
      },
      {
          "testcase_name":
              "check_local_init_op",
          "max_train_steps":
              10,
          "candidate_pool":
              lambda head, feature_columns, optimizer: {
                  "expect_local_init_op":
                      AutoEnsembleSubestimator(CheckLocalInitOpEstimator()),
              },
          "want_loss":
              1.0,
      },
      {
          "testcase_name":
              "check_iteration_count",
          "max_train_steps":
              10,
          "candidate_pool":
              lambda head, feature_columns, optimizer:
              (lambda config, iteration_number: {
                  "dnn":
                      AutoEnsembleSubestimator(
                          tf.estimator.DNNEstimator(
                              head=head,
                              feature_columns=feature_columns,
                              optimizer=optimizer,
                              hidden_units=[3])),
                  "linear":
                      AutoEnsembleSubestimator(
                          tf.estimator.LinearEstimator(
                              head=head,
                              feature_columns=feature_columns,
                              optimizer=optimizer)),
              }),
          "want_loss":
              .209,
      },
      {
          "testcase_name":
              "check_has_asset",
          "max_train_steps":
              10,
          "candidate_pool":
              lambda head, feature_columns, optimizer: {
                  "linear":
                      lambda subconfig: AutoEnsembleSubestimator(
                          tf.estimator.LinearEstimator(
                              head=head,
                              feature_columns=feature_columns,
                              optimizer=optimizer,
                              config=subconfig)),
                  "check_asset_1":
                      lambda subconfig: AutoEnsembleSubestimator(
                          CheckAssetEstimator(config=subconfig)),
                  "check_asset_2":
                      lambda subconfig: AutoEnsembleSubestimator(
                          CheckAssetEstimator(config=subconfig)),
              },
          "want_loss":
              .209,
      })
  # pylint: enable=g-long-lambda
  @tf_compat.skip_for_tf2
  def test_auto_ensemble_estimator_lifecycle(self,
                                             candidate_pool,
                                             want_loss,
                                             max_train_steps=30):
    features = {"input_1": [[1., 0.]]}
    labels = [[1.]]

    run_config = tf.estimator.RunConfig(tf_random_seed=42)
    head = tf.contrib.estimator.regression_head(
        loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)

    optimizer = tf.train.GradientDescentOptimizer(learning_rate=.01)
    feature_columns = [tf.feature_column.numeric_column("input_1", shape=[2])]

    def train_input_fn():
      input_features = {}
      for key, feature in features.items():
        input_features[key] = tf.constant(feature, name=key)
      input_labels = tf.constant(labels, name="labels")
      return input_features, input_labels

    def test_input_fn():
      input_features = tf.data.Dataset.from_tensors([
          tf.constant(features["input_1"])
      ]).make_one_shot_iterator().get_next()
      return {"input_1": input_features}, None

    estimator = AutoEnsembleEstimator(
        head=head,
        candidate_pool=candidate_pool(head, feature_columns, optimizer),
        max_iteration_steps=10,
        force_grow=True,
        model_dir=self.test_subdirectory,
        config=run_config)

    # Train for three iterations.
    estimator.train(input_fn=train_input_fn, max_steps=max_train_steps)

    # Evaluate.
    eval_results = estimator.evaluate(input_fn=train_input_fn, steps=1)

    self.assertAllClose(max_train_steps, eval_results["global_step"])
    self.assertAllClose(want_loss, eval_results["loss"], atol=.3)

    # Predict.
    predictions = estimator.predict(input_fn=test_input_fn)
    for prediction in predictions:
      self.assertIsNotNone(prediction["predictions"])

    # Export SavedModel.
    def serving_input_fn():
      """Input fn for serving export, starting from serialized example."""
      serialized_example = tf.placeholder(
          dtype=tf.string, shape=(None), name="serialized_example")
      for key, value in features.items():
        features[key] = tf.constant(value)
      return export.SupervisedInputReceiver(
          features=features,
          labels=tf.constant(labels),
          receiver_tensors=serialized_example)

    export_dir_base = os.path.join(self.test_subdirectory, "export")
    export_saved_model_fn = getattr(estimator, "export_saved_model", None)
    if not callable(export_saved_model_fn):
      export_saved_model_fn = estimator.export_savedmodel
    export_saved_model_fn(
        export_dir_base=export_dir_base,
        serving_input_receiver_fn=serving_input_fn)

  @tf_compat.skip_for_tf2
  def test_last_layer_fn(self):
    head = tf.contrib.estimator.regression_head(
        loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)
    optimizer = tf.train.GradientDescentOptimizer(learning_rate=.01)
    feature_columns = [tf.feature_column.numeric_column("input_1", shape=[2])]
    cand_pool = [
        tf.estimator.DNNEstimator(
            head=head,
            feature_columns=feature_columns,
            optimizer=optimizer,
            hidden_units=[3])
    ]
    input_features = {}
    features = {"input_1": [[1., 0.]]}
    labels = [[1.]]
    for key, feature in features.items():
      input_features[key] = tf.constant(feature, name=key)
    input_labels = tf.constant(labels, name="labels")

    class _FakeSummary(object):
      """A fake adanet.Summary."""

      def scalar(self, name, tensor, family=None):
        del name, tensor, family
        return "fake_scalar"

      @contextlib.contextmanager
      def current_scope(self):
        yield

    def _adanet_last_layer_fn(estimator_spec):
      del estimator_spec
      return input_labels

    # Call with custom last_layer_fn which simply returns the labels tensor.
    generator = _GeneratorFromCandidatePool(
        cand_pool, logits_fn=None, last_layer_fn=_adanet_last_layer_fn)
    candidates = generator.generate_candidates(
        previous_ensemble=None,
        iteration_number=None,
        previous_ensemble_reports=None,
        all_reports=None,
        config=tf.estimator.RunConfig())
    subnetwork = candidates[0].build_subnetwork(input_features,
                                                input_labels, None, False, 1,
                                                _FakeSummary(), None)

    self.assertEqual(input_labels, subnetwork.last_layer)

  @tf_compat.skip_for_tf2
  def test_extra_checkpoint_saver_hook(self):
    """Tests b/122795064."""

    features = {"input_1": [[1., 0.]]}
    labels = [[1.]]

    run_config = tf.estimator.RunConfig(tf_random_seed=42)
    head = tf.contrib.estimator.binary_classification_head(
        loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)

    optimizer = tf.train.GradientDescentOptimizer(learning_rate=.01)
    feature_columns = [tf.feature_column.numeric_column("input_1", shape=[2])]

    estimator = AutoEnsembleEstimator(
        head=head,
        candidate_pool=[
            tf.estimator.LinearClassifier(
                n_classes=2,
                feature_columns=feature_columns,
                optimizer=optimizer),
            tf.estimator.DNNClassifier(
                n_classes=2,
                feature_columns=feature_columns,
                optimizer=optimizer,
                hidden_units=[3]),
        ],
        max_iteration_steps=3,
        force_grow=True,
        model_dir=self.test_subdirectory,
        config=run_config)

    ckpt_dir = os.path.join(self.test_subdirectory)
    hooks = [tf.train.CheckpointSaverHook(ckpt_dir, save_steps=1)]

    def train_input_fn():
      input_features = {}
      for key, feature in features.items():
        input_features[key] = tf.constant(feature, name=key)
      input_labels = tf.constant(labels, name="labels")
      return input_features, input_labels

    estimator.train(input_fn=train_input_fn, max_steps=6, hooks=hooks)


if __name__ == "__main__":
  tf.test.main()
