/* Copyright 2018 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.
==============================================================================*/
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/shape_inference_testutil.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/platform/test.h"

namespace tensorflow {
namespace {

class RaggedGatherOpTest : public ::tensorflow::OpsTestBase {
 protected:
  // Builds the tensorflow test graph for RaggedGather.
  template <typename VALUE_TYPE, typename INDEX_TYPE>
  void BuildRaggedGatherGraph(
      const TensorShape& indices_shape, const std::vector<INDEX_TYPE>& indices,
      const std::vector<std::vector<int64>>& params_nested_splits,
      const TensorShape& params_dense_values_shape,
      const gtl::ArraySlice<VALUE_TYPE> params_dense_values) {
    const auto& value_dtype = DataTypeToEnum<VALUE_TYPE>::v();
    const auto& index_dtype = DataTypeToEnum<INDEX_TYPE>::v();
    int64 PARAMS_RAGGED_RANK = params_nested_splits.size();
    int64 num_splits = PARAMS_RAGGED_RANK + indices_shape.dims() - 1;
    TF_ASSERT_OK(
        NodeDefBuilder("tested_op", "RaggedGather")
            .Input(FakeInput(PARAMS_RAGGED_RANK))  // params_nested_splits
            .Input(FakeInput(value_dtype))         // params_dense_values
            .Input(FakeInput(index_dtype))         // indices
            .Attr("PARAMS_RAGGED_RANK", PARAMS_RAGGED_RANK)
            .Attr("OUTPUT_RAGGED_RANK", num_splits)
            .Attr("Tvalues", value_dtype)
            .Attr("Tindices", index_dtype)
            .Finalize(node_def()));
    TF_ASSERT_OK(InitOp());
    for (const auto& splits : params_nested_splits) {
      int64 splits_size = splits.size();
      AddInputFromArray<int64>(TensorShape({splits_size}), splits);
    }
    AddInputFromArray<VALUE_TYPE>(params_dense_values_shape,
                                  params_dense_values);
    AddInputFromArray<INDEX_TYPE>(indices_shape, indices);
  }
};

TEST_F(RaggedGatherOpTest, RaggedGather) {
  // indices = [2, 1, 0, 3]
  // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
  // params.shape = [4, None]
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({4}),                     // indices.shape
      {2, 1, 0, 3},                         // indices
      {{0, 3, 3, 7, 9}},                    // params_nested_splits
      TensorShape({9}),                     // params_dense_values.shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9}  // params_dense_values
  );

  TF_ASSERT_OK(RunOpKernel());

  // Expected: [[.4, .5, .6, .7], [.1, .2, .3], [], [.8, .9]]
  test::ExpectTensorEqual<int64>(*GetOutput(0),
                                 test::AsTensor<int64>({0, 4, 4, 7, 9}));
  test::ExpectTensorNear<float>(
      *GetOutput(1),
      test::AsTensor<float>({.4, .5, .6, .7, .1, .2, .3, .8, .9}), 0.1);
}

TEST_F(RaggedGatherOpTest, RaggedGather_3DParams) {
  // indices = [2, 1, 0, 2, 3]
  // params = [[[]], [[.1, 2], [.3]], [], [[.4, .5], [.6, .7, .8]], [[.9]]]
  // params.shape = [5, None, None]
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({5}),                             // indices.shape
      {2, 1, 0, 2, 3},                              // indices
      {{0, 1, 3, 3, 5, 6}, {0, 0, 2, 3, 5, 8, 9}},  // params_nested_splits
      TensorShape({9}),                             // params_dense_values.shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9}          // params_dense_values
  );

  TF_ASSERT_OK(RunOpKernel());

  // Expected: [[], [[.1, 2], [.3]], [[]], [], [[.4, .5], [.6, .7, .8]]]
  test::ExpectTensorEqual<int64>(*GetOutput(0),
                                 test::AsTensor<int64>({0, 0, 2, 3, 3, 5}));
  test::ExpectTensorEqual<int64>(*GetOutput(1),
                                 test::AsTensor<int64>({0, 2, 3, 3, 5, 8}));
  test::ExpectTensorNear<float>(
      *GetOutput(2), test::AsTensor<float>({.1, .2, .3, .4, .5, .6, .7, .8}),
      0.1);
}

TEST_F(RaggedGatherOpTest, RaggedGather_4DParams) {
  // indices = [2, 1, 0, 2]
  // params = [[[]], [[[1, 2], [3, 4], [5, 6]], [[7, 8]]], []]
  // params.shape = [4, None, None, 2]
  BuildRaggedGatherGraph<int32, int32>(
      TensorShape({4}),              // indices.shape
      {2, 1, 0, 2},                  // indices
      {{0, 1, 3, 3}, {0, 0, 3, 4}},  // params_nested_splits
      TensorShape({4, 2}),           // params_dense_values.shape
      {1, 2, 3, 4, 5, 6, 7, 8}       // params_dense_values
  );

  TF_ASSERT_OK(RunOpKernel());

  // Expected: [[],
  //            [[[1, 2], [3, 4], [5, 6]], [[7, 8]]],
  //            [[]],
  //            []]
  test::ExpectTensorEqual<int64>(*GetOutput(0),
                                 test::AsTensor<int64>({0, 0, 2, 3, 3}));
  test::ExpectTensorEqual<int64>(*GetOutput(1),
                                 test::AsTensor<int64>({0, 3, 4, 4}));
  test::ExpectTensorEqual<int32>(
      *GetOutput(2),
      test::AsTensor<int32>({1, 2, 3, 4, 5, 6, 7, 8}, TensorShape({4, 2})));
}

TEST_F(RaggedGatherOpTest, RaggedGather_2DIndices) {
  // indices = [[2, 1], [0, 3]]
  // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({2, 2}),                  // indices.shape
      {2, 1, 0, 3},                         // indices
      {{0, 3, 3, 7, 9}},                    // params_nested_splits
      TensorShape({9}),                     // params_dense_values.shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9}  // params_dense_values
  );

  TF_ASSERT_OK(RunOpKernel());

  // Expected: [ [ [.4, .5, .6, .7], [.1, .2, .3] ],
  //             [ [],               [.8, .9]     ] ]
  test::ExpectTensorEqual<int64>(*GetOutput(0),
                                 test::AsTensor<int64>({0, 2, 4}));
  test::ExpectTensorEqual<int64>(*GetOutput(1),
                                 test::AsTensor<int64>({0, 4, 4, 7, 9}));
  test::ExpectTensorNear<float>(
      *GetOutput(2),
      test::AsTensor<float>({.4, .5, .6, .7, .1, .2, .3, .8, .9}), 0.1);
}

TEST_F(RaggedGatherOpTest, RaggedGather_ScalarIndices) {
  // indices = 2
  // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({}),                      // indices.shape
      {2},                                  // indices
      {{0, 3, 3, 7, 9}},                    // params_nested_splits
      TensorShape({9}),                     // params_dense_values.shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9}  // params_dense_values
  );
  TF_ASSERT_OK(RunOpKernel());

  // Expected: [.4, .5, .6, .7]
  test::ExpectTensorNear<float>(*GetOutput(0),
                                test::AsTensor<float>({.4, .5, .6, .7}), 0.1);
}

TEST_F(RaggedGatherOpTest, RaggedGather_OutOfBounds) {
  // indices = [2, 10]
  // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({2}),                     // indices.shape
      {2, 10},                              // indices
      {{0, 3, 3, 7, 9}},                    // params_nested_splits
      TensorShape({9}),                     // params_dense_values.shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9}  // params_dense_values
  );
  EXPECT_EQ("indices[1] = 10 is not in [0, 4)", RunOpKernel().error_message());
}

TEST_F(RaggedGatherOpTest, InvalidSplitsNotSorted) {
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({2}),                     // indices.shape
      {0, 2},                               // indices
      {{0, 3, 5, 2, 9}},                    // params_nested_splits
      TensorShape({9}),                     // params_dense_values.shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9}  // params_dense_values
  );
  EXPECT_EQ("Ragged splits must be sorted", RunOpKernel().error_message());
}

TEST_F(RaggedGatherOpTest, InvalidSplitsNegative) {
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({2}),                     // indices.shape
      {0, 2},                               // indices
      {{-1, 3, 2, 7, 9}},                   // params_nested_splits
      TensorShape({9}),                     // params_dense_values.shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9}  // params_dense_values
  );
  EXPECT_EQ("Ragged splits must be non-negative",
            RunOpKernel().error_message());
}

TEST_F(RaggedGatherOpTest, InvalidSplitsEmpty) {
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({0}),  // indices.shape
      {},                // indices
      {{}},              // params_nested_splits
      TensorShape({0}),  // params_dense_values.shape
      {}                 // params_dense_values
  );
  EXPECT_EQ("Ragged splits may not be empty", RunOpKernel().error_message());
}

TEST_F(RaggedGatherOpTest, InvalidSplitsTooBig) {
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({2}),                     // indices.shape
      {0, 2},                               // indices
      {{0, 20, 40, 80, 100}},               // params_nested_splits
      TensorShape({9}),                     // params_dense_values.shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9}  // params_dense_values
  );
  EXPECT_EQ("Ragged splits must not point past values",
            RunOpKernel().error_message());
}

TEST_F(RaggedGatherOpTest, BadValuesShape) {
  BuildRaggedGatherGraph<float, int32>(
      TensorShape({0}),  // indices.shape
      {},                // indices
      {{0}},             // params_nested_splits
      TensorShape({}),   // params_dense_values.shape
      {.1}               // params_dense_values
  );
  EXPECT_EQ("params.rank must be nonzero", RunOpKernel().error_message());
}

TEST_F(RaggedGatherOpTest, ShapeFn) {
  // RaggedGather(param_splits+, param_values, indices) -> [splits+, values]
  ShapeInferenceTestOp op("RaggedGather");

  (*op.node_def.mutable_attr())["PARAMS_RAGGED_RANK"].set_i(1);
  (*op.node_def.mutable_attr())["OUTPUT_RAGGED_RANK"].set_i(1);
  INFER_OK(op, "?;?;?", "[?];?");
  INFER_OK(op, "[?];[?];[?]", "[?];[?]");
  INFER_OK(op, "[?];[?,?,?];[?]", "[?];[?,d1_1,d1_2]");
  INFER_OK(op, "[5];[10];[15]", "[?];[?]");
  INFER_OK(op, "[5];[10,2];[15]", "[?];[?,d1_1]");
  INFER_ERROR("Shape must be rank 1 but is rank 0", op, "[5];[];[]");
  INFER_ERROR("Shape must be rank 1 but is rank 2", op, "[1,2];[];[5]");

  (*op.node_def.mutable_attr())["PARAMS_RAGGED_RANK"].set_i(2);
  (*op.node_def.mutable_attr())["OUTPUT_RAGGED_RANK"].set_i(2);
  INFER_OK(op, "?;?;?;?", "[?];[?];?");
  INFER_OK(op, "[?];[?];[?];[?]", "[?];[?];[?]");
  INFER_OK(op, "[?];[?];[?,?,?];[?]", "[?];[?];[?,d2_1,d2_2]");
  INFER_OK(op, "[5];[10];[15];[20]", "[?];[?];[?]");

  (*op.node_def.mutable_attr())["PARAMS_RAGGED_RANK"].set_i(1);
  (*op.node_def.mutable_attr())["OUTPUT_RAGGED_RANK"].set_i(2);
  INFER_OK(op, "?;?;?", "[?];[?];?");
  INFER_OK(op, "[?];[?];[?,?]", "[?];[?];[?]");
  INFER_OK(op, "[?];[?,?,?];[?,?]", "[?];[?];[?,d1_1,d1_2]");
  INFER_OK(op, "[15];[20];[5,10]", "[?];[?];[?]");
  INFER_OK(op, "[15];[20,2];[5,10]", "[?];[?];[?,d1_1]");

  (*op.node_def.mutable_attr())["PARAMS_RAGGED_RANK"].set_i(1);
  (*op.node_def.mutable_attr())["OUTPUT_RAGGED_RANK"].set_i(0);
  INFER_OK(op, "[?];[?];[]", "[?]");
}

}  // namespace
}  // namespace tensorflow