/**
 * Copyright 2019 Huawei Technologies Co., Ltd
 *
 * 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 <iostream>
#include <memory>
#include <vector>

#include "common/common.h"
#include "dataset/core/client.h"
#include "dataset/core/tensor.h"
#include "dataset/engine/datasetops/source/image_folder_op.h"
#include "dataset/kernels/image/decode_op.h"
#include "dataset/kernels/image/resize_op.h"
#include "dataset/kernels/tensor_op.h"
#include "utils/log_adapter.h"

using namespace mindspore::dataset;
using mindspore::LogStream;
using mindspore::MsLogLevel::INFO;

namespace mindspore {
namespace dataset {
namespace test {
class NoOp : public TensorOp {
 public:
    NoOp() {};

    ~NoOp() {};

    Status Compute(const std::shared_ptr<Tensor> &input, std::shared_ptr<Tensor> *output) override {
      *output = std::move(input);
      return Status::OK();
    };

    void Print(std::ostream &out) const override { out << "NoOp"; };
};

class ThreeToOneOp : public TensorOp {
 public:
    ThreeToOneOp() {};

    ~ThreeToOneOp() {};

    uint32_t NumInput() override { return 3; }
    // Compute function that holds the actual implementation of the operation.
    Status Compute(const std::vector<std::shared_ptr<Tensor>> &input,
                   std::vector<std::shared_ptr<Tensor>> *output) override {
      output->push_back(input[0]);
      return Status::OK();
    };

    void Print(std::ostream &out) const override { out << "ThreeToOneOp"; };
};

class OneToThreeOp : public TensorOp {
 public:
    OneToThreeOp() {};

    ~OneToThreeOp() {};

  uint32_t NumOutput() override { return 3; }

    // Compute function that holds the actual implementation of the operation.
    // Simply pushing the same shared pointer of the first element of input vector three times.
    Status Compute(const std::vector<std::shared_ptr<Tensor>> &input,
                   std::vector<std::shared_ptr<Tensor>> *output) override {
      output->push_back(input[0]);
      output->push_back(input[0]);
      output->push_back(input[0]);
      return Status::OK();
    };

    void Print(std::ostream &out) const override { out << "OneToThreeOp"; };
};
}  // namespace test
}  // namespace dataset
}  // namespace mindspore


class MindDataTestMapOp : public UT::DatasetOpTesting {
 public:
    void SetUp() override {
      DatasetOpTesting::SetUp();
      dataset_path_ = datasets_root_path_ + "" + "/testDataset2";

      GlobalInit();

      // Start with an empty execution tree
      my_tree_ = std::make_shared<ExecutionTree>();
    }

    std::shared_ptr<StorageOp> CreateStorageOp() {
      std::shared_ptr<StorageOp> my_storage_op;
      StorageOp::Builder builder;
      builder.SetDatasetFilesDir(dataset_path_)
          .SetColumnsToLoad({"image", "label", "A", "B"})
          .SetRowsPerBuffer(2)
          .SetWorkerConnectorSize(2)
          .SetNumWorkers(2);
      Status rc = builder.Build(&my_storage_op);
      EXPECT_TRUE(rc.IsOk());
      return my_storage_op;
    }

    std::shared_ptr<ExecutionTree> my_tree_;
 private:
    std::string dataset_path_;
};

std::shared_ptr<ImageFolderOp> ImageFolder(int64_t num_works, int64_t rows, int64_t conns, std::string path,
                                           bool shuf = false, std::shared_ptr<Sampler> sampler = nullptr,
                                           std::map<std::string, int32_t> map = {}, bool decode = false);

std::shared_ptr<ExecutionTree> Build(std::vector<std::shared_ptr<DatasetOp>> ops);

// TestByPosition scenario:
//    StorageOp reads a dataset that have column ordering |image|label|A|B|.
//    A TensorOp that does nothing picks the label column and output a column also named label.
//    Thus, based on the new MapOp behaviour, the column ordering will be |image|label|A|B|.
//    Verify the column ordering based on the Tensor properties matching to that of in the schema file.
TEST_F(MindDataTestMapOp, TestByPosition) {
  Status rc;
  MS_LOG(INFO) << "Doing TestByPosition.";

  // Note: The above storage config yields 5 buffers, each with 2 rows, for a total
  // of 10 rows.
  auto my_storage_op = this->CreateStorageOp();
  rc = my_tree_->AssociateNode(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  auto my_no_op = std::make_shared<mindspore::dataset::test::NoOp>();
  std::vector<std::shared_ptr<TensorOp>> my_func_list;
  my_func_list.push_back(my_no_op);
  std::shared_ptr<MapOp> my_map_op;
  MapOp::Builder builder;
  builder.SetInColNames({"label"})
      .SetOutColNames({})
      .SetTensorFuncs(std::move(my_func_list))
      .SetNumWorkers(100);
  rc = builder.Build(&my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_map_op->AddChild(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssignRoot(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());


  // Based on the schema file, create the golden result to compare with.
  std::vector<DataType::Type> golden_types({
                                              DataType::Type::DE_UINT8,
                                              DataType::Type::DE_INT64,
                                              DataType::Type::DE_FLOAT32,
                                              DataType::Type::DE_INT64}
  );

  std::vector<uint64_t> golden_ranks({3, 1, 4, 1});

  std::vector<TensorShape> golden_shapes({
                                            TensorShape({3, 4, 2}),
                                            TensorShape({7}),
                                            TensorShape({1, 13, 14, 12}),
                                            TensorShape({9})}
  );

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorRow tensor_list;
  rc = di.FetchNextTensorRow(&tensor_list);
  EXPECT_TRUE(rc.IsOk());
  EXPECT_EQ(tensor_list.size(), 4);
  for (uint32_t i = 0; i < tensor_list.size(); i++) {
    EXPECT_EQ(tensor_list[i]->type(), golden_types[i]);
    EXPECT_EQ(tensor_list[i]->Rank(), golden_ranks[i]);
    EXPECT_EQ(tensor_list[i]->shape(), golden_shapes[i]);
    EXPECT_NE(tensor_list[i]->GetBuffer(), nullptr);
  }
}

// TestAsMap scenario:
//    StorageOp reads a dataset that have column ordering |image|label|A|B|.
//    A TensorOp that does nothing picks the "image" column and produces a column named "X".
//    Thus, based on the new MapOp behaviour, the column ordering will be |X|label|A|B|.
//    Verify that the "image" column is removed and "X" column is added.
TEST_F(MindDataTestMapOp, TestAsMap) {
  Status rc;
  MS_LOG(INFO) << "Doing TestAsMap.";

  // Note: The above storage config yields 5 buffers, each with 2 rows, for a total of 10 rows.
  auto my_storage_op = this->CreateStorageOp();
  rc = my_tree_->AssociateNode(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  auto my_no_op = std::make_shared<mindspore::dataset::test::NoOp>();
  std::vector<std::shared_ptr<TensorOp>> my_func_list;
  my_func_list.push_back(my_no_op);
  std::shared_ptr<MapOp> my_map_op;
  MapOp::Builder builder;
  builder.SetInColNames({"image"})
      .SetOutColNames({"X"})
      .SetTensorFuncs(std::move(my_func_list))
      .SetNumWorkers(1);
  rc = builder.Build(&my_map_op);
  rc = my_tree_->AssociateNode(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_map_op->AddChild(my_storage_op);
  EXPECT_TRUE(rc.IsOk());

  // Assign the tree root
  rc = my_tree_->AssignRoot(my_map_op);
  EXPECT_TRUE(rc.IsOk());

  // Now prepare the tree
  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorMap tensor_map;
  rc = di.GetNextAsMap(&tensor_map);
  EXPECT_TRUE(rc.IsOk());
  EXPECT_EQ(tensor_map.size(), 4);
  EXPECT_EQ(tensor_map.find("image"), tensor_map.end());
  EXPECT_NE(tensor_map.find("label"), tensor_map.end());
  EXPECT_NE(tensor_map.find("X"), tensor_map.end());
  EXPECT_NE(tensor_map.find("A"), tensor_map.end());
  EXPECT_NE(tensor_map.find("B"), tensor_map.end());
}

// Test3to1 scenario:
//    StorageOp reads a dataset that have column ordering |image|label|A|B|.
//    A 3-to-1 TensorOp picks the columns [image, A, B] and produce a column named "X".
//    Thus, based on the new MapOp behaviour, the column ordering will be |X|label|.
//    Verify that the only columns "X" and "label" exist.
TEST_F(MindDataTestMapOp, Test3to1) {
  Status rc;
  MS_LOG(INFO) << "Doing Test3to1.";

  // Note: The above storage config yields 5 buffers, each with 2 rows, for a total of 10 rows.
  auto my_storage_op = this->CreateStorageOp();
  rc = my_tree_->AssociateNode(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  auto my_op = std::make_shared<mindspore::dataset::test::ThreeToOneOp>();
  std::vector<std::shared_ptr<TensorOp>> my_func_list;
  my_func_list.push_back(my_op);
  std::shared_ptr<MapOp> my_map_op;
  MapOp::Builder builder;
  builder.SetInColNames({"image", "A", "B"})
      .SetOutColNames({"X"})
      .SetTensorFuncs(std::move(my_func_list))
      .SetNumWorkers(1);
  rc = builder.Build(&my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_map_op->AddChild(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssignRoot(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorMap tensor_map;
  rc = di.GetNextAsMap(&tensor_map);
  EXPECT_TRUE(rc.IsOk());
  while (!tensor_map.empty()) {
    EXPECT_EQ(tensor_map.size(), 2);
    EXPECT_EQ(tensor_map.find("image"), tensor_map.end());
    EXPECT_NE(tensor_map.find("label"), tensor_map.end());
    EXPECT_NE(tensor_map.find("X"), tensor_map.end());
    EXPECT_EQ(tensor_map.find("A"), tensor_map.end());
    EXPECT_EQ(tensor_map.find("B"), tensor_map.end());
    rc = di.GetNextAsMap(&tensor_map);
    EXPECT_TRUE(rc.IsOk());
  }
}

// Test1to3 scenario:
//    StorageOp reads a dataset that have column ordering |image|label|A|B|.
//    A 1-to-3 TensorOp picks the columns [image] and produce a column named [X, Y, Z].
//    Thus, based on the new MapOp behaviour, the column ordering will be |X|Y|Z|label|A|B|.
//    Verify that the only columns X, Y, Z are added (to the front) and followed by columns label, A, B..
TEST_F(MindDataTestMapOp, Test1to3) {
  Status rc;
  MS_LOG(INFO) << "Doing Test1to3.";

  // Note: The above storage config yields 5 buffers, each with 2 rows, for a total of 10 rows.
  auto my_storage_op = this->CreateStorageOp();
  rc = my_tree_->AssociateNode(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  auto my_op = std::make_shared<mindspore::dataset::test::OneToThreeOp>();
  std::vector<std::shared_ptr<TensorOp>> my_func_list;
  my_func_list.push_back(my_op);
  std::shared_ptr<MapOp> my_map_op;
  MapOp::Builder builder;
  builder.SetInColNames({"image"})
      .SetOutColNames({"X", "Y", "Z"})
      .SetTensorFuncs(std::move(my_func_list))
      .SetNumWorkers(1);
  rc = builder.Build(&my_map_op);
  rc = my_tree_->AssociateNode(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_map_op->AddChild(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssignRoot(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorMap tensor_map;
  rc = di.GetNextAsMap(&tensor_map);
  EXPECT_TRUE(rc.IsOk());
  EXPECT_EQ(tensor_map.size(), 6);
  EXPECT_EQ(tensor_map.find("image"), tensor_map.end());
  EXPECT_NE(tensor_map.find("label"), tensor_map.end());
  EXPECT_NE(tensor_map.find("A"), tensor_map.end());
  EXPECT_NE(tensor_map.find("B"), tensor_map.end());
  EXPECT_NE(tensor_map.find("X"), tensor_map.end());
  EXPECT_NE(tensor_map.find("Y"), tensor_map.end());
  EXPECT_NE(tensor_map.find("Z"), tensor_map.end());

  // Getting the next row as vector (by position).
  TensorRow tensor_list;
  rc =di.FetchNextTensorRow(&tensor_list);
  EXPECT_TRUE(rc.IsOk());

  // Based on the schema file, create the golden result to compare with.
  std::vector<DataType::Type> golden_types({DataType::Type::DE_UINT8, DataType::Type::DE_UINT8,
                                            DataType::Type::DE_UINT8, DataType::Type::DE_INT64,
                                            DataType::Type::DE_FLOAT32, DataType::Type::DE_INT64}
  );

  std::vector<uint64_t> golden_ranks({3, 3, 3, 1, 4, 1});

  std::vector<TensorShape> golden_shapes({TensorShape({3, 4, 2}), TensorShape({3, 4, 2}), TensorShape({3, 4, 2}),
        TensorShape({7}), TensorShape({1, 13, 14, 12}), TensorShape({9})} );

  while (!tensor_list.empty()) {
    for (uint32_t i = 0; i < tensor_list.size(); i++) {
      EXPECT_EQ(tensor_list[i]->type(), golden_types[i]);
      EXPECT_EQ(tensor_list[i]->Rank(), golden_ranks[i]);
      EXPECT_EQ(tensor_list[i]->shape(), golden_shapes[i]);
      EXPECT_NE(tensor_list[i]->GetBuffer(), nullptr);
    }
    rc = di.FetchNextTensorRow(&tensor_list);
    EXPECT_TRUE(rc.IsOk());
  }
}

// TestMultiTensorOp scenario:
//    StorageOp reads a dataset that have column ordering |image|label|A|B|.
//    A series of 3-to-1 and 1-to-3 TensorOps are applied to [image, A, B] and
//    produce final output columns [X, Y, Z].
//    Based on the new MapOp behaviour, the column ordering will be |X|Y|Z|label|.
TEST_F(MindDataTestMapOp, TestMultiTensorOp) {
  Status rc;
  MS_LOG(INFO) << "Doing TestMultiTensorOp.";

  // Note: The above storage config yields 5 buffers, each with 2 rows, for a total of 10 rows.
  auto my_storage_op = this->CreateStorageOp();
  rc = my_tree_->AssociateNode(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  auto my_op1 = std::make_shared<mindspore::dataset::test::ThreeToOneOp>();
  auto my_op2 = std::make_shared<mindspore::dataset::test::OneToThreeOp>();
  std::vector<std::shared_ptr<TensorOp>> my_func_list;
  my_func_list.push_back(my_op1);
  my_func_list.push_back(my_op2);
  std::shared_ptr<MapOp> my_map_op;
  MapOp::Builder builder;
  builder.SetInColNames({"image", "A", "B"})
      .SetOutColNames({"X", "Y", "Z"})
      .SetTensorFuncs(std::move(my_func_list))
      .SetNumWorkers(1);
  rc = builder.Build(&my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_map_op->AddChild(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssignRoot(my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorMap tensor_map;
  rc = di.GetNextAsMap(&tensor_map);
  EXPECT_TRUE(rc.IsOk());
  while (!tensor_map.empty()) {
    EXPECT_EQ(tensor_map.size(), 4);
    EXPECT_EQ(tensor_map.find("image"), tensor_map.end());
    EXPECT_EQ(tensor_map.find("A"), tensor_map.end());
    EXPECT_EQ(tensor_map.find("B"), tensor_map.end());
    EXPECT_NE(tensor_map.find("label"), tensor_map.end());
    EXPECT_NE(tensor_map.find("X"), tensor_map.end());
    EXPECT_NE(tensor_map.find("Y"), tensor_map.end());
    EXPECT_NE(tensor_map.find("Z"), tensor_map.end());

    // XYZ are Tensor shared_ptr to image, so it should have the same shape as image column.
    EXPECT_EQ(tensor_map["X"]->shape(), TensorShape({3, 4, 2}));
    EXPECT_EQ(tensor_map["Y"]->shape(), TensorShape({3, 4, 2}));
    EXPECT_EQ(tensor_map["Z"]->shape(), TensorShape({3, 4, 2}));
    rc = di.GetNextAsMap(&tensor_map);
    EXPECT_TRUE(rc.IsOk());
  }
}

TEST_F(MindDataTestMapOp, TestStorageRepeatMap) {
  Status rc;
  MS_LOG(INFO) << "Doing TestStorageRepeatMap.";
  uint32_t num_repeats = 3;

  // Note: The above storage config yields 5 buffers, each with 2 rows, for a total
  // of 10 rows.
  auto my_storage_op = this->CreateStorageOp();
  rc = my_tree_->AssociateNode(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  auto my_no_op = std::make_shared<mindspore::dataset::test::NoOp>();
  std::vector<std::shared_ptr<TensorOp>> my_func_list;
  my_func_list.push_back(my_no_op);

  std::shared_ptr<RepeatOp> my_repeat_op;
  rc = RepeatOp::Builder(num_repeats).Build(&my_repeat_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_repeat_op);
  EXPECT_TRUE(rc.IsOk());

  std::shared_ptr<MapOp> my_map_op;
  MapOp::Builder builder;
  builder.SetInColNames({"label"})
    .SetOutColNames({})
    .SetTensorFuncs(std::move(my_func_list))
    .SetNumWorkers(5);
  rc = builder.Build(&my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_map_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_map_op->AddChild(my_repeat_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_repeat_op->AddChild(my_storage_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_tree_->AssignRoot(my_map_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorRow tensor_list;
  rc = di.FetchNextTensorRow(&tensor_list);
  EXPECT_TRUE(rc.IsOk());
  EXPECT_EQ(tensor_list.size(), 4);
  uint32_t row_count = 0;
  while (!tensor_list.empty()) {
    row_count++;
    MS_LOG(INFO) << "row_count: " << row_count << ".";
    rc = di.FetchNextTensorRow(&tensor_list);
    EXPECT_TRUE(rc.IsOk());

  }
  ASSERT_EQ(row_count, 10 * num_repeats);
}

TEST_F(MindDataTestMapOp, TestStorageMapRepeat) {
  Status rc;
  MS_LOG(INFO) << "Doing TestStorageMapRepeat.";
  uint32_t num_repeats = 3;

  // Note: The above storage config yields 5 buffers, each with 2 rows, for a total
  // of 10 rows.
  auto my_storage_op = this->CreateStorageOp();
  rc = my_tree_->AssociateNode(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  auto my_no_op = std::make_shared<mindspore::dataset::test::NoOp>();
  std::vector<std::shared_ptr<TensorOp>> my_func_list;
  my_func_list.push_back(my_no_op);

  std::shared_ptr<RepeatOp> my_repeat_op;
  rc = RepeatOp::Builder(num_repeats).Build(&my_repeat_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_repeat_op);
  EXPECT_TRUE(rc.IsOk());

  std::shared_ptr<MapOp> my_map_op;
  MapOp::Builder builder;
  builder.SetInColNames({"label"})
    .SetOutColNames({})
    .SetTensorFuncs(std::move(my_func_list))
    .SetNumWorkers(50);
  rc = builder.Build(&my_map_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_map_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_repeat_op->AddChild(my_map_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_map_op->AddChild(my_storage_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_tree_->AssignRoot(my_repeat_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorRow tensor_list;
  rc = di.FetchNextTensorRow(&tensor_list);
  EXPECT_TRUE(rc.IsOk());
  EXPECT_EQ(tensor_list.size(), 4);
  uint32_t row_count = 0;
  while (!tensor_list.empty()) {
    row_count++;
    MS_LOG(INFO) << "row_count: " << row_count << ".";
    rc = di.FetchNextTensorRow(&tensor_list);
    EXPECT_TRUE(rc.IsOk());
  }
  ASSERT_EQ(row_count, 10 * num_repeats);
}

TEST_F(MindDataTestMapOp, Storage_Decode_Repeat_Resize) {
  Status rc;
  MS_LOG(INFO) << "Doing Storage_Decode_Repeat_Resize.";
  uint32_t num_repeats = 2;

  std::string dataset_path_ = datasets_root_path_ + "/" + "test_tf_file_3_images";
  std::shared_ptr<StorageOp> my_storage_op;
  StorageOp::Builder sobuilder;
  sobuilder.SetDatasetFilesDir(dataset_path_)
    .SetColumnsToLoad({"image", "label"})
    .SetRowsPerBuffer(2)
    .SetWorkerConnectorSize(2)
    .SetNumWorkers(2);
  rc = sobuilder.Build(&my_storage_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_tree_->AssociateNode(my_storage_op);
  EXPECT_TRUE(rc.IsOk());
  auto decode_op = std::make_shared<DecodeOp>();
  std::vector<std::shared_ptr<TensorOp>> my_func_list;
  my_func_list.push_back(decode_op);

  std::shared_ptr<RepeatOp> my_repeat_op;
  rc = RepeatOp::Builder(num_repeats).Build(&my_repeat_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_repeat_op);
  EXPECT_TRUE(rc.IsOk());

  std::shared_ptr<MapOp> my_map_decode_op;
  MapOp::Builder builder;
  builder.SetInColNames({"image"})
    .SetOutColNames({})
    .SetTensorFuncs(std::move(my_func_list))
    .SetNumWorkers(4);
  rc = builder.Build(&my_map_decode_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_map_decode_op);
  EXPECT_TRUE(rc.IsOk());


  auto resize_op = std::make_shared<ResizeOp>(300, 300);
  std::vector<std::shared_ptr<TensorOp>> my_func_list2;
  my_func_list2.push_back(resize_op);
  std::shared_ptr<MapOp> my_map_resize_op;
  MapOp::Builder builder2;
  builder2.SetInColNames({"image"})
    .SetOutColNames({})
    .SetTensorFuncs(std::move(my_func_list2))
    .SetNumWorkers(5);
  rc = builder2.Build(&my_map_resize_op);
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->AssociateNode(my_map_resize_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_map_decode_op->AddChild(my_storage_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_repeat_op->AddChild(my_map_decode_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_map_resize_op->AddChild(my_repeat_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_tree_->AssignRoot(my_map_resize_op);
  EXPECT_TRUE(rc.IsOk());

  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorRow tensor_list;
  rc = di.FetchNextTensorRow(&tensor_list);
  EXPECT_TRUE(rc.IsOk());
  EXPECT_EQ(tensor_list.size(), 2);
  uint32_t row_count = 0;
  while (!tensor_list.empty()) {
    row_count++;
    rc = di.FetchNextTensorRow(&tensor_list);
    EXPECT_TRUE(rc.IsOk());
  }

  ASSERT_EQ(row_count, 6);
}

TEST_F(MindDataTestMapOp, ImageFolder_Decode_Repeat_Resize) {
  Status rc;
  MS_LOG(INFO) << "Doing ImageFolder_Decode_Repeat_Resize.";

  std::string folder_path = datasets_root_path_ + "/testPK/data";

  uint32_t num_repeats = 2;
  std::shared_ptr<RepeatOp> repeat_op;
  rc = RepeatOp::Builder(num_repeats).Build(&repeat_op);
  EXPECT_TRUE(rc.IsOk());

  auto decode_op = std::make_shared<DecodeOp>();
  std::vector<std::shared_ptr<TensorOp>> func_list;
  func_list.push_back(decode_op);

  std::shared_ptr<MapOp> map_decode_map;
  MapOp::Builder map_decode_builder;
  map_decode_builder.SetInColNames({"image"})
    .SetOutColNames({})
    .SetTensorFuncs(func_list)
    .SetNumWorkers(4);
  rc = map_decode_builder.Build(&map_decode_map);
  EXPECT_TRUE(rc.IsOk());

  auto resize_op = std::make_shared<ResizeOp>(300, 300);
  std::vector<std::shared_ptr<TensorOp>> func_list2;
  func_list2.push_back(resize_op);
  std::shared_ptr<MapOp> map_resize_op;
  MapOp::Builder map_resize_builder;
  map_resize_builder.SetInColNames({"image"})
    .SetOutColNames({})
    .SetTensorFuncs(func_list2)
    .SetNumWorkers(5);
  rc = map_resize_builder.Build(&map_resize_op);
  EXPECT_TRUE(rc.IsOk());

  my_tree_ = Build({ImageFolder(16, 2, 32, folder_path, false), map_decode_map, repeat_op, map_resize_op});
  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorMap tensor_map;
  di.GetNextAsMap(&tensor_map);
  EXPECT_TRUE(rc.IsOk());
  uint64_t i = 0;
  int32_t label = 0;
  int32_t img_class[] = {0, 1, 2, 3};
  std::string result;
  while (tensor_map.size() != 0) {
    tensor_map["label"]->GetItemAt<int32_t>(&label, {});
    MS_LOG(DEBUG) << "row:" << i << "\tlabel:" << label << "\n";
    EXPECT_TRUE(img_class[(i % 44) / 11] == label);
    // Dump all the image into string, to be used as a comparison later.
    result.append((char *)tensor_map["image"]->GetBuffer(), (int64_t)tensor_map["image"]->Size());
    di.GetNextAsMap(&tensor_map);
    i++;
  }
  EXPECT_TRUE(i == 88);

  // Part-2 : creating mapop with performance mode = false, to check if the result is the same
  // as when performance mode = true.
  rc = RepeatOp::Builder(num_repeats).Build(&repeat_op);
  EXPECT_TRUE(rc.IsOk());

  map_decode_builder.SetInColNames({"image"})
    .SetOutColNames({})
    .SetTensorFuncs(func_list)
    .SetNumWorkers(14)
    .SetPerformanceMode(false);
  rc = map_decode_builder.Build(&map_decode_map);
  EXPECT_TRUE(rc.IsOk());

  map_resize_builder.SetInColNames({"image"})
    .SetOutColNames({})
    .SetTensorFuncs(func_list2)
    .SetNumWorkers(15)
    .SetPerformanceMode(false);
  rc = map_resize_builder.Build(&map_resize_op);
  EXPECT_TRUE(rc.IsOk());

  auto my_tree_2 = Build({ImageFolder(16, 2, 32, folder_path, false), map_decode_map, repeat_op, map_resize_op});

  rc = my_tree_2->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_2->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di2(my_tree_2);
  di2.GetNextAsMap(&tensor_map);
  EXPECT_TRUE(rc.IsOk());
  i = 0;
  label = 0;
  std::string result2;
  while (tensor_map.size() != 0) {
    tensor_map["label"]->GetItemAt<int32_t>(&label, {});
    MS_LOG(DEBUG) << "row:" << i << "\tlabel:" << label << "\n";
    EXPECT_TRUE(img_class[(i % 44) / 11] == label);
    result2.append((char *)tensor_map["image"]->GetBuffer(), (int64_t)tensor_map["image"]->Size());
    di2.GetNextAsMap(&tensor_map);
    i++;
  }
  EXPECT_TRUE(i == 88);

  EXPECT_EQ(result.size(), result2.size());
  EXPECT_EQ(result, result2);
}


TEST_F(MindDataTestMapOp, ImageFolder_Decode_Repeat_Resize_NoInputColumns) {
  Status rc;
  MS_LOG(INFO) << "Doing ImageFolder_Decode_Repeat_Resize_NoInputColumns.";

  std::string folder_path = datasets_root_path_ + "/testPK/data";

  uint32_t num_repeats = 2;
  std::shared_ptr<RepeatOp> repeat_op;
  rc = RepeatOp::Builder(num_repeats).Build(&repeat_op);
  EXPECT_TRUE(rc.IsOk());

  auto decode_op = std::make_shared<DecodeOp>();
  std::vector<std::shared_ptr<TensorOp>> func_list;
  func_list.push_back(decode_op);

  std::shared_ptr<MapOp> map_decode_map;
  MapOp::Builder map_decode_builder;
  map_decode_builder.SetInColNames({})
    .SetOutColNames({})
    .SetTensorFuncs(func_list)
    .SetNumWorkers(4);
  rc = map_decode_builder.Build(&map_decode_map);
  EXPECT_TRUE(rc.IsOk());

  auto resize_op = std::make_shared<ResizeOp>(300, 300);
  std::vector<std::shared_ptr<TensorOp>> func_list2;
  func_list2.push_back(resize_op);
  std::shared_ptr<MapOp> map_resize_op;
  MapOp::Builder map_resize_builder;
  map_resize_builder.SetTensorFuncs(func_list2).SetNumWorkers(5);
  rc = map_resize_builder.Build(&map_resize_op);
  EXPECT_TRUE(rc.IsOk());

  my_tree_ = Build({ImageFolder(16, 2, 32, folder_path, false), map_decode_map, repeat_op, map_resize_op});
  rc = my_tree_->Prepare();
  EXPECT_TRUE(rc.IsOk());
  rc = my_tree_->Launch();
  EXPECT_TRUE(rc.IsOk());

  // Start the loop of reading tensors from our pipeline
  DatasetIterator di(my_tree_);
  TensorMap tensor_map;
  di.GetNextAsMap(&tensor_map);
  EXPECT_TRUE(rc.IsOk());
  uint64_t i = 0;
  int32_t label = 0;
  int32_t img_class[] = {0, 1, 2, 3};
  std::string result;
  while (tensor_map.size() != 0) {
    tensor_map["label"]->GetItemAt<int32_t>(&label, {});
    EXPECT_TRUE(img_class[(i % 44) / 11] == label);
    di.GetNextAsMap(&tensor_map);
    i++;
  }
  EXPECT_TRUE(i == 88);
}