/* Copyright 2015 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/util/tensor_slice_set.h"

#include <vector>
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"

namespace tensorflow {

namespace checkpoint {

namespace {

// A simple test: we have a 2-d tensor of shape 4 X 5 that looks like this:
//
//   0   1   2   3   4
//   5   6   7   8   9
//  10  11  12  13  14
//  15  16  17  18  19
//
// We assume this is a row-major matrix.
//
// We store the tensor in a couple of slices and verify that we can recover all
// of them.
TEST(TensorSliceSetTest, QueryTwoD) {
  TensorShape shape({4, 5});

  TensorSliceSet tss(shape, DT_FLOAT);
  // We store a few slices.

  // Slice #1 is the top two rows:
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  const float src_1[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  TensorSlice slice_1 = TensorSlice::ParseOrDie("0,2:-");
  TF_CHECK_OK(tss.Register(slice_1, "", src_1));

  // Slice #2 is the bottom left corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //  10  11  12   .   .
  //  15  16  17   .   .
  const float src_2[] = {10, 11, 12, 15, 16, 17};
  TensorSlice slice_2 = TensorSlice::ParseOrDie("2,2:0,3");
  TF_CHECK_OK(tss.Register(slice_2, "", src_2));

  // Slice #3 is the bottom right corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .  18  19
  const float src_3[] = {18, 19};
  TensorSlice slice_3 = TensorSlice::ParseOrDie("3,1:3,2");
  TF_CHECK_OK(tss.Register(slice_3, "", src_3));

  // Notice that we leave a hole in the tensor
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   . (13) (14)
  //   .   .   .   .   .

  // Now we query some of the slices

  // Slice #1 is an exact match
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("0,2:-");
    float expected[] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
    float results[10];
    EXPECT_TRUE(tss.Query(s, results));
    for (int i = 0; i < 10; ++i) {
      EXPECT_EQ(expected[i], results[i]);
    }
  }

  // Slice #2 is a subset match
  //   .   .   .   .   .
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,1:-");
    float expected[] = {5, 6, 7, 8, 9};
    float results[5];
    EXPECT_TRUE(tss.Query(s, results));
    for (int i = 0; i < 5; ++i) {
      EXPECT_EQ(expected[i], results[i]);
    }
  }

  // Slice #3 is a more complicated match: it needs the combination of a couple
  // of slices
  //   .   .   .   .   .
  //   5   6   7   .   .
  //  10  11  12   .   .
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,2:0,3");
    float expected[] = {5, 6, 7, 10, 11, 12};
    float results[6];
    EXPECT_TRUE(tss.Query(s, results));
    for (int i = 0; i < 6; ++i) {
      EXPECT_EQ(expected[i], results[i]);
    }
  }

  // Slice #4 includes the hole and so there is no match
  //   .   .   .   .   .
  //   .   .   7   8   9
  //   .   .  12  13  14
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,2:2,3");
    float results[6];
    EXPECT_FALSE(tss.Query(s, results));
  }
}

// Testing the meta version of the tensor slice set.
TEST(TensorSliceSetTest, QueryMetaTwoD) {
  TensorShape shape({4, 5});

  TensorSliceSet tss(shape, DT_INT32);
  // We store a few slices.

  // Slice #1 is the top two rows:
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  TensorSlice slice_1 = TensorSlice::ParseOrDie("0,2:-");
  TF_CHECK_OK(tss.Register(slice_1, "slice_1", nullptr));

  // Slice #2 is the bottom left corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //  10  11  12   .   .
  //  15  16  17   .   .
  TensorSlice slice_2 = TensorSlice::ParseOrDie("2,2:0,3");
  TF_CHECK_OK(tss.Register(slice_2, "slice_2", nullptr));

  // Slice #3 is the bottom right corner
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   .  18  19
  TensorSlice slice_3 = TensorSlice::ParseOrDie("3,1:3,2");
  TF_CHECK_OK(tss.Register(slice_3, "slice_3", nullptr));

  // Notice that we leave a hole in the tensor
  //   .   .   .   .   .
  //   .   .   .   .   .
  //   .   .   . (13) (14)
  //   .   .   .   .   .

  // Now we query some of the slices

  // Slice #1 is an exact match
  //   0   1   2   3   4
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  // We just need slice_1 for this
  {
    TensorSlice s = TensorSlice::ParseOrDie("0,2:-");
    std::vector<std::pair<TensorSlice, string>> results;
    EXPECT_TRUE(tss.QueryMeta(s, &results));
    EXPECT_EQ(1, results.size());
    EXPECT_EQ("0,2:-", results[0].first.DebugString());
    EXPECT_EQ("slice_1", results[0].second);
  }

  // Slice #2 is a subset match
  //   .   .   .   .   .
  //   5   6   7   8   9
  //   .   .   .   .   .
  //   .   .   .   .   .
  // We just need slice_1 for this
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,1:-");
    std::vector<std::pair<TensorSlice, string>> results;
    EXPECT_TRUE(tss.QueryMeta(s, &results));
    EXPECT_EQ(1, results.size());
    EXPECT_EQ("0,2:-", results[0].first.DebugString());
    EXPECT_EQ("slice_1", results[0].second);
  }

  // Slice #3 is a more complicated match: it needs the combination of a couple
  // of slices
  //   .   .   .   .   .
  //   5   6   7   .   .
  //  10  11  12   .   .
  //   .   .   .   .   .
  // We need both slice_1 and slice_2 for this.
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,2:0,3");
    std::vector<std::pair<TensorSlice, string>> results;
    EXPECT_TRUE(tss.QueryMeta(s, &results));
    EXPECT_EQ(2, results.size());
    EXPECT_EQ("2,2:0,3", results[0].first.DebugString());
    EXPECT_EQ("slice_2", results[0].second);
    EXPECT_EQ("0,2:-", results[1].first.DebugString());
    EXPECT_EQ("slice_1", results[1].second);
  }

  // Slice #4 includes the hole and so there is no match
  //   .   .   .   .   .
  //   .   .   7   8   9
  //   .   .  12  13  14
  //   .   .   .   .   .
  {
    TensorSlice s = TensorSlice::ParseOrDie("1,2:2,3");
    std::vector<std::pair<TensorSlice, string>> results;
    EXPECT_FALSE(tss.QueryMeta(s, &results));
    EXPECT_EQ(0, results.size());
  }
}

static void BM_RegisterOneByOne(int parts) {
  TensorShape shape({parts, 41});
  TensorSliceSet slice_set(shape, DT_INT32);
  for (int i = 0; i < parts; ++i) {
    TensorSlice part({{i, 1}, {0, -1}});
    TF_CHECK_OK(slice_set.Register(part, part.DebugString(), nullptr));
  }
}

BENCHMARK(BM_RegisterOneByOne);

}  // namespace

}  // namespace checkpoint

}  // namespace tensorflow