# Copyright 2019 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 dense_features."""

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

from absl.testing import parameterized
import numpy as np

from tensorflow.python.client import session
from tensorflow.python.eager import backprop
from tensorflow.python.feature_column import feature_column_v2 as fc
from tensorflow.python.feature_column import sequence_feature_column as sfc
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import test_util
from tensorflow.python.keras import combinations
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras.feature_column import dense_features as df
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import partitioned_variables
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.platform import test


def _initialized_session(config=None):
  sess = session.Session(config=config)
  sess.run(variables_lib.global_variables_initializer())
  sess.run(lookup_ops.tables_initializer())
  return sess


class DenseFeaturesTest(keras_parameterized.TestCase):

  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
  def test_retrieving_input(self):
    features = {'a': [0.]}
    dense_features = df.DenseFeatures(fc.numeric_column('a'))
    inputs = self.evaluate(dense_features(features))
    self.assertAllClose([[0.]], inputs)

  @combinations.generate(combinations.combine(mode=['eager']))
  def test_reuses_variables(self):
    sparse_input = sparse_tensor.SparseTensor(
        indices=((0, 0), (1, 0), (2, 0)),
        values=(0, 1, 2),
        dense_shape=(3, 3))

    # Create feature columns (categorical and embedding).
    categorical_column = fc.categorical_column_with_identity(
        key='a', num_buckets=3)
    embedding_dimension = 2

    def _embedding_column_initializer(shape, dtype, partition_info=None):
      del shape  # unused
      del dtype  # unused
      del partition_info  # unused
      embedding_values = (
          (1, 0),  # id 0
          (0, 1),  # id 1
          (1, 1))  # id 2
      return embedding_values

    embedding_column = fc.embedding_column(
        categorical_column,
        dimension=embedding_dimension,
        initializer=_embedding_column_initializer)

    dense_features = df.DenseFeatures([embedding_column])
    features = {'a': sparse_input}

    inputs = dense_features(features)
    variables = dense_features.variables

    # Sanity check: test that the inputs are correct.
    self.assertAllEqual([[1, 0], [0, 1], [1, 1]], inputs)

    # Check that only one variable was created.
    self.assertEqual(1, len(variables))

    # Check that invoking dense_features on the same features does not create
    # additional variables
    _ = dense_features(features)
    self.assertEqual(1, len(variables))
    self.assertIs(variables[0], dense_features.variables[0])

  @combinations.generate(combinations.combine(mode=['eager']))
  def test_dense_feature_with_partitioner(self):
    sparse_input = sparse_tensor.SparseTensor(
        indices=((0, 0), (1, 0), (2, 0), (3, 0)),
        values=(0, 1, 3, 2),
        dense_shape=(4, 4))

    # Create feature columns (categorical and embedding).
    categorical_column = fc.categorical_column_with_identity(
        key='a', num_buckets=4)
    embedding_dimension = 2

    def _embedding_column_initializer(shape, dtype, partition_info=None):
      offset = partition_info._var_offset[0]
      del shape  # unused
      del dtype  # unused
      if offset == 0:
        embedding_values = (
            (1, 0),  # id 0
            (0, 1))  # id 1
      else:
        embedding_values = (
            (1, 1),  # id 2
            (2, 2))  # id 3
      return embedding_values

    embedding_column = fc.embedding_column(
        categorical_column,
        dimension=embedding_dimension,
        initializer=_embedding_column_initializer)

    dense_features = df.DenseFeatures(
        [embedding_column],
        partitioner=partitioned_variables.fixed_size_partitioner(2))
    features = {'a': sparse_input}

    inputs = dense_features(features)
    variables = dense_features.variables

    # Sanity check: test that the inputs are correct.
    self.assertAllEqual([[1, 0], [0, 1], [2, 2], [1, 1]], inputs)

    # Check that only one variable was created.
    self.assertEqual(2, len(variables))

    # Check that invoking dense_features on the same features does not create
    # additional variables
    _ = dense_features(features)
    self.assertEqual(2, len(variables))
    self.assertIs(variables[0], dense_features.variables[0])
    self.assertIs(variables[1], dense_features.variables[1])

  @combinations.generate(combinations.combine(mode=['eager']))
  def test_feature_column_dense_features_gradient(self):
    sparse_input = sparse_tensor.SparseTensor(
        indices=((0, 0), (1, 0), (2, 0)),
        values=(0, 1, 2),
        dense_shape=(3, 3))

    # Create feature columns (categorical and embedding).
    categorical_column = fc.categorical_column_with_identity(
        key='a', num_buckets=3)
    embedding_dimension = 2

    def _embedding_column_initializer(shape, dtype, partition_info=None):
      del shape  # unused
      del dtype  # unused
      del partition_info  # unused
      embedding_values = (
          (1, 0),  # id 0
          (0, 1),  # id 1
          (1, 1))  # id 2
      return embedding_values

    embedding_column = fc.embedding_column(
        categorical_column,
        dimension=embedding_dimension,
        initializer=_embedding_column_initializer)

    dense_features = df.DenseFeatures([embedding_column])
    features = {'a': sparse_input}

    def scale_matrix():
      matrix = dense_features(features)
      return 2 * matrix

    # Sanity check: Verify that scale_matrix returns the correct output.
    self.assertAllEqual([[2, 0], [0, 2], [2, 2]], scale_matrix())

    # Check that the returned gradient is correct.
    grad_function = backprop.implicit_grad(scale_matrix)
    grads_and_vars = grad_function()
    indexed_slice = grads_and_vars[0][0]
    gradient = grads_and_vars[0][0].values

    self.assertAllEqual([0, 1, 2], indexed_slice.indices)
    self.assertAllEqual([[2, 2], [2, 2], [2, 2]], gradient)

  def test_raises_if_empty_feature_columns(self):
    with self.assertRaisesRegex(ValueError,
                                'feature_columns must not be empty'):
      df.DenseFeatures(feature_columns=[])(features={})

  def test_should_be_dense_column(self):
    with self.assertRaisesRegex(ValueError, 'must be a .*DenseColumn'):
      df.DenseFeatures(feature_columns=[
          fc.categorical_column_with_hash_bucket('wire_cast', 4)
      ])(
          features={
              'a': [[0]]
          })

  def test_does_not_support_dict_columns(self):
    with self.assertRaisesRegex(
        ValueError, 'Expected feature_columns to be iterable, found dict.'):
      df.DenseFeatures(feature_columns={'a': fc.numeric_column('a')})(
          features={
              'a': [[0]]
          })

  def test_bare_column(self):
    with ops.Graph().as_default():
      features = features = {'a': [0.]}
      net = df.DenseFeatures(fc.numeric_column('a'))(features)

      self.evaluate(variables_lib.global_variables_initializer())
      self.evaluate(lookup_ops.tables_initializer())

      self.assertAllClose([[0.]], self.evaluate(net))

  def test_column_generator(self):
    with ops.Graph().as_default():
      features = features = {'a': [0.], 'b': [1.]}
      columns = (fc.numeric_column(key) for key in features)
      net = df.DenseFeatures(columns)(features)

      self.evaluate(variables_lib.global_variables_initializer())
      self.evaluate(lookup_ops.tables_initializer())

      self.assertAllClose([[0., 1.]], self.evaluate(net))

  def test_raises_if_duplicate_name(self):
    with self.assertRaisesRegex(
        ValueError, 'Duplicate feature column name found for columns'):
      df.DenseFeatures(
          feature_columns=[fc.numeric_column('a'),
                           fc.numeric_column('a')])(
                               features={
                                   'a': [[0]]
                               })

  def test_one_column(self):
    price = fc.numeric_column('price')
    with ops.Graph().as_default():
      features = {'price': [[1.], [5.]]}
      net = df.DenseFeatures([price])(features)

      self.evaluate(variables_lib.global_variables_initializer())
      self.evaluate(lookup_ops.tables_initializer())

      self.assertAllClose([[1.], [5.]], self.evaluate(net))

  def test_multi_dimension(self):
    price = fc.numeric_column('price', shape=2)
    with ops.Graph().as_default():
      features = {'price': [[1., 2.], [5., 6.]]}
      net = df.DenseFeatures([price])(features)

      self.evaluate(variables_lib.global_variables_initializer())
      self.evaluate(lookup_ops.tables_initializer())

      self.assertAllClose([[1., 2.], [5., 6.]], self.evaluate(net))

  def test_compute_output_shape(self):
    price1 = fc.numeric_column('price1', shape=2)
    price2 = fc.numeric_column('price2', shape=4)
    with ops.Graph().as_default():
      features = {
          'price1': [[1., 2.], [5., 6.]],
          'price2': [[3., 4., 5., 6.], [7., 8., 9., 10.]]
      }
      dense_features = df.DenseFeatures([price1, price2])
      self.assertEqual((None, 6), dense_features.compute_output_shape((None,)))
      net = dense_features(features)

      self.evaluate(variables_lib.global_variables_initializer())
      self.evaluate(lookup_ops.tables_initializer())

      self.assertAllClose([[1., 2., 3., 4., 5., 6.], [5., 6., 7., 8., 9., 10.]],
                          self.evaluate(net))

  def test_raises_if_shape_mismatch(self):
    price = fc.numeric_column('price', shape=2)
    with ops.Graph().as_default():
      features = {'price': [[1.], [5.]]}
      with self.assertRaisesRegex(
          Exception,
          r'Cannot reshape a tensor with 2 elements to shape \[2,2\]'):
        df.DenseFeatures([price])(features)

  def test_reshaping(self):
    price = fc.numeric_column('price', shape=[1, 2])
    with ops.Graph().as_default():
      features = {'price': [[[1., 2.]], [[5., 6.]]]}
      net = df.DenseFeatures([price])(features)

      self.evaluate(variables_lib.global_variables_initializer())
      self.evaluate(lookup_ops.tables_initializer())

      self.assertAllClose([[1., 2.], [5., 6.]], self.evaluate(net))

  def test_multi_column(self):
    price1 = fc.numeric_column('price1', shape=2)
    price2 = fc.numeric_column('price2')
    with ops.Graph().as_default():
      features = {'price1': [[1., 2.], [5., 6.]], 'price2': [[3.], [4.]]}
      net = df.DenseFeatures([price1, price2])(features)

      self.evaluate(variables_lib.global_variables_initializer())
      self.evaluate(lookup_ops.tables_initializer())

      self.assertAllClose([[1., 2., 3.], [5., 6., 4.]], self.evaluate(net))

  def test_cols_to_output_tensors(self):
    price1 = fc.numeric_column('price1', shape=2)
    price2 = fc.numeric_column('price2')
    with ops.Graph().as_default():
      cols_dict = {}
      features = {'price1': [[1., 2.], [5., 6.]], 'price2': [[3.], [4.]]}
      dense_features = df.DenseFeatures([price1, price2])
      net = dense_features(features, cols_dict)

      self.evaluate(variables_lib.global_variables_initializer())
      self.evaluate(lookup_ops.tables_initializer())

      self.assertAllClose([[1., 2.], [5., 6.]],
                          self.evaluate(cols_dict[price1]))
      self.assertAllClose([[3.], [4.]], self.evaluate(cols_dict[price2]))
      self.assertAllClose([[1., 2., 3.], [5., 6., 4.]], self.evaluate(net))

  def test_column_order(self):
    price_a = fc.numeric_column('price_a')
    price_b = fc.numeric_column('price_b')
    with ops.Graph().as_default():
      features = {
          'price_a': [[1.]],