/* Copyright 2023 The MediaPipe Authors.

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 "mediapipe/tasks/cc/text/custom_ops/ragged/ragged_tensor_to_tensor_tflite.h"

#include <initializer_list>
#include <string>
#include <vector>

#include "flatbuffers/flexbuffers.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/interpreter.h"
#include "tensorflow/lite/kernels/internal/tensor.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/schema/schema_generated.h"

namespace mediapipe::tflite_operations {
namespace {

using ::tflite::TensorType;
using ::tflite::TensorType_FLOAT32;
using ::tflite::TensorType_INT32;

class RaggedTensorToTensorOpModel : public tflite::SingleOpModel {
 public:
  RaggedTensorToTensorOpModel(int output_shape_dims,
                              std::initializer_list<int> values_shape,
                              std::initializer_list<std::initializer_list<int>>
                                  partition_tensors_shapes,
                              std::vector<std::string> partition_types,
                              TensorType value_type = TensorType_FLOAT32,
                              TensorType index_type = TensorType_INT32) {
    // A structure to collect shapes for the input.
    std::vector<std::vector<int>> shapes;
    input_shape_ = AddInput(index_type);
    shapes.push_back({output_shape_dims});
    input_values_ = AddInput(value_type);
    shapes.emplace_back(values_shape);
    input_default_values_ = AddInput(value_type);
    shapes.push_back({1});
    for (const auto& p : partition_tensors_shapes) {
      partition_tensors_.push_back(AddInput(TensorType_INT32));
      shapes.emplace_back(p);
    }
    output_ = AddOutput(value_type);

    flexbuffers::Builder fbb;
    size_t start = fbb.StartMap();
    {
      size_t start = fbb.StartVector("row_partition_types");
      for (const auto& s : partition_types) {
        fbb.String(s);
      }
      fbb.EndVector(start, /*typed=*/true, /*fixed=*/false);
    }
    fbb.Int("num_row_partition_tensors", partition_types.size());
    fbb.EndMap(start);
    fbb.Finish();
    SetCustomOp("RaggedTensorToTensor", fbb.GetBuffer(),
                tflite_operations::Register_RAGGED_TENSOR_TO_TENSOR);
    BuildInterpreter(shapes);
  }

  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }

  std::vector<float> GetOutputFloat() { return ExtractVector<float>(output_); }
  std::vector<int32> GetOutputInt() { return ExtractVector<int32>(output_); }

  void InvokeFloat(const std::vector<int>& shape,
                   const std::vector<float>& values, float default_value,
                   const std::vector<std::vector<int>>& partition_values) {
    PopulateTensor(input_shape_, shape);
    PopulateTensor(input_values_, values);
    PopulateTensor(input_default_values_, {default_value});
    for (int i = 0; i < partition_values.size(); ++i) {
      PopulateTensor(partition_tensors_[i], partition_values[i]);
    }
    SingleOpModel::Invoke();
  }
  void InvokeInt(const std::vector<int>& shape,
                 const std::vector<int32>& values, int32 default_value,
                 const std::vector<std::vector<int>>& partition_values) {
    PopulateTensor(input_shape_, shape);
    PopulateTensor(input_values_, values);
    PopulateTensor(input_default_values_, {default_value});
    for (int i = 0; i < partition_values.size(); ++i) {
      PopulateTensor(partition_tensors_[i], partition_values[i]);
    }
    SingleOpModel::Invoke();
  }

 private:
  int input_shape_;
  int input_values_;
  int input_default_values_;
  std::vector<int> partition_tensors_;
  int output_;
};

TEST(RaggedTensorToTensorTest, RaggedTensorToTensor) {
  // indices = [2, 1, 0, 3]
  // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
  // params.shape = [4, None]
  RaggedTensorToTensorOpModel model(
      2,           // output_shape_dims
      {9},         // values_shape
      {{1}, {9}},  // partition_tensors_shapes
      std::vector<std::string>({"FIRST_DIM_SIZE", "VALUE_ROWIDS"}));
  model.InvokeFloat({4, 4},                                // shape
                    {.1, .2, .3, .4, .5, .6, .7, .8, .9},  // values
                    1.5,                                   // default_value
                    std::vector<std::vector<int>>(
                        {std::vector<int>({4}),
                         std::vector<int>({0, 0, 0, 2, 2, 2, 2, 3, 3})}));
  EXPECT_THAT(model.GetOutputShape(), testing::ElementsAreArray({4, 4}));
  EXPECT_THAT(model.GetOutputFloat(),
              testing::ElementsAreArray({.1, .2, .3, 1.5, 1.5, 1.5, 1.5, 1.5,
                                         .4, .5, .6, .7, .8, .9, 1.5, 1.5}));
}

TEST(RaggedTensorToTensorTest, RaggedTensorToTensorRowSplits) {
  // indices = [2, 1, 0, 3]
  // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
  RaggedTensorToTensorOpModel model(2,      // output_shape_dims
                                    {9},    // values_shape
                                    {{5}},  // partition_tensors_shapes
                                    std::vector<std::string>({"ROW_SPLITS"}));
  model.InvokeFloat(
      {4, 4},                                // shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9},  // values
      1.5,                                   // default_value
      std::vector<std::vector<int>>({std::vector<int>({0, 3, 3, 7, 9})}));
  EXPECT_THAT(model.GetOutputShape(), testing::ElementsAreArray({4, 4}));
  EXPECT_THAT(model.GetOutputFloat(),
              testing::ElementsAreArray({.1, .2, .3, 1.5, 1.5, 1.5, 1.5, 1.5,
                                         .4, .5, .6, .7, .8, .9, 1.5, 1.5}));
}

TEST(RaggedTensorToTensorTest, RaggedTensorToTensor_3DParams) {
  // params = [
  //           [[]],
  //           [[.1, .2], [.3]],
  //           [],
  //           [[.4, .5], [.6, .7, .8]],
  //           [[.9]]
  //          ]
  RaggedTensorToTensorOpModel model(
      3,                // output_shape_dims
      {9},              // values_shape
      {{1}, {6}, {9}},  // partition_tensors_shapes
      std::vector<std::string>(
          {"FIRST_DIM_SIZE", "VALUE_ROWIDS", "VALUE_ROWIDS"}));
  model.InvokeFloat(
      {5, 2, 3},                             // shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9},  // values
      1.5,                                   // default_value
      std::vector<std::vector<int>>(
          {std::vector<int>({5}), std::vector<int>({0, 1, 1, 3, 3, 4}),
           std::vector<int>({1, 1, 2, 3, 3, 4, 4, 4, 5})}));

  EXPECT_THAT(model.GetOutputShape(), testing::ElementsAreArray({5, 2, 3}));
  EXPECT_THAT(model.GetOutputFloat(),
              testing::ElementsAreArray({1.5, 1.5, 1.5, 1.5, 1.5, 1.5, .1,  .2,
                                         1.5, .3,  1.5, 1.5, 1.5, 1.5, 1.5, 1.5,
                                         1.5, 1.5, .4,  .5,  1.5, .6,  .7,  .8,
                                         .9,  1.5, 1.5, 1.5, 1.5, 1.5}));
}

TEST(RaggedTensorToTensorOpTest, RaggedTensorToTensor_3DParamsRowSplits) {
  // params = [
  //           [[]],
  //           [[.1, .2], [.3]],
  //           [],
  //           [[.4, .5], [.6, .7, .8]],
  //           [[.9]]
  //          ]
  RaggedTensorToTensorOpModel model(
      3,           // output_shape_dims
      {9},         // values_shape
      {{6}, {7}},  // partition_tensors_shapes
      std::vector<std::string>({"ROW_SPLITS", "ROW_SPLITS"}));
  model.InvokeFloat(
      {5, 2, 3},                             // shape
      {.1, .2, .3, .4, .5, .6, .7, .8, .9},  // values
      1.5,                                   // default_value
      std::vector<std::vector<int>>({std::vector<int>({0, 1, 3, 3, 5, 6}),
                                     std::vector<int>({0, 0, 2, 3, 5, 8, 9})}));
  EXPECT_THAT(model.GetOutputShape(), testing::ElementsAreArray({5, 2, 3}));
  EXPECT_THAT(model.GetOutputFloat(),
              testing::ElementsAreArray({1.5, 1.5, 1.5, 1.5, 1.5, 1.5, .1,  .2,
                                         1.5, .3,  1.5, 1.5, 1.5, 1.5, 1.5, 1.5,
                                         1.5, 1.5, .4,  .5,  1.5, .6,  .7,  .8,
                                         .9,  1.5, 1.5, 1.5, 1.5, 1.5}));
}

TEST(RaggedTensorToTensorTest, RaggedTensorToTensor_3DParamsRowSplits2) {
  // params = [
  //           [[0, 1, 2], []],
  //           [],
  //           [[3]]
  //          ]

  RaggedTensorToTensorOpModel model(
      3,           // output_shape_dims
      {4},         // values_shape
      {{4}, {4}},  // partition_tensors_shapes
      std::vector<std::string>({"ROW_SPLITS", "ROW_SPLITS"}), TensorType_INT32);
  model.InvokeInt(
      {3, 2, 3},     // shape
      {0, 1, 2, 3},  // values
      5,             // default_value
      std::vector<std::vector<int>>(
          {std::vector<int>({0, 2, 2, 3}), std::vector<int>({0, 3, 3, 4})}));

  EXPECT_THAT(model.GetOutputShape(), testing::ElementsAreArray({3, 2, 3}));

  EXPECT_THAT(model.GetOutputInt(),
              testing::ElementsAreArray(
                  {0, 1, 2, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3, 5, 5, 5, 5, 5}));
}

TEST(RaggedTensorToTensorTest, RaggedTensorToTensorContractExpanded) {
  // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
  RaggedTensorToTensorOpModel model(
      2,           // output_shape_dims
      {9},         // values_shape
      {{1}, {9}},  // partition_tensors_shapes
      std::vector<std::string>({"FIRST_DIM_SIZE", "VALUE_ROWIDS"}));
  model.InvokeFloat({3, 5},                                // shape
                    {.1, .2, .3, .4, .5, .6, .7, .8, .9},  // values
                    1.5,                                   // default_value
                    std::vector<std::vector<int>>(
                        {std::vector<int>({4}),
                         std::vector<int>({0, 0, 0, 2, 2, 2, 2, 3, 3})}));
  EXPECT_THAT(model.GetOutputShape(), testing::ElementsAreArray({3, 5}));

  EXPECT_THAT(model.GetOutputFloat(),
              testing::ElementsAreArray({.1, .2, .3, 1.5, 1.5,     //
                                         1.5, 1.5, 1.5, 1.5, 1.5,  //
                                         .4, .5, .6, .7, 1.5}));
}

// Adds a dense dimension.
TEST(RaggedTensorToTensorTest, RaggedTensorToTensorContractExpandedDense) {
  // params = [[.1, .2, .3], [], [.4, .5, .6, .7], [.8, .9]]
  RaggedTensorToTensorOpModel model(
      3,           // output_shape_dims
      {9, 2},      // values_shape
      {{1}, {9}},  // partition_tensors_shapes
      std::vector<std::string>({"FIRST_DIM_SIZE", "VALUE_ROWIDS"}));

  model.InvokeFloat({3, 5, 2},  // shape
                    {.1, 1.1, .2, 1.2, .3, 1.3, .4, 1.4, .5, 1.5, .6, 1.6, .7,
                     1.7, .8, 1.8, .9, 1.9},  // values
                    1.5,                      // default_value
                    std::vector<std::vector<int>>(
                        {std::vector<int>({4}),
                         std::vector<int>({0, 0, 0, 2, 2, 2, 2, 3, 3})}));

  EXPECT_THAT(model.GetOutputShape(), testing::ElementsAreArray({3, 5, 2}));
  EXPECT_THAT(model.GetOutputFloat(),
              testing::ElementsAreArray(
                  {.1,  1.1, .2,  1.2, .3,  1.3, 1.5, 1.5, 1.5, 1.5,  //
                   1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5, 1.5,  //
                   .4,  1.4, .5,  1.5, .6,  1.6, .7,  1.7, 1.5, 1.5}));
}
}  // namespace
}  // namespace mediapipe::tflite_operations
