# 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.
# ==============================================================================
"""Test configs for strided_slice operators."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import tensorflow.compat.v1 as tf
from tensorflow.lite.testing.zip_test_utils import create_tensor_data
from tensorflow.lite.testing.zip_test_utils import make_zip_of_tests
from tensorflow.lite.testing.zip_test_utils import register_make_test_function
from tensorflow.lite.testing.zip_test_utils import TF_TYPE_INFO


def _make_strided_slice_tests(options, test_parameters, expected_tf_failures=0):
  """Utility function to make strided_slice_tests based on parameters."""

  def build_graph(parameters):
    """Build graph for stride_slice test."""
    input_tensor = tf.compat.v1.placeholder(
        dtype=parameters["dtype"],
        name="input",
        shape=parameters["input_shape"])
    if parameters["constant_indices"]:
      begin = parameters["begin"]
      end = parameters["end"]
      strides = parameters["strides"]
      tensors = [input_tensor]
    else:
      begin = tf.compat.v1.placeholder(
          dtype=parameters["index_type"],
          name="begin",
          shape=[len(parameters["begin"])])
      end = tf.compat.v1.placeholder(
          dtype=parameters["index_type"],
          name="end",
          shape=[len(parameters["end"])])
      strides = None
      if parameters["strides"] is not None:
        strides = tf.compat.v1.placeholder(
            dtype=parameters["index_type"],
            name="strides",
            shape=[len(parameters["strides"])])
      tensors = [input_tensor, begin, end]
      if strides is not None:
        tensors.append(strides)
    out = tf.strided_slice(
        input_tensor,
        begin,
        end,
        strides,
        begin_mask=parameters["begin_mask"],
        end_mask=parameters["end_mask"])
    return tensors, [out]

  def build_inputs(parameters, sess, inputs, outputs):
    """Build inputs for stride_slice test."""
    input_values = create_tensor_data(
        parameters["dtype"],
        parameters["input_shape"],
        min_value=-1,
        max_value=1)
    index_type = TF_TYPE_INFO[parameters["index_type"]][0]
    values = [input_values]
    if not parameters["constant_indices"]:
      begin_values = np.array(parameters["begin"]).astype(index_type)
      end_values = np.array(parameters["end"]).astype(index_type)
      stride_values = (
          np.array(parameters["strides"]).astype(index_type)
          if parameters["strides"] is not None else None)
      values.append(begin_values)
      values.append(end_values)
      if stride_values is not None:
        values.append(stride_values)

    return values, sess.run(outputs, feed_dict=dict(zip(inputs, values)))

  make_zip_of_tests(
      options,
      test_parameters,
      build_graph,
      build_inputs,
      expected_tf_failures=expected_tf_failures)


@register_make_test_function()
def make_strided_slice_tests(options):
  """Make a set of tests to do strided_slice."""

  # TODO(soroosh): add test/support for uint8.
  test_parameters = [
      # 4-D (basic cases with const/non-const indices).
      {
          "dtype": [tf.float32, tf.int32, tf.int64, tf.bool],
          "index_type": [tf.int32],
          "input_shape": [[12, 2, 2, 5]],
          "strides": [None, [2, 1, 3, 1]],
          "begin": [[0, 0, 0, 0]],
          "end": [[12, 2, 2, 5]],
          "begin_mask": [None],
          "end_mask": [None],
          "shrink_axis_mask": [None],
          "constant_indices": [False, True],
          "fully_quantize": [False],
      },
      # 4-D with non-trivial begin & end.
      {
          "dtype": [tf.float32],
          "index_type": [tf.int32],
          "input_shape": [[12, 2, 2, 5]],
          "begin": [[0, 0, 0, 0], [1, 0, 1, 0]],
          "end": [[8, 2, 2, 3], [12, 2, 2, 5]],
          "strides": [None, [2, 1, 3, 1]],
          "begin_mask": [None, 8],
          "end_mask": [None, 3],
          "shrink_axis_mask": [None, 15, -1],
          "constant_indices": [True],
          "fully_quantize": [False],
      },
      # Begin, end, strides dim are different from input shape
      {
          "dtype": [tf.float32],
          "index_type": [tf.int32],
          "input_shape": [[12, 2, 2, 5]],
          "begin": [[0]],
          "end": [[1]],
          "strides": [None, [1]],
          "begin_mask": [0],
          "end_mask": [0],
          "shrink_axis_mask": [1],
          "constant_indices": [True, False],
          "fully_quantize": [False],
      },
      # 2-D
      {
          "dtype": [tf.float32],
          "index_type": [tf.int32],
          "input_shape": [[2, 3]],
          "begin": [[0, 0]],
          "end": [[2, 2]],
          "strides": [None, [2, 2]],
          "begin_mask": [None, 1, 2],
          "end_mask": [None, 1, 2],
          "shrink_axis_mask": [None, 1, 2, 3, -1],
          "constant_indices": [False, True],
          "fully_quantize": [False],
      },
      # Negative strides
      {
          "dtype": [tf.float32],
          "index_type": [tf.int32],
          "input_shape": [[2, 3]],
          "begin": [[0, -1]],
          "end": [[2, -3]],
          "strides": [[1, -1]],
          "begin_mask": [None, 1, 2],
          "end_mask": [None, 1, 2],
          "shrink_axis_mask": [None, 1, 2, 3, -1],
          "constant_indices": [False],
          "fully_quantize": [False],
      },
      # 4-D (cases with const indices and batchsize of 1).
      {
          "dtype": [tf.float32],
          "index_type": [tf.int32],
          "input_shape": [[1, 2, 2, 5]],
          "strides": [None, [1, 1, 1, 1]],
          "begin": [[0, 0, 0, 0], [0, 1, 1, 3]],
          "end": [[1, 2, 2, 5], [1, 2, 2, 4]],
          "begin_mask": [None],
          "end_mask": [None],
          "shrink_axis_mask": [None],
          "constant_indices": [True],
          "fully_quantize": [True],
      },
      # Begin, end, strides dim are different from input shape
      {
          "dtype": [tf.float32],
          "index_type": [tf.int32],
          "input_shape": [[12, 2, 2, 5]],
          "begin": [[0]],
          "end": [[1]],
          "strides": [None, [1]],
          "begin_mask": [0],
          "end_mask": [0],
          "shrink_axis_mask": [1],
          "constant_indices": [True],
          "fully_quantize": [True],
      },
  ]

  if options.use_experimental_converter:
    test_parameters = test_parameters + [
        # Begin equal to input dim.
        {
            "dtype": [tf.float32],
            "index_type": [tf.int32],
            "input_shape": [[1, 1, 2]],
            "begin": [[1]],
            "end": [[0]],
            "strides": [[1]],
            "begin_mask": [0],
            "end_mask": [1],
            "shrink_axis_mask": [0],
            "constant_indices": [True, False],
            "fully_quantize": [False],
        },
        {
            "dtype": [tf.float32],
            "index_type": [tf.int32],
            "input_shape": [[1, 1, 2]],
            "begin": [[1, 0, 0]],
            "end": [[0, -1, -1]],
            "strides": [[1, 1, 1]],
            "begin_mask": [6],
            "end_mask": [7],
            "shrink_axis_mask": [0],
            "constant_indices": [True, False],
            "fully_quantize": [False],
        },
        # String input.
        {
            "dtype": [tf.string],
            "index_type": [tf.int32],
            "input_shape": [[12, 2, 2, 5]],
            "begin": [[0, 0, 0, 0]],
            "end": [[8, 2, 2, 3]],
            "strides": [[2, 1, 3, 1]],
            "begin_mask": [8],
            "end_mask": [3],
            "shrink_axis_mask": [None, -1],
            "constant_indices": [True, False],
            "fully_quantize": [False],
        }
    ]
  _make_strided_slice_tests(options, test_parameters, expected_tf_failures=2)


@register_make_test_function()
def make_strided_slice_1d_exhaustive_tests(options):
  """Make a set of exhaustive tests for 1D strided_slice."""
  test_parameters = [
      # 1-D Exhaustive
      {
          "dtype": [tf.float32],
          "index_type": [tf.int32],
          "input_shape": [[3]],
          "begin": [[-2], [-1], [0], [1], [2]],
          "end": [[-2], [-1], [0], [1], [2]],
          "strides": [[-2], [-1], [1], [2]],
          "begin_mask": [0, 1],
          "end_mask": [0, 1],
          "shrink_axis_mask": [0],
          "constant_indices": [False],
      },
  ]
  _make_strided_slice_tests(options, test_parameters)
