// Copyright (c) ONNX Project Contributors
//
// SPDX-License-Identifier: Apache-2.0

#include "onnx_op/operator_sets_text.h"

#include <gtest/gtest.h>

#include <vector>

using namespace ONNX_LIGHT_NAMESPACE;

namespace Test {

constexpr size_t kExpectedTextSchemaCount = 5;

static const core::schema::LightOpSchema *
FindByVersion(const std::vector<core::schema::LightOpSchema> &schemas, int version) {
  for (const auto &schema : schemas) {
    if (schema.since_version() == version) {
      return &schema;
    }
  }
  return nullptr;
}

TEST(OnnxOpTextRegistrationTest, ReturnsStringConcatSchema) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> string_concat_schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory("StringConcat");

  EXPECT_EQ(schemas.size(), kExpectedTextSchemaCount);

  const core::schema::LightOpSchema *const string_concat_v20 =
      FindByVersion(string_concat_schemas, 20);
  ASSERT_NE(nullptr, string_concat_v20);

  EXPECT_EQ(string_concat_v20->domain(), "ai.onnx");
  EXPECT_EQ(string_concat_v20->inputs().size(), 2u);
  EXPECT_EQ(string_concat_v20->outputs().size(), 1u);
  EXPECT_EQ(string_concat_v20->type_constraints().size(), 1u);

  EXPECT_EQ(string_concat_v20->inputs()[0].name, "X");
  EXPECT_EQ(string_concat_v20->inputs()[0].description, "Tensor to prepend in concatenation");
  EXPECT_EQ(string_concat_v20->inputs()[0].type, "T");

  EXPECT_EQ(string_concat_v20->inputs()[1].name, "Y");
  EXPECT_EQ(string_concat_v20->inputs()[1].description, "Tensor to append in concatenation");
  EXPECT_EQ(string_concat_v20->inputs()[1].type, "T");

  EXPECT_EQ(string_concat_v20->outputs()[0].name, "Z");
  EXPECT_EQ(string_concat_v20->outputs()[0].description, "Concatenated string tensor");
  EXPECT_EQ(string_concat_v20->outputs()[0].type, "T");

  EXPECT_EQ(string_concat_v20->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(string_concat_v20->type_constraints()[0].allowed_type_strs,
            std::vector<core::schema::TensorType>{core::schema::TensorType::kString});
  EXPECT_EQ(string_concat_v20->type_constraints()[0].description,
            "Inputs and outputs must be UTF-8 strings");

  EXPECT_NE(string_concat_v20->doc().find("StringConcat concatenates string tensors elementwise"),
            std::string::npos);
}

TEST(OnnxOpTextRegistrationTest, ReturnsStringSplitSchema) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> string_split_schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory("StringSplit");

  const core::schema::LightOpSchema *const string_split_v20 =
      FindByVersion(string_split_schemas, 20);
  ASSERT_NE(nullptr, string_split_v20);

  EXPECT_EQ(string_split_v20->domain(), "ai.onnx");
  EXPECT_EQ(string_split_v20->inputs().size(), 1u);
  EXPECT_EQ(string_split_v20->outputs().size(), 2u);
  EXPECT_EQ(string_split_v20->type_constraints().size(), 3u);

  EXPECT_EQ(string_split_v20->inputs()[0].name, "X");
  EXPECT_EQ(string_split_v20->inputs()[0].type, "T1");

  EXPECT_EQ(string_split_v20->outputs()[0].name, "Y");
  EXPECT_EQ(string_split_v20->outputs()[0].type, "T2");
  EXPECT_EQ(string_split_v20->outputs()[1].name, "Z");
  EXPECT_EQ(string_split_v20->outputs()[1].type, "T3");

  EXPECT_EQ(string_split_v20->type_constraints()[0].type_param_str, "T1");
  EXPECT_EQ(string_split_v20->type_constraints()[0].allowed_type_strs,
            std::vector<core::schema::TensorType>{core::schema::TensorType::kString});
  EXPECT_EQ(string_split_v20->type_constraints()[1].type_param_str, "T2");
  EXPECT_EQ(string_split_v20->type_constraints()[1].allowed_type_strs,
            std::vector<core::schema::TensorType>{core::schema::TensorType::kString});
  EXPECT_EQ(string_split_v20->type_constraints()[2].type_param_str, "T3");
  EXPECT_EQ(string_split_v20->type_constraints()[2].allowed_type_strs,
            std::vector<core::schema::TensorType>{core::schema::TensorType::kInt64});

  EXPECT_NE(string_split_v20->doc().find("StringSplit splits a string tensor"), std::string::npos);
}

TEST(OnnxOpTextRegistrationTest, ReturnsRegexFullMatchSchema) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> regex_full_match_schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory("RegexFullMatch");

  const core::schema::LightOpSchema *const regex_v20 = FindByVersion(regex_full_match_schemas, 20);
  ASSERT_NE(nullptr, regex_v20);

  EXPECT_EQ(regex_v20->domain(), "ai.onnx");
  EXPECT_EQ(regex_v20->inputs().size(), 1u);
  EXPECT_EQ(regex_v20->outputs().size(), 1u);
  EXPECT_EQ(regex_v20->type_constraints().size(), 2u);

  EXPECT_EQ(regex_v20->inputs()[0].name, "X");
  EXPECT_EQ(regex_v20->inputs()[0].type, "T1");
  EXPECT_EQ(regex_v20->outputs()[0].name, "Y");
  EXPECT_EQ(regex_v20->outputs()[0].type, "T2");

  EXPECT_EQ(regex_v20->type_constraints()[0].type_param_str, "T1");
  EXPECT_EQ(regex_v20->type_constraints()[0].allowed_type_strs,
            std::vector<core::schema::TensorType>{core::schema::TensorType::kString});
  EXPECT_EQ(regex_v20->type_constraints()[1].type_param_str, "T2");
  EXPECT_EQ(regex_v20->type_constraints()[1].allowed_type_strs,
            std::vector<core::schema::TensorType>{core::schema::TensorType::kBool});

  EXPECT_NE(regex_v20->doc().find("RegexFullMatch performs a full regex match"), std::string::npos);
}

TEST(OnnxOpTextRegistrationTest, ReturnsStringNormalizerSchema) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> string_normalizer_schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory("StringNormalizer");

  const core::schema::LightOpSchema *const normalizer_v10 =
      FindByVersion(string_normalizer_schemas, 10);
  ASSERT_NE(nullptr, normalizer_v10);

  EXPECT_EQ(normalizer_v10->domain(), "ai.onnx");
  EXPECT_EQ(normalizer_v10->inputs().size(), 1u);
  EXPECT_EQ(normalizer_v10->outputs().size(), 1u);
  EXPECT_EQ(normalizer_v10->type_constraints().size(), 0u);

  EXPECT_EQ(normalizer_v10->inputs()[0].name, "X");
  EXPECT_EQ(normalizer_v10->inputs()[0].type, "tensor(string)");
  EXPECT_EQ(normalizer_v10->outputs()[0].name, "Y");
  EXPECT_EQ(normalizer_v10->outputs()[0].type, "tensor(string)");

  EXPECT_NE(normalizer_v10->doc().find("StringNormalization performs string operations"),
            std::string::npos);
}

TEST(OnnxOpTextRegistrationTest, ReturnsTfIdfVectorizerSchema) {
  const std::vector<core::schema::LightOpSchema> tf_idf_schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory("TfIdfVectorizer");

  const core::schema::LightOpSchema *const tf_idf_v9 = FindByVersion(tf_idf_schemas, 9);
  ASSERT_NE(nullptr, tf_idf_v9);

  EXPECT_EQ(tf_idf_v9->domain(), "ai.onnx");
  EXPECT_EQ(tf_idf_v9->name(), "TfIdfVectorizer");
  EXPECT_EQ(tf_idf_v9->inputs().size(), 1u);
  EXPECT_EQ(tf_idf_v9->outputs().size(), 1u);
  EXPECT_EQ(tf_idf_v9->type_constraints().size(), 2u);
  EXPECT_EQ(tf_idf_v9->attributes().size(), 9u);

  EXPECT_EQ(tf_idf_v9->inputs()[0].name, "X");
  EXPECT_EQ(tf_idf_v9->inputs()[0].type, "T");
  EXPECT_EQ(tf_idf_v9->outputs()[0].name, "Y");
  EXPECT_EQ(tf_idf_v9->outputs()[0].type, "T1");

  EXPECT_EQ(tf_idf_v9->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(tf_idf_v9->type_constraints()[0].allowed_type_strs,
            (std::vector<core::schema::TensorType>{core::schema::TensorType::kString,
                                                   core::schema::TensorType::kInt32,
                                                   core::schema::TensorType::kInt64}));
  EXPECT_EQ(tf_idf_v9->type_constraints()[1].type_param_str, "T1");
  EXPECT_EQ(tf_idf_v9->type_constraints()[1].allowed_type_strs,
            std::vector<core::schema::TensorType>{core::schema::TensorType::kFloat});

  EXPECT_NE(tf_idf_v9->doc().find("This transform extracts n-grams"), std::string::npos);
}

TEST(OnnxOpTextRegistrationTest, StripDocsRemovesDocumentation) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::text::GetAllOnnxOpTextSchemasWithHistory(/*op_type=*/"", /*init_doc=*/false);
  EXPECT_EQ(schemas.size(), kExpectedTextSchemaCount);
  for (const auto &schema : schemas) {
    EXPECT_EQ(schema.doc(), "");
  }
}

} // namespace Test
