/* 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 <functional>
#include <memory>

#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/kernels/ops_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"

namespace tensorflow {
namespace {

class RollOpTest : public OpsTestBase {
 protected:
  void MakeOp(DataType data_type, DataType index_type) {
    TF_ASSERT_OK(NodeDefBuilder("myop", "Roll")
                     .Input(FakeInput(data_type))
                     .Input(FakeInput(index_type))
                     .Input(FakeInput(index_type))
                     .Finalize(node_def()));
    TF_ASSERT_OK(InitOp());
  }
};

TEST_F(RollOpTest, ScalarIndices) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({5}), {0, 1, 2, 3, 4});
  AddInputFromArray<int32>(TensorShape({}), {3});
  AddInputFromArray<int32>(TensorShape({}), {0});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({5}));
  test::FillValues<float>(&expected, {2, 3, 4, 0, 1});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, ScalarIndices_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT32);

  // Feed and run
  AddInputFromArray<string>(TensorShape({5}), {"a", "b", "c", "d", "e"});
  AddInputFromArray<int32>(TensorShape({}), {3});
  AddInputFromArray<int32>(TensorShape({}), {0});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({5}));
  test::FillValues<string>(&expected, {"c", "d", "e", "a", "b"});
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, ScalarIndices_Complex) {
  MakeOp(DT_COMPLEX64, DT_INT32);

  // Feed and run
  AddInputFromArray<std::complex<float>>(
      TensorShape({5}), {std::complex<float>(0, 10), std::complex<float>(1, 11),
                         std::complex<float>(2, 12), std::complex<float>(3, 13),
                         std::complex<float>(4, 14)});
  AddInputFromArray<int32>(TensorShape({}), {3});
  AddInputFromArray<int32>(TensorShape({}), {0});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_COMPLEX64, TensorShape({5}));
  test::FillValues<std::complex<float>>(
      &expected, {std::complex<float>(2, 12), std::complex<float>(3, 13),
                  std::complex<float>(4, 14), std::complex<float>(0, 10),
                  std::complex<float>(1, 11)});
  test::ExpectTensorEqual<std::complex<float>>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Simple_TwoD32) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({3, 5}),
                           {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14});
  AddInputFromArray<int32>(TensorShape({2}), {2, -1});
  AddInputFromArray<int32>(TensorShape({2}), {0, 1});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({3, 5}));
  test::FillValues<float>(&expected,
                          {6, 7, 8, 9, 5, 11, 12, 13, 14, 10, 1, 2, 3, 4, 0});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Simple_TwoD32_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT32);

  // Feed and run
  AddInputFromArray<string>(TensorShape({3, 5}),
                            {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j",
                             "k", "l", "m", "n", "o"});
  AddInputFromArray<int32>(TensorShape({2}), {2, -1});
  AddInputFromArray<int32>(TensorShape({2}), {0, 1});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({3, 5}));
  test::FillValues<string>(&expected, {"g", "h", "i", "j", "f", "l", "m", "n",
                                       "o", "k", "b", "c", "d", "e", "a"});
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Simple_ThreeD32) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({2, 2, 3}),
                           {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11});
  AddInputFromArray<int32>(TensorShape({3}), {1, -1, -1});
  AddInputFromArray<int32>(TensorShape({3}), {0, 1, 2});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 2, 3}));
  test::FillValues<float>(&expected, {10, 11, 9, 7, 8, 6, 4, 5, 3, 1, 2, 0});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Simple_ThreeD32_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT32);

  // Feed and run
  AddInputFromArray<string>(
      TensorShape({2, 2, 3}),
      {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"});
  AddInputFromArray<int32>(TensorShape({3}), {1, -1, -1});
  AddInputFromArray<int32>(TensorShape({3}), {0, 1, 2});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({2, 2, 3}));
  test::FillValues<string>(
      &expected, {"k", "l", "j", "h", "i", "g", "e", "f", "d", "b", "c", "a"});
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Simple_TwoD64) {
  MakeOp(DT_FLOAT, DT_INT64);

  // Feed and run
  AddInputFromArray<float>(TensorShape({5, 3}),
                           {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14});
  AddInputFromArray<int64>(TensorShape({2}), {-1, 4});
  AddInputFromArray<int64>(TensorShape({2}), {0, 1});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({5, 3}));
  test::FillValues<float>(&expected,
                          {5, 3, 4, 8, 6, 7, 11, 9, 10, 14, 12, 13, 2, 0, 1});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Simple_TwoD64_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT64);

  // Feed and run
  AddInputFromArray<string>(TensorShape({5, 3}),
                            {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j",
                             "k", "l", "m", "n", "o"});
  AddInputFromArray<int64>(TensorShape({2}), {-1, 4});
  AddInputFromArray<int64>(TensorShape({2}), {0, 1});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({5, 3}));
  test::FillValues<string>(&expected, {"f", "d", "e", "i", "g", "h", "l", "j",
                                       "k", "o", "m", "n", "c", "a", "b"});
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Simple_ThreeD64) {
  MakeOp(DT_FLOAT, DT_INT64);

  // Feed and run
  AddInputFromArray<float>(TensorShape({4, 1, 3}),
                           {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11});
  AddInputFromArray<int64>(TensorShape({3}), {4, 3, 2});
  AddInputFromArray<int64>(TensorShape({3}), {0, 1, 2});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({4, 1, 3}));
  test::FillValues<float>(&expected, {1, 2, 0, 4, 5, 3, 7, 8, 6, 10, 11, 9});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Simple_ThreeD64_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT64);

  // Feed and run
  AddInputFromArray<string>(
      TensorShape({4, 1, 3}),
      {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"});
  AddInputFromArray<int64>(TensorShape({3}), {4, 3, 2});
  AddInputFromArray<int64>(TensorShape({3}), {0, 1, 2});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({4, 1, 3}));
  test::FillValues<string>(
      &expected, {"b", "c", "a", "e", "f", "d", "h", "i", "g", "k", "l", "j"});
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, ZeroShift_ThreeD32) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({2, 2, 3}),
                           {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11});
  AddInputFromArray<int32>(TensorShape({3}), {0, 0, 0});
  AddInputFromArray<int32>(TensorShape({3}), {0, 1, 2});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({2, 2, 3}));
  test::FillValues<float>(&expected, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, ZeroShift_ThreeD32_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT32);

  // Feed and run
  AddInputFromArray<string>(
      TensorShape({2, 2, 3}),
      {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"});
  AddInputFromArray<int32>(TensorShape({3}), {0, 0, 0});
  AddInputFromArray<int32>(TensorShape({3}), {0, 1, 2});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({2, 2, 3}));
  test::FillValues<string>(
      &expected, {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"});
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, ZeroSize_ThreeD32) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({5, 0, 0}), {});
  AddInputFromArray<int32>(TensorShape({}), {1});
  AddInputFromArray<int32>(TensorShape({}), {0});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({5, 0, 0}));
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, ZeroSize_ThreeD32_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT32);

  // Feed and run
  AddInputFromArray<string>(TensorShape({5, 0, 0}), {});
  AddInputFromArray<int32>(TensorShape({}), {1});
  AddInputFromArray<int32>(TensorShape({}), {0});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({5, 0, 0}));
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, OneSize_ThreeD32) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({1, 1, 1}), {5});
  AddInputFromArray<int32>(TensorShape({}), {1});
  AddInputFromArray<int32>(TensorShape({}), {0});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({1, 1, 1}));
  test::FillValues<float>(&expected, {5});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, OneSize_ThreeD32_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT32);

  // Feed and run
  AddInputFromArray<string>(TensorShape({1, 1, 1}), {"a"});
  AddInputFromArray<int32>(TensorShape({}), {1});
  AddInputFromArray<int32>(TensorShape({}), {0});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({1, 1, 1}));
  test::FillValues<string>(&expected, {"a"});
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, MultiShifts_TwoD32) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({3, 5}),
                           {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14});
  AddInputFromArray<int32>(TensorShape({4}), {-2, 2, -1, 1});
  AddInputFromArray<int32>(TensorShape({4}), {1, 0, 0, 1});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_FLOAT, TensorShape({3, 5}));
  test::FillValues<float>(&expected,
                          {11, 12, 13, 14, 10, 1, 2, 3, 4, 0, 6, 7, 8, 9, 5});
  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, MultiShifts_TwoD32_NoMemcpy) {
  MakeOp(DT_STRING, DT_INT32);

  // Feed and run
  AddInputFromArray<string>(TensorShape({3, 5}),
                            {"a", "b", "c", "d", "e", "f", "g", "h", "i", "j",
                             "k", "l", "m", "n", "o"});
  AddInputFromArray<int32>(TensorShape({4}), {-2, 2, -1, 1});
  AddInputFromArray<int32>(TensorShape({4}), {1, 0, 0, 1});
  TF_ASSERT_OK(RunOpKernel());

  // Check the output.
  Tensor expected(allocator(), DT_STRING, TensorShape({3, 5}));
  test::FillValues<string>(&expected, {"l", "m", "n", "o", "k", "b", "c", "d",
                                       "e", "a", "g", "h", "i", "j", "f"});
  test::ExpectTensorEqual<string>(expected, *GetOutput(0));
}

TEST_F(RollOpTest, Error_InputMustBeVectorOrHigher) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({}), {7});
  AddInputFromArray<int32>(TensorShape({}), {1});
  AddInputFromArray<int32>(TensorShape({}), {0});
  Status s = RunOpKernel();
  EXPECT_TRUE(
      str_util::StrContains(s.ToString(), "input must be 1-D or higher"))
      << s;
}

TEST_F(RollOpTest, Error_AxisMustBeScalarOrVector) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({2, 2}), {1, 2, 3, 4});
  AddInputFromArray<int32>(TensorShape({}), {1});
  AddInputFromArray<int32>(TensorShape({1, 2}), {0, 1});
  Status s = RunOpKernel();
  EXPECT_TRUE(str_util::StrContains(s.ToString(),
                                    "axis must be a scalar or a 1-D vector"))
      << s;
}

TEST_F(RollOpTest, Error_ShiftMustBeScalarOrVector) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({2, 2}), {1, 2, 3, 4});
  AddInputFromArray<int32>(TensorShape({1, 2}), {0, 1});
  AddInputFromArray<int32>(TensorShape({}), {1});
  Status s = RunOpKernel();
  EXPECT_TRUE(str_util::StrContains(s.ToString(),
                                    "shift must be a scalar or a 1-D vector"))
      << s;
}

TEST_F(RollOpTest, Error_ShiftAndAxisMustBeSameSize) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({2, 2}), {1, 2, 3, 4});
  AddInputFromArray<int32>(TensorShape({1}), {1});
  AddInputFromArray<int32>(TensorShape({2}), {0, 1});
  Status s = RunOpKernel();
  EXPECT_TRUE(str_util::StrContains(s.ToString(),
                                    "shift and axis must have the same size"))
      << s;
}

TEST_F(RollOpTest, Error_AxisOutOfRange) {
  MakeOp(DT_FLOAT, DT_INT32);

  // Feed and run
  AddInputFromArray<float>(TensorShape({4}), {1, 2, 3, 4});
  AddInputFromArray<int32>(TensorShape({}), {1});
  AddInputFromArray<int32>(TensorShape({}), {1});
  Status s = RunOpKernel();
  EXPECT_TRUE(str_util::StrContains(s.ToString(), "is out of range")) << s;
}

// isd - (inner shift dimension) The inner most dimension to be shifted.
//    All outer dimensions will also be shifted for testing.
static Graph* RollGraph(const TensorShape& shape, int isd) {
  Graph* g = new Graph(OpRegistry::Global());
  Tensor input(DT_FLOAT, shape);
  input.flat<float>().setRandom();
  const int dims = static_cast<int>(input.dims());
  Tensor shift(DT_INT32, TensorShape({dims}));
  for (int i = 0; i < dims; i++) {
    // shift the inner shift dimension and all outer dimensions
    shift.flat<int32>()(i) = (i <= isd) ? 2 : 0;
  }
  Tensor axis(DT_INT32, TensorShape({dims}));
  for (int i = 0; i < dims; i++) {
    axis.flat<int32>()(i) = i;
  }
  test::graph::Roll(g, test::graph::Constant(g, input),
                    test::graph::Constant(g, shift),
                    test::graph::Constant(g, axis));
  return g;
}

#define BM_ROLL_OUTER(DEVICE)                                                 \
  static void BM_##DEVICE##_roll_outer(int iters, int rows, int columns) {    \
    TensorShape shape{rows, columns};                                         \
    const int64 num_items = static_cast<int64>(iters) * shape.num_elements(); \
    testing::ItemsProcessed(num_items);                                       \
    testing::BytesProcessed(num_items * sizeof(float));                       \
    testing::UseRealTime();                                                   \
    test::Benchmark(#DEVICE, RollGraph(shape, 0)).Run(iters);                 \
  }                                                                           \
  BENCHMARK(BM_##DEVICE##_roll_outer)                                         \
      ->ArgPair(256, 256)                                                     \
      ->ArgPair(512, 512)                                                     \
      ->ArgPair(1024, 1024)                                                   \
      ->ArgPair(2048, 2048)

#define BM_ROLL_ALL(DEVICE)                                                   \
  static void BM_##DEVICE##_roll_all(int iters, int rows, int columns) {      \
    TensorShape shape{rows, columns};                                         \
    const int64 num_items = static_cast<int64>(iters) * shape.num_elements(); \
    testing::ItemsProcessed(num_items);                                       \
    testing::BytesProcessed(num_items * sizeof(float));                       \
    testing::UseRealTime();                                                   \
    test::Benchmark(#DEVICE, RollGraph(shape, 1)).Run(iters);                 \
  }                                                                           \
  BENCHMARK(BM_##DEVICE##_roll_all)                                           \
      ->ArgPair(256, 256)                                                     \
      ->ArgPair(512, 512)                                                     \
      ->ArgPair(1024, 1024)                                                   \
      ->ArgPair(2048, 2048)

BM_ROLL_OUTER(cpu);
BM_ROLL_ALL(cpu);
}  // namespace
}  // namespace tensorflow