# 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.
# ==============================================================================
"""Tests for v1 version of dnn_linear_combined.py."""

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

import shutil
import tempfile

from absl.testing import parameterized
import numpy as np
import six
import tensorflow as tf
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
from tensorflow.python.feature_column import feature_column
from tensorflow.python.feature_column import feature_column_v2
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow_estimator.python.estimator import estimator
from tensorflow_estimator.python.estimator.canned import dnn_linear_combined
from tensorflow_estimator.python.estimator.canned import prediction_keys
from tensorflow_estimator.python.estimator.canned.v1 import dnn_testing_utils_v1
from tensorflow_estimator.python.estimator.canned.v1 import linear_testing_utils_v1
from tensorflow_estimator.python.estimator.export import export
from tensorflow_estimator.python.estimator.inputs import numpy_io
from tensorflow_estimator.python.estimator.inputs import pandas_io

try:
  # pylint: disable=g-import-not-at-top
  import pandas as pd
  HAS_PANDAS = True
except IOError:
  # Pandas writes a temporary file during import. If it fails, don't use pandas.
  HAS_PANDAS = False
except ImportError:
  HAS_PANDAS = False

# This is so that we can easily switch between feature_column and
# feature_column_v2 for testing.
# Note that following V2 version of tests are for feature_column_v2, not the v2
# version of canned estimator.
feature_column.numeric_column = feature_column._numeric_column
feature_column.categorical_column_with_hash_bucket = feature_column._categorical_column_with_hash_bucket  # pylint: disable=line-too-long
feature_column.categorical_column_with_vocabulary_list = feature_column._categorical_column_with_vocabulary_list  # pylint: disable=line-too-long
feature_column.categorical_column_with_vocabulary_file = feature_column._categorical_column_with_vocabulary_file  # pylint: disable=line-too-long
feature_column.embedding_column = feature_column._embedding_column


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyModelFnTest(dnn_testing_utils_v1.BaseDNNModelFnTest,
                         tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNModelFnTest.__init__(self,
                                                     self._dnn_only_model_fn)

  def _dnn_only_model_fn(self,
                         features,
                         labels,
                         mode,
                         head,
                         hidden_units,
                         feature_columns,
                         optimizer='Adagrad',
                         activation_fn=tf.nn.relu,
                         dropout=None,
                         input_layer_partitioner=None,
                         config=None):
    return dnn_linear_combined._dnn_linear_combined_model_fn(
        features=features,
        labels=labels,
        mode=mode,
        head=head,
        linear_feature_columns=[],
        dnn_hidden_units=hidden_units,
        dnn_feature_columns=feature_columns,
        dnn_optimizer=optimizer,
        dnn_activation_fn=activation_fn,
        dnn_dropout=dropout,
        input_layer_partitioner=input_layer_partitioner,
        config=config)


# A function to mimic linear-regressor init reuse same tests.
def _linear_regressor_fn(feature_columns,
                         model_dir=None,
                         label_dimension=1,
                         weight_column=None,
                         optimizer='Ftrl',
                         config=None,
                         partitioner=None,
                         sparse_combiner='sum'):
  return dnn_linear_combined.DNNLinearCombinedRegressor(
      model_dir=model_dir,
      linear_feature_columns=feature_columns,
      linear_optimizer=optimizer,
      label_dimension=label_dimension,
      weight_column=weight_column,
      input_layer_partitioner=partitioner,
      config=config,
      linear_sparse_combiner=sparse_combiner)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorPartitionerTest(
    linear_testing_utils_v1.BaseLinearRegressorPartitionerTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorPartitionerTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorPartitionerV2Test(
    linear_testing_utils_v1.BaseLinearRegressorPartitionerTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorPartitionerTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorEvaluationTest(
    linear_testing_utils_v1.BaseLinearRegressorEvaluationTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorEvaluationTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorEvaluationV2Test(
    linear_testing_utils_v1.BaseLinearRegressorEvaluationTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorEvaluationTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorPredictTest(
    linear_testing_utils_v1.BaseLinearRegressorPredictTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorPredictTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorPredictV2Test(
    linear_testing_utils_v1.BaseLinearRegressorPredictTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorPredictTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorIntegrationTest(
    linear_testing_utils_v1.BaseLinearRegressorIntegrationTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorIntegrationTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorIntegrationV2Test(
    linear_testing_utils_v1.BaseLinearRegressorIntegrationTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorIntegrationTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorTrainingTest(
    linear_testing_utils_v1.BaseLinearRegressorTrainingTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorTrainingTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyRegressorTrainingV2Test(
    linear_testing_utils_v1.BaseLinearRegressorTrainingTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearRegressorTrainingTest.__init__(
        self, _linear_regressor_fn, fc_lib=feature_column_v2)


def _linear_classifier_fn(feature_columns,
                          model_dir=None,
                          n_classes=2,
                          weight_column=None,
                          label_vocabulary=None,
                          optimizer='Ftrl',
                          config=None,
                          partitioner=None,
                          sparse_combiner='sum'):
  return dnn_linear_combined.DNNLinearCombinedClassifier(
      model_dir=model_dir,
      linear_feature_columns=feature_columns,
      linear_optimizer=optimizer,
      n_classes=n_classes,
      weight_column=weight_column,
      label_vocabulary=label_vocabulary,
      input_layer_partitioner=partitioner,
      config=config,
      linear_sparse_combiner=sparse_combiner)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyClassifierTrainingTest(
    linear_testing_utils_v1.BaseLinearClassifierTrainingTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearClassifierTrainingTest.__init__(
        self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyClassifierTrainingV2Test(
    linear_testing_utils_v1.BaseLinearClassifierTrainingTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearClassifierTrainingTest.__init__(
        self,
        linear_classifier_fn=_linear_classifier_fn,
        fc_lib=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyClassifierClassesEvaluationTest(
    linear_testing_utils_v1.BaseLinearClassifierEvaluationTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearClassifierEvaluationTest.__init__(
        self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyClassifierClassesEvaluationV2Test(
    linear_testing_utils_v1.BaseLinearClassifierEvaluationTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearClassifierEvaluationTest.__init__(
        self,
        linear_classifier_fn=_linear_classifier_fn,
        fc_lib=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyClassifierPredictTest(
    linear_testing_utils_v1.BaseLinearClassifierPredictTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearClassifierPredictTest.__init__(
        self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyClassifierPredictV2Test(
    linear_testing_utils_v1.BaseLinearClassifierPredictTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearClassifierPredictTest.__init__(
        self,
        linear_classifier_fn=_linear_classifier_fn,
        fc_lib=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyClassifierIntegrationTest(
    linear_testing_utils_v1.BaseLinearClassifierIntegrationTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearClassifierIntegrationTest.__init__(
        self, linear_classifier_fn=_linear_classifier_fn, fc_lib=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class LinearOnlyClassifierIntegrationV2Test(
    linear_testing_utils_v1.BaseLinearClassifierIntegrationTest,
    tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    linear_testing_utils_v1.BaseLinearClassifierIntegrationTest.__init__(
        self,
        linear_classifier_fn=_linear_classifier_fn,
        fc_lib=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNLinearCombinedRegressorIntegrationTest(tf.test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def _test_complete_flow_helper(self, linear_feature_columns,
                                 dnn_feature_columns, feature_spec,
                                 train_input_fn, eval_input_fn,
                                 predict_input_fn, input_dimension,
                                 label_dimension, batch_size):
    est = dnn_linear_combined.DNNLinearCombinedRegressor(
        linear_feature_columns=linear_feature_columns,
        dnn_hidden_units=(2, 2),
        dnn_feature_columns=dnn_feature_columns,
        label_dimension=label_dimension,
        model_dir=self._model_dir)

    # TRAIN
    num_steps = 10
    est.train(train_input_fn, steps=num_steps)

    # EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(num_steps, scores[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertIn('loss', six.iterkeys(scores))

    # PREDICT
    predictions = np.array([
        x[prediction_keys.PredictionKeys.PREDICTIONS]
        for x in est.predict(predict_input_fn)
    ])
    self.assertAllEqual((batch_size, label_dimension), predictions.shape)

    # EXPORT
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_saved_model(tempfile.mkdtemp(),
                                        serving_input_receiver_fn)
    self.assertTrue(tf.compat.v1.gfile.Exists(export_dir))

  def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
                          input_dimension, label_dimension, batch_size,
                          fc_impl):
    linear_feature_columns = [
        fc_impl.numeric_column('x', shape=(input_dimension,))
    ]
    dnn_feature_columns = [
        fc_impl.numeric_column('x', shape=(input_dimension,))
    ]
    feature_columns = linear_feature_columns + dnn_feature_columns
    feature_spec = tf.compat.v1.feature_column.make_parse_example_spec(
        feature_columns)
    self._test_complete_flow_helper(linear_feature_columns, dnn_feature_columns,
                                    feature_spec, train_input_fn, eval_input_fn,
                                    predict_input_fn, input_dimension,
                                    label_dimension, batch_size)

  def _test_complete_flow_mix1(self, train_input_fn, eval_input_fn,
                               predict_input_fn, input_dimension,
                               label_dimension, batch_size, fc_impl):
    del fc_impl
    linear_feature_columns = [
        tf.feature_column.numeric_column('x', shape=(input_dimension,))
    ]
    dnn_feature_columns = [
        feature_column.numeric_column('x', shape=(input_dimension,))
    ]
    feature_columns = linear_feature_columns + dnn_feature_columns
    feature_spec = tf.compat.v1.feature_column.make_parse_example_spec(
        feature_columns)
    self._test_complete_flow_helper(linear_feature_columns, dnn_feature_columns,
                                    feature_spec, train_input_fn, eval_input_fn,
                                    predict_input_fn, input_dimension,
                                    label_dimension, batch_size)

  def _test_complete_flow_mix2(self, train_input_fn, eval_input_fn,
                               predict_input_fn, input_dimension,
                               label_dimension, batch_size, fc_impl):
    del fc_impl
    linear_feature_columns = [
        feature_column.numeric_column('x', shape=(input_dimension,))
    ]
    dnn_feature_columns = [
        tf.feature_column.numeric_column('x', shape=(input_dimension,))
    ]
    feature_columns = linear_feature_columns + dnn_feature_columns
    feature_spec = tf.compat.v1.feature_column.make_parse_example_spec(
        feature_columns)
    self._test_complete_flow_helper(linear_feature_columns, dnn_feature_columns,
                                    feature_spec, train_input_fn, eval_input_fn,
                                    predict_input_fn, input_dimension,
                                    label_dimension, batch_size)

  def _test_numpy_input_fn_helper(self, fc_impl, fn_to_run):
    """Tests complete flow with numpy_input_fn."""
    label_dimension = 2
    batch_size = 10
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)
    # learn y = x
    train_input_fn = numpy_io.numpy_input_fn(
        x={'x': data},
        y=data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': data}, y=data, batch_size=batch_size, shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': data}, batch_size=batch_size, shuffle=False)

    fn_to_run(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size,
        fc_impl=fc_impl)

  def test_numpy_input_fn_basic(self, fc_impl):
    self._test_numpy_input_fn_helper(fc_impl, self._test_complete_flow)

  def test_numpy_input_fn_mix1(self, fc_impl):
    self._test_numpy_input_fn_helper(fc_impl, self._test_complete_flow_mix1)

  def test_numpy_input_fn_mix2(self, fc_impl):
    self._test_numpy_input_fn_helper(fc_impl, self._test_complete_flow_mix2)

  def _test_pandas_input_fn_helper(self, fc_impl, fn_to_run):
    """Tests complete flow with pandas_input_fn."""
    if not HAS_PANDAS:
      return
    label_dimension = 1
    batch_size = 10
    data = np.linspace(0., 2., batch_size, dtype=np.float32)
    x = pd.DataFrame({'x': data})
    y = pd.Series(data)
    train_input_fn = pandas_io.pandas_input_fn(
        x=x, y=y, batch_size=batch_size, num_epochs=None, shuffle=True)
    eval_input_fn = pandas_io.pandas_input_fn(
        x=x, y=y, batch_size=batch_size, shuffle=False)
    predict_input_fn = pandas_io.pandas_input_fn(
        x=x, batch_size=batch_size, shuffle=False)

    fn_to_run(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size,
        fc_impl=fc_impl)

  def test_pandas_input_fn_basic(self, fc_impl):
    self._test_pandas_input_fn_helper(fc_impl, self._test_complete_flow)

  def test_pandas_input_fn_mix1(self, fc_impl):
    self._test_pandas_input_fn_helper(fc_impl, self._test_complete_flow_mix1)

  def test_pandas_input_fn_mix2(self, fc_impl):
    self._test_pandas_input_fn_helper(fc_impl, self._test_complete_flow_mix2)

  def _test_input_fn_from_parse_example_helper(self, fc_impl, fn_to_run):
    """Tests complete flow with input_fn constructed from parse_example."""
    label_dimension = 2
    batch_size = 10
    data = np.linspace(0., 2., batch_size * label_dimension, dtype=np.float32)
    data = data.reshape(batch_size, label_dimension)

    serialized_examples = []
    for datum in data:
      example = example_pb2.Example(
          features=feature_pb2.Features(
              feature={
                  'x':
                      feature_pb2.Feature(
                          float_list=feature_pb2.FloatList(value=datum)),
                  'y':
                      feature_pb2.Feature(
                          float_list=feature_pb2.FloatList(value=datum)),
              }))
      serialized_examples.append(example.SerializeToString())

    feature_spec = {
        'x': tf.io.FixedLenFeature([label_dimension], tf.dtypes.float32),
        'y': tf.io.FixedLenFeature([label_dimension], tf.dtypes.float32),
    }

    def _train_input_fn():
      feature_map = tf.compat.v1.io.parse_example(serialized_examples,
                                                  feature_spec)
      features = linear_testing_utils_v1.queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels

    def _eval_input_fn():
      feature_map = tf.compat.v1.io.parse_example(
          tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = linear_testing_utils_v1.queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels

    def _predict_input_fn():
      feature_map = tf.compat.v1.io.parse_example(
          tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = linear_testing_utils_v1.queue_parsed_features(feature_map)
      features.pop('y')
      return features, None

    fn_to_run(
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        input_dimension=label_dimension,
        label_dimension=label_dimension,
        batch_size=batch_size,
        fc_impl=fc_impl)

  def test_input_fn_from_parse_example_basic(self, fc_impl):
    self._test_input_fn_from_parse_example_helper(fc_impl,
                                                  self._test_complete_flow)

  def test_input_fn_from_parse_example_mix1(self, fc_impl):
    self._test_input_fn_from_parse_example_helper(fc_impl,
                                                  self._test_complete_flow_mix1)

  def test_input_fn_from_parse_example_mix2(self, fc_impl):
    self._test_input_fn_from_parse_example_helper(fc_impl,
                                                  self._test_complete_flow_mix2)


# A function to mimic dnn-classifier init reuse same tests.
def _dnn_classifier_fn(hidden_units,
                       feature_columns,
                       model_dir=None,
                       n_classes=2,
                       weight_column=None,
                       label_vocabulary=None,
                       optimizer='Adagrad',
                       config=None,
                       input_layer_partitioner=None):
  return dnn_linear_combined.DNNLinearCombinedClassifier(
      model_dir=model_dir,
      dnn_hidden_units=hidden_units,
      dnn_feature_columns=feature_columns,
      dnn_optimizer=optimizer,
      n_classes=n_classes,
      weight_column=weight_column,
      label_vocabulary=label_vocabulary,
      input_layer_partitioner=input_layer_partitioner,
      config=config)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyClassifierEvaluateTest(
    dnn_testing_utils_v1.BaseDNNClassifierEvaluateTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNClassifierEvaluateTest.__init__(
        self, _dnn_classifier_fn, fc_impl=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyClassifierEvaluateV2Test(
    dnn_testing_utils_v1.BaseDNNClassifierEvaluateTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNClassifierEvaluateTest.__init__(
        self, _dnn_classifier_fn, fc_impl=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyClassifierPredictTest(
    dnn_testing_utils_v1.BaseDNNClassifierPredictTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNClassifierPredictTest.__init__(
        self, _dnn_classifier_fn, fc_impl=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyClassifierPredictV2Test(
    dnn_testing_utils_v1.BaseDNNClassifierPredictTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNClassifierPredictTest.__init__(
        self, _dnn_classifier_fn, fc_impl=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyClassifierTrainTest(
    dnn_testing_utils_v1.BaseDNNClassifierTrainTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNClassifierTrainTest.__init__(
        self, _dnn_classifier_fn, fc_impl=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyClassifierTrainV2Test(
    dnn_testing_utils_v1.BaseDNNClassifierTrainTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNClassifierTrainTest.__init__(
        self, _dnn_classifier_fn, fc_impl=feature_column_v2)


# A function to mimic dnn-regressor init reuse same tests.
def _dnn_regressor_fn(hidden_units,
                      feature_columns,
                      model_dir=None,
                      label_dimension=1,
                      weight_column=None,
                      optimizer='Adagrad',
                      config=None,
                      input_layer_partitioner=None):
  return dnn_linear_combined.DNNLinearCombinedRegressor(
      model_dir=model_dir,
      dnn_hidden_units=hidden_units,
      dnn_feature_columns=feature_columns,
      dnn_optimizer=optimizer,
      label_dimension=label_dimension,
      weight_column=weight_column,
      input_layer_partitioner=input_layer_partitioner,
      config=config)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyRegressorEvaluateTest(
    dnn_testing_utils_v1.BaseDNNRegressorEvaluateTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNRegressorEvaluateTest.__init__(
        self, _dnn_regressor_fn, fc_impl=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyRegressorEvaluateV2Test(
    dnn_testing_utils_v1.BaseDNNRegressorEvaluateTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNRegressorEvaluateTest.__init__(
        self, _dnn_regressor_fn, fc_impl=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyRegressorPredictTest(
    dnn_testing_utils_v1.BaseDNNRegressorPredictTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNRegressorPredictTest.__init__(
        self, _dnn_regressor_fn, fc_impl=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyRegressorPredictV2Test(
    dnn_testing_utils_v1.BaseDNNRegressorPredictTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNRegressorPredictTest.__init__(
        self, _dnn_regressor_fn, fc_impl=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyRegressorTrainTest(dnn_testing_utils_v1.BaseDNNRegressorTrainTest,
                                tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNRegressorTrainTest.__init__(
        self, _dnn_regressor_fn, fc_impl=feature_column)


@test_util.run_v1_only('Tests v1 only symbols')
class DNNOnlyRegressorTrainV2Test(
    dnn_testing_utils_v1.BaseDNNRegressorTrainTest, tf.test.TestCase):

  def __init__(self, methodName='runTest'):  # pylint: disable=invalid-name
    tf.test.TestCase.__init__(self, methodName)
    dnn_testing_utils_v1.BaseDNNRegressorTrainTest.__init__(
        self, _dnn_regressor_fn, fc_impl=feature_column_v2)


@test_util.run_v1_only('Tests v1 only symbols')
@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNLinearCombinedClassifierIntegrationTest(tf.test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      tf.compat.v1.summary.FileWriterCache.clear()
      shutil.rmtree(self._model_dir)

  def _as_label(self, data_in_float):
    return np.rint(data_in_float).astype(np.int64)

  def _test_complete_flow(self, train_input_fn, eval_input_fn, predict_input_fn,
                          input_dimension, n_classes, batch_size, fc_impl):
    linear_feature_columns = [
        fc_impl.numeric_column('x', shape=(input_dimension,))
    ]
    dnn_feature_columns = [
        fc_impl.numeric_column('x', shape=(input_dimension,))
    ]
    feature_columns = linear_feature_columns + dnn_feature_columns
    est = dnn_linear_combined.DNNLinearCombinedClassifier(
        linear_feature_columns=linear_feature_columns,
        dnn_hidden_units=(2, 2),
        dnn_feature_columns=dnn_feature_columns,
        n_classes=n_classes,
        model_dir=self._model_dir)

    # TRAIN
    num_steps = 10
    est.train(train_input_fn, steps=num_steps)

    # EVALUTE
    scores = est.evaluate(eval_input_fn)
    self.assertEqual(num_steps, scores[tf.compat.v1.GraphKeys.GLOBAL_STEP])
    self.assertIn('loss', six.iterkeys(scores))

    # PREDICT
    predicted_proba = np.array([
        x[prediction_keys.PredictionKeys.PROBABILITIES]
        for x in est.predict(predict_input_fn)
    ])
    self.assertAllEqual((batch_size, n_classes), predicted_proba.shape)

    # EXPORT
    feature_spec = tf.compat.v1.feature_column.make_parse_example_spec(
        feature_columns)
    serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
        feature_spec)
    export_dir = est.export_saved_model(tempfile.mkdtemp(),
                                        serving_input_receiver_fn)
    self.assertTrue(tf.compat.v1.gfile.Exists(export_dir))

  def test_numpy_input_fn(self, fc_impl):
    """Tests complete flow with numpy_input_fn."""
    n_classes = 3
    input_dimension = 2
    batch_size = 10
    data = np.linspace(
        0., n_classes - 1., batch_size * input_dimension, dtype=np.float32)
    x_data = data.reshape(batch_size, input_dimension)
    y_data = self._as_label(np.reshape(data[:batch_size], (batch_size, 1)))
    # learn y = x
    train_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data},
        y=y_data,
        batch_size=batch_size,
        num_epochs=None,
        shuffle=True)
    eval_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data}, y=y_data, batch_size=batch_size, shuffle=False)
    predict_input_fn = numpy_io.numpy_input_fn(
        x={'x': x_data}, batch_size=batch_size, shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size,
        fc_impl=fc_impl)

  def test_pandas_input_fn(self, fc_impl):
    """Tests complete flow with pandas_input_fn."""
    if not HAS_PANDAS:
      return
    input_dimension = 1
    n_classes = 2
    batch_size = 10
    data = np.linspace(0., n_classes - 1., batch_size, dtype=np.float32)
    x = pd.DataFrame({'x': data})
    y = pd.Series(self._as_label(data))
    train_input_fn = pandas_io.pandas_input_fn(
        x=x, y=y, batch_size=batch_size, num_epochs=None, shuffle=True)
    eval_input_fn = pandas_io.pandas_input_fn(
        x=x, y=y, batch_size=batch_size, shuffle=False)
    predict_input_fn = pandas_io.pandas_input_fn(
        x=x, batch_size=batch_size, shuffle=False)

    self._test_complete_flow(
        train_input_fn=train_input_fn,
        eval_input_fn=eval_input_fn,
        predict_input_fn=predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size,
        fc_impl=fc_impl)

  def test_input_fn_from_parse_example(self, fc_impl):
    """Tests complete flow with input_fn constructed from parse_example."""
    input_dimension = 2
    n_classes = 3
    batch_size = 10
    data = np.linspace(
        0., n_classes - 1., batch_size * input_dimension, dtype=np.float32)
    data = data.reshape(batch_size, input_dimension)

    serialized_examples = []
    for datum in data:
      example = example_pb2.Example(
          features=feature_pb2.Features(
              feature={
                  'x':
                      feature_pb2.Feature(
                          float_list=feature_pb2.FloatList(value=datum)),
                  'y':
                      feature_pb2.Feature(
                          int64_list=feature_pb2.Int64List(
                              value=self._as_label(datum[:1]))),
              }))
      serialized_examples.append(example.SerializeToString())

    feature_spec = {
        'x': tf.io.FixedLenFeature([input_dimension], tf.dtypes.float32),
        'y': tf.io.FixedLenFeature([1], tf.dtypes.int64),
    }

    def _train_input_fn():
      feature_map = tf.compat.v1.io.parse_example(serialized_examples,
                                                  feature_spec)
      features = linear_testing_utils_v1.queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels

    def _eval_input_fn():
      feature_map = tf.compat.v1.io.parse_example(
          tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = linear_testing_utils_v1.queue_parsed_features(feature_map)
      labels = features.pop('y')
      return features, labels

    def _predict_input_fn():
      feature_map = tf.compat.v1.io.parse_example(
          tf.compat.v1.train.limit_epochs(serialized_examples, num_epochs=1),
          feature_spec)
      features = linear_testing_utils_v1.queue_parsed_features(feature_map)
      features.pop('y')
      return features, None

    self._test_complete_flow(
        train_input_fn=_train_input_fn,
        eval_input_fn=_eval_input_fn,
        predict_input_fn=_predict_input_fn,
        input_dimension=input_dimension,
        n_classes=n_classes,
        batch_size=batch_size,
        fc_impl=fc_impl)


@test_util.run_v1_only('Tests v1 only symbols')
@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNLinearCombinedTests(tf.test.TestCase):

  def setUp(self):
    self._model_dir = tempfile.mkdtemp()

  def tearDown(self):
    if self._model_dir:
      shutil.rmtree(self._model_dir)

  def _mock_optimizer(self, real_optimizer, var_name_prefix):
    """Verifies global_step is None and var_names start with given prefix."""

    def _minimize(loss, global_step=None, var_list=None):
      self.assertIsNone(global_step)
      trainable_vars = var_list or tf.compat.v1.get_collection(
          tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES)
      var_names = [var.name for var in trainable_vars]
      self.assertTrue(
          all([name.startswith(var_name_prefix) for name in var_names]))
      # var is used to check this op called by training.
      with ops.name_scope(''):
        var = tf.Variable(0., name=(var_name_prefix + '_called'))
      with tf.control_dependencies([var.assign(100.)]):
        return real_optimizer.minimize(loss, global_step, var_list)

    optimizer_mock = tf.compat.v1.test.mock.NonCallableMagicMock(
        spec=tf.compat.v1.train.Optimizer, wraps=real_optimizer)
    optimizer_mock.minimize = tf.compat.v1.test.mock.MagicMock(wraps=_minimize)

    return optimizer_mock

  def test_train_op_calls_both_dnn_and_linear(self, fc_impl):
    opt = tf.compat.v1.train.GradientDescentOptimizer(1.)
    x_column = fc_impl.numeric_column('x')
    input_fn = numpy_io.numpy_input_fn(
        x={'x': np.array([[0.], [1.]])},
        y=np.array([[0.], [1.]]),
        batch_size=1,
        shuffle=False)
    est = dnn_linear_combined.DNNLinearCombinedClassifier(
        linear_feature_columns=[x_column],
        # verifies linear_optimizer is used only for linear part.
        linear_optimizer=self._mock_optimizer(opt, 'linear'),
        dnn_hidden_units=(2, 2),
        dnn_feature_columns=[x_column],
        # verifies dnn_optimizer is used only for linear part.
        dnn_optimizer=self._mock_optimizer(opt, 'dnn'),
        model_dir=self._model_dir)
    est.train(input_fn, steps=1)
    # verifies train_op fires linear minimize op
    self.assertEqual(100.,
                     tf.train.load_variable(self._model_dir, 'linear_called'))
    # verifies train_op fires dnn minimize op
    self.assertEqual(100., tf.train.load_variable(self._model_dir,
                                                  'dnn_called'))

  def test_dnn_and_linear_logits_are_added(self, fc_impl):
    with tf.Graph().as_default():
      tf.Variable([[1.0]], name='linear/linear_model/x/weights')
      tf.Variable([2.0], name='linear/linear_model/bias_weights')
      tf.Variable([[3.0]], name='dnn/hiddenlayer_0/kernel')
      tf.Variable([4.0], name='dnn/hiddenlayer_0/bias')
      tf.Variable([[5.0]], name='dnn/logits/kernel')
      tf.Variable([6.0], name='dnn/logits/bias')
      tf.Variable(1, name='global_step', dtype=tf.dtypes.int64)
      linear_testing_utils_v1.save_variables_to_ckpt(self._model_dir)

    x_column = fc_impl.numeric_column('x')
    est = dnn_linear_combined.DNNLinearCombinedRegressor(
        linear_feature_columns=[x_column],
        dnn_hidden_units=[1],
        dnn_feature_columns=[x_column],
        model_dir=self._model_dir)
    input_fn = numpy_io.numpy_input_fn(
        x={'x': np.array([[10.]])}, batch_size=1, shuffle=False)
    # linear logits = 10*1 + 2 = 12
    # dnn logits = (10*3 + 4)*5 + 6 = 176
    # logits = dnn + linear = 176 + 12 = 188
    self.assertAllClose({
        prediction_keys.PredictionKeys.PREDICTIONS: [188.],
    }, next(est.predict(input_fn=input_fn)))


@test_util.run_v1_only('Tests v1 only symbols')
@parameterized.parameters((feature_column,), (feature_column_v2,))
class DNNLinearCombinedWarmStartingTest(tf.test.TestCase):

  def setUp(self):
    # Create a directory to save our old checkpoint and vocabularies to.
    self._ckpt_and_vocab_dir = tempfile.mkdtemp()

    # Make a dummy input_fn.
    def _input_fn():
      features = {
          'age': [[23.], [31.]],
          'city': [['Palo Alto'], ['Mountain View']],
      }
      return features, [0, 1]

    self._input_fn = _input_fn

  def tearDown(self):
    # Clean up checkpoint / vocab dir.
    tf.compat.v1.summary.FileWriterCache.clear()
    shutil.rmtree(self._ckpt_and_vocab_dir)

  def test_classifier_basic_warm_starting(self, fc_impl):
    """Tests correctness of DNNLinearCombinedClassifier default warm-start."""
    age = fc_impl.numeric_column('age')
    city = fc_impl.embedding_column(
        fc_impl.categorical_column_with_vocabulary_list(
            'city', vocabulary_list=['Mountain View', 'Palo Alto']),
        dimension=5)

    # Create a DNNLinearCombinedClassifier and train to save a checkpoint.
    dnn_lc_classifier = dnn_linear_combined.DNNLinearCombinedClassifier(
        linear_feature_columns=[age],
        dnn_feature_columns=[city],
        dnn_hidden_units=[256, 128],
        model_dir=self._ckpt_and_vocab_dir,
        n_classes=4,
        linear_optimizer='SGD',
        dnn_optimizer='SGD')
    dnn_lc_classifier.train(input_fn=self._input_fn, max_steps=1)

    # Create a second DNNLinearCombinedClassifier, warm-started from the first.
    # Use a learning_rate = 0.0 optimizer to check values (use SGD so we don't
    # have accumulator values that change).
    warm_started_dnn_lc_classifier = (
        dnn_linear_combined.DNNLinearCombinedClassifier(
            linear_feature_columns=[age],
            dnn_feature_columns=[city],
            dnn_hidden_units=[256, 128],
            n_classes=4,
            linear_optimizer=tf.compat.v1.train.GradientDescentOptimizer(
                learning_rate=0.0),
            dnn_optimizer=tf.compat.v1.train.GradientDescentOptimizer(
                learning_rate=0.0),
            warm_start_from=dnn_lc_classifier.model_dir))

    warm_started_dnn_lc_classifier.train(input_fn=self._input_fn, max_steps=1)
    for variable_name in warm_started_dnn_lc_classifier.get_variable_names():
      self.assertAllClose(
          dnn_lc_classifier.get_variable_value(variable_name),
          warm_started_dnn_lc_classifier.get_variable_value(variable_name))

  def test_regressor_basic_warm_starting(self, fc_impl):
    """Tests correctness of DNNLinearCombinedRegressor default warm-start."""
    age = fc_impl.numeric_column('age')
    city = fc_impl.embedding_column(
        fc_impl.categorical_column_with_vocabulary_list(
            'city', vocabulary_list=['Mountain View', 'Palo Alto']),
        dimension=5)

    # Create a DNNLinearCombinedRegressor and train to save a checkpoint.
    dnn_lc_regressor = dnn_linear_combined.DNNLinearCombinedRegressor(
        linear_feature_columns=[age],
        dnn_feature_columns=[city],
        dnn_hidden_units=[256, 128],
        model_dir=self._ckpt_and_vocab_dir,
        linear_optimizer='SGD',
        dnn_optimizer='SGD')
    dnn_lc_regressor.train(input_fn=self._input_fn, max_steps=1)

    # Create a second DNNLinearCombinedRegressor, warm-started from the first.
    # Use a learning_rate = 0.0 optimizer to check values (use SGD so we don't
    # have accumulator values that change).
    warm_started_dnn_lc_regressor = (
        dnn_linear_combined.DNNLinearCombinedRegressor(
            linear_feature_columns=[age],
            dnn_feature_columns=[city],
            dnn_hidden_units=[256, 128],
            linear_optimizer=tf.compat.v1.train.GradientDescentOptimizer(
                learning_rate=0.0),
            dnn_optimizer=tf.compat.v1.train.GradientDescentOptimizer(
                learning_rate=0.0),
            warm_start_from=dnn_lc_regressor.model_dir))

    warm_started_dnn_lc_regressor.train(input_fn=self._input_fn, max_steps=1)
    for variable_name in warm_started_dnn_lc_regressor.get_variable_names():
      self.assertAllClose(
          dnn_lc_regressor.get_variable_value(variable_name),
          warm_started_dnn_lc_regressor.get_variable_value(variable_name))

  def test_warm_starting_selective_variables(self, fc_impl):
    """Tests selecting variables to warm-start."""
    age = fc_impl.numeric_column('age')
    city = fc_impl.embedding_column(
        fc_impl.categorical_column_with_vocabulary_list(
            'city', vocabulary_list=['Mountain View', 'Palo Alto']),
        dimension=5)

    # Create a DNNLinearCombinedClassifier and train to save a checkpoint.
    dnn_lc_classifier = dnn_linear_combined.DNNLinearCombinedClassifier(
        linear_feature_columns=[age],
        dnn_feature_columns=[city],
        dnn_hidden_units=[256, 128],
        model_dir=self._ckpt_and_vocab_dir,
        n_classes=4,
        linear_optimizer='SGD',
        dnn_optimizer='SGD')
    dnn_lc_classifier.train(input_fn=self._input_fn, max_steps=1)

    # Create a second DNNLinearCombinedClassifier, warm-started from the first.
    # Use a learning_rate = 0.0 optimizer to check values (use SGD so we don't
    # have accumulator values that change).
    warm_started_dnn_lc_classifier = (
        dnn_linear_combined.DNNLinearCombinedClassifier(
            linear_feature_columns=[age],
            dnn_feature_columns=[city],
            dnn_hidden_units=[256, 128],
            n_classes=4,
            linear_optimizer=tf.compat.v1.train.GradientDescentOptimizer(
                learning_rate=0.0),
            dnn_optimizer=tf.compat.v1.train.GradientDescentOptimizer(
                learning_rate=0.0),
            # The provided regular expression will only warm-start the deep
            # portion of the model.
            warm_start_from=estimator.WarmStartSettings(
                ckpt_to_initialize_from=dnn_lc_classifier.model_dir,
                vars_to_warm_start='.*(dnn).*')))

    warm_started_dnn_lc_classifier.train(input_fn=self._input_fn, max_steps=1)
    for variable_name in warm_started_dnn_lc_classifier.get_variable_names():
      if 'dnn' in variable_name:
        self.assertAllClose(
            dnn_lc_classifier.get_variable_value(variable_name),
            warm_started_dnn_lc_classifier.get_variable_value(variable_name))
      elif 'linear' in variable_name:
        linear_values = warm_started_dnn_lc_classifier.get_variable_value(
            variable_name)
        # Since they're not warm-started, the linear weights will be
        # zero-initialized.
        self.assertAllClose(np.zeros_like(linear_values), linear_values)


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