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

#include "onnx_op/operator_sets_traditionalml.h"

#include <gtest/gtest.h>

#include <vector>

using namespace ONNX_LIGHT_NAMESPACE;

namespace Test {

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(OnnxOpTraditionalMLRegistrationTest, ReturnsArrayFeatureExtractorAndLabelEncoderSchemas) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> array_feature_extractor_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("ArrayFeatureExtractor");
  const std::vector<core::schema::LightOpSchema> binarizer_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("Binarizer");
  const std::vector<core::schema::LightOpSchema> label_encoder_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("LabelEncoder");

  EXPECT_EQ(schemas.size(), 23u);

  const core::schema::LightOpSchema *const array_feature_extractor_v1 =
      FindByVersion(array_feature_extractor_schemas, 1);
  ASSERT_NE(nullptr, array_feature_extractor_v1);
  EXPECT_EQ(array_feature_extractor_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(array_feature_extractor_v1->inputs().size(), 2u);
  EXPECT_EQ(array_feature_extractor_v1->outputs().size(), 1u);
  EXPECT_EQ(array_feature_extractor_v1->inputs()[0].name, "X");
  EXPECT_EQ(array_feature_extractor_v1->inputs()[0].type, "T");
  EXPECT_EQ(array_feature_extractor_v1->inputs()[1].name, "Y");
  EXPECT_EQ(array_feature_extractor_v1->inputs()[1].type, "tensor(int64)");
  EXPECT_EQ(array_feature_extractor_v1->outputs()[0].name, "Z");
  EXPECT_EQ(array_feature_extractor_v1->outputs()[0].type, "T");
  EXPECT_EQ(array_feature_extractor_v1->type_constraints().size(), 1u);
  EXPECT_EQ(array_feature_extractor_v1->type_constraints()[0].allowed_type_strs.size(), 5u);
  EXPECT_EQ(array_feature_extractor_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kFloat);
  EXPECT_EQ(array_feature_extractor_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kDouble);
  EXPECT_EQ(array_feature_extractor_v1->type_constraints()[0].allowed_type_strs[2],
            core::schema::TensorType::kInt64);
  EXPECT_EQ(array_feature_extractor_v1->type_constraints()[0].allowed_type_strs[3],
            core::schema::TensorType::kInt32);
  EXPECT_EQ(array_feature_extractor_v1->type_constraints()[0].allowed_type_strs[4],
            core::schema::TensorType::kString);

  const core::schema::LightOpSchema *const binarizer_v1 = FindByVersion(binarizer_schemas, 1);
  ASSERT_NE(nullptr, binarizer_v1);
  EXPECT_EQ(binarizer_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(binarizer_v1->inputs().size(), 1u);
  EXPECT_EQ(binarizer_v1->outputs().size(), 1u);
  EXPECT_EQ(binarizer_v1->inputs()[0].name, "X");
  EXPECT_EQ(binarizer_v1->inputs()[0].type, "T");
  EXPECT_EQ(binarizer_v1->outputs()[0].name, "Y");
  EXPECT_EQ(binarizer_v1->outputs()[0].type, "T");
  EXPECT_EQ(binarizer_v1->type_constraints().size(), 1u);
  EXPECT_EQ(binarizer_v1->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(binarizer_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
  EXPECT_EQ(binarizer_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kFloat);
  EXPECT_EQ(binarizer_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kDouble);
  EXPECT_EQ(binarizer_v1->type_constraints()[0].allowed_type_strs[2],
            core::schema::TensorType::kInt64);
  EXPECT_EQ(binarizer_v1->type_constraints()[0].allowed_type_strs[3],
            core::schema::TensorType::kInt32);

  const core::schema::LightOpSchema *const label_encoder_v4 =
      FindByVersion(label_encoder_schemas, 4);
  ASSERT_NE(nullptr, label_encoder_v4);
  EXPECT_EQ(label_encoder_v4->domain(), "ai.onnx.ml");
  EXPECT_EQ(label_encoder_v4->inputs().size(), 1u);
  EXPECT_EQ(label_encoder_v4->outputs().size(), 1u);
  EXPECT_EQ(label_encoder_v4->type_constraints().size(), 2u);
  EXPECT_EQ(label_encoder_v4->type_constraints()[0].allowed_type_strs,
            label_encoder_v4->type_constraints()[1].allowed_type_strs);
  EXPECT_EQ(label_encoder_v4->type_constraints()[0].allowed_type_strs.size(), 6u);
  EXPECT_EQ(label_encoder_v4->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kString);
  EXPECT_EQ(label_encoder_v4->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kInt64);
  EXPECT_EQ(label_encoder_v4->type_constraints()[0].allowed_type_strs[2],
            core::schema::TensorType::kFloat);
  EXPECT_EQ(label_encoder_v4->type_constraints()[0].allowed_type_strs[3],
            core::schema::TensorType::kInt32);
  EXPECT_EQ(label_encoder_v4->type_constraints()[0].allowed_type_strs[4],
            core::schema::TensorType::kInt16);
  EXPECT_EQ(label_encoder_v4->type_constraints()[0].allowed_type_strs[5],
            core::schema::TensorType::kDouble);
  EXPECT_EQ(label_encoder_v4->inputs()[0].name, "X");
  EXPECT_EQ(label_encoder_v4->outputs()[0].name, "Y");
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsZipMapSchemaWithoutShapeInference) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> zip_map_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("ZipMap");

  const core::schema::LightOpSchema *const zipmap_v1 = FindByVersion(zip_map_schemas, 1);
  ASSERT_NE(nullptr, zipmap_v1);
  EXPECT_EQ(zipmap_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(zipmap_v1->inputs().size(), 1u);
  EXPECT_EQ(zipmap_v1->outputs().size(), 1u);
  EXPECT_EQ(zipmap_v1->type_constraints().size(), 1u);
  EXPECT_EQ(zipmap_v1->inputs()[0].name, "X");
  EXPECT_EQ(zipmap_v1->inputs()[0].type, "tensor(float)");
  EXPECT_EQ(zipmap_v1->outputs()[0].name, "Z");
  EXPECT_EQ(zipmap_v1->outputs()[0].type, "T");
  EXPECT_EQ(zipmap_v1->type_constraints()[0].allowed_type_strs.size(), 2u);
  EXPECT_EQ(zipmap_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kSeqMapStringFloat);
  EXPECT_EQ(zipmap_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kSeqMapInt64Float);
  EXPECT_STREQ(core::schema::ToTypeString(zipmap_v1->type_constraints()[0].allowed_type_strs[0]),
               "seq(map(string, float))");
  EXPECT_STREQ(core::schema::ToTypeString(zipmap_v1->type_constraints()[0].allowed_type_strs[1]),
               "seq(map(int64, float))");
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsTreeEnsembleSchema) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> tree_ensemble_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("TreeEnsemble");

  const core::schema::LightOpSchema *const tree_ensemble_v5 =
      FindByVersion(tree_ensemble_schemas, 5);
  ASSERT_NE(nullptr, tree_ensemble_v5);
  EXPECT_EQ(tree_ensemble_v5->domain(), "ai.onnx.ml");
  EXPECT_EQ(tree_ensemble_v5->inputs().size(), 1u);
  EXPECT_EQ(tree_ensemble_v5->outputs().size(), 1u);
  EXPECT_EQ(tree_ensemble_v5->type_constraints().size(), 1u);
  EXPECT_EQ(tree_ensemble_v5->inputs()[0].name, "X");
  EXPECT_EQ(tree_ensemble_v5->inputs()[0].type, "T");
  EXPECT_EQ(tree_ensemble_v5->outputs()[0].name, "Y");
  EXPECT_EQ(tree_ensemble_v5->outputs()[0].type, "T");
  EXPECT_EQ(tree_ensemble_v5->type_constraints()[0].allowed_type_strs.size(), 3u);
  EXPECT_EQ(tree_ensemble_v5->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kFloat);
  EXPECT_EQ(tree_ensemble_v5->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kDouble);
  EXPECT_EQ(tree_ensemble_v5->type_constraints()[0].allowed_type_strs[2],
            core::schema::TensorType::kFloat16);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsTreeEnsembleClassifierSchemaHistory) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> tree_ensemble_classifier_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("TreeEnsembleClassifier");

  for (int version : {1, 3, 5}) {
    const core::schema::LightOpSchema *const schema =
        FindByVersion(tree_ensemble_classifier_schemas, version);
    ASSERT_NE(nullptr, schema) << "version=" << version;
    EXPECT_EQ(schema->domain(), "ai.onnx.ml");
    EXPECT_EQ(schema->inputs().size(), 1u);
    EXPECT_EQ(schema->outputs().size(), 2u);
    EXPECT_EQ(schema->inputs()[0].name, "X");
    EXPECT_EQ(schema->inputs()[0].type, "T1");
    EXPECT_EQ(schema->outputs()[0].name, "Y");
    EXPECT_EQ(schema->outputs()[0].type, "T2");
    EXPECT_EQ(schema->outputs()[1].name, "Z");
    EXPECT_EQ(schema->outputs()[1].type, "tensor(float)");
    EXPECT_EQ(schema->type_constraints().size(), 2u);
    EXPECT_EQ(schema->type_constraints()[0].type_param_str, "T1");
    EXPECT_EQ(schema->type_constraints()[0].allowed_type_strs.size(), 4u);
    EXPECT_EQ(schema->type_constraints()[1].type_param_str, "T2");
    EXPECT_EQ(schema->type_constraints()[1].allowed_type_strs.size(), 2u);
    EXPECT_EQ(schema->type_constraints()[1].allowed_type_strs[0],
              core::schema::TensorType::kString);
    EXPECT_EQ(schema->type_constraints()[1].allowed_type_strs[1], core::schema::TensorType::kInt64);
  }
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsTreeEnsembleRegressorSchemaHistory) {
  const std::vector<core::schema::LightOpSchema> schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory();
  const std::vector<core::schema::LightOpSchema> tree_ensemble_regressor_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("TreeEnsembleRegressor");

  for (int version : {1, 3, 5}) {
    const core::schema::LightOpSchema *const schema =
        FindByVersion(tree_ensemble_regressor_schemas, version);
    ASSERT_NE(nullptr, schema) << "version=" << version;
    EXPECT_EQ(schema->domain(), "ai.onnx.ml");
    EXPECT_EQ(schema->inputs().size(), 1u);
    EXPECT_EQ(schema->outputs().size(), 1u);
    EXPECT_EQ(schema->inputs()[0].name, "X");
    EXPECT_EQ(schema->inputs()[0].type, "T");
    EXPECT_EQ(schema->outputs()[0].name, "Y");
    EXPECT_EQ(schema->outputs()[0].type, "tensor(float)");
    EXPECT_EQ(schema->type_constraints().size(), 1u);
    EXPECT_EQ(schema->type_constraints()[0].type_param_str, "T");
    EXPECT_EQ(schema->type_constraints()[0].allowed_type_strs.size(), 4u);
    EXPECT_EQ(schema->type_constraints()[0].allowed_type_strs[0], core::schema::TensorType::kFloat);
    EXPECT_EQ(schema->type_constraints()[0].allowed_type_strs[1],
              core::schema::TensorType::kDouble);
    EXPECT_EQ(schema->type_constraints()[0].allowed_type_strs[2], core::schema::TensorType::kInt64);
    EXPECT_EQ(schema->type_constraints()[0].allowed_type_strs[3], core::schema::TensorType::kInt32);
  }
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsOneHotEncoderSchema) {
  const std::vector<core::schema::LightOpSchema> one_hot_encoder_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("OneHotEncoder");

  const core::schema::LightOpSchema *const one_hot_encoder_v1 =
      FindByVersion(one_hot_encoder_schemas, 1);
  ASSERT_NE(nullptr, one_hot_encoder_v1);
  EXPECT_EQ(one_hot_encoder_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(one_hot_encoder_v1->inputs().size(), 1u);
  EXPECT_EQ(one_hot_encoder_v1->outputs().size(), 1u);
  EXPECT_EQ(one_hot_encoder_v1->inputs()[0].name, "X");
  EXPECT_EQ(one_hot_encoder_v1->inputs()[0].type, "T");
  EXPECT_EQ(one_hot_encoder_v1->outputs()[0].name, "Y");
  EXPECT_EQ(one_hot_encoder_v1->outputs()[0].type, "tensor(float)");
  EXPECT_EQ(one_hot_encoder_v1->type_constraints().size(), 1u);
  EXPECT_EQ(one_hot_encoder_v1->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(one_hot_encoder_v1->type_constraints()[0].allowed_type_strs.size(), 5u);
  EXPECT_EQ(one_hot_encoder_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kString);
  EXPECT_EQ(one_hot_encoder_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kInt64);
  EXPECT_EQ(one_hot_encoder_v1->type_constraints()[0].allowed_type_strs[2],
            core::schema::TensorType::kInt32);
  EXPECT_EQ(one_hot_encoder_v1->type_constraints()[0].allowed_type_strs[3],
            core::schema::TensorType::kFloat);
  EXPECT_EQ(one_hot_encoder_v1->type_constraints()[0].allowed_type_strs[4],
            core::schema::TensorType::kDouble);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsSVMClassifierSchema) {
  const std::vector<core::schema::LightOpSchema> svm_classifier_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("SVMClassifier");

  const core::schema::LightOpSchema *const svm_classifier_v1 =
      FindByVersion(svm_classifier_schemas, 1);
  ASSERT_NE(nullptr, svm_classifier_v1);
  EXPECT_EQ(svm_classifier_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(svm_classifier_v1->inputs().size(), 1u);
  EXPECT_EQ(svm_classifier_v1->outputs().size(), 2u);
  EXPECT_EQ(svm_classifier_v1->inputs()[0].name, "X");
  EXPECT_EQ(svm_classifier_v1->inputs()[0].type, "T1");
  EXPECT_EQ(svm_classifier_v1->outputs()[0].name, "Y");
  EXPECT_EQ(svm_classifier_v1->outputs()[0].type, "T2");
  EXPECT_EQ(svm_classifier_v1->outputs()[1].name, "Z");
  EXPECT_EQ(svm_classifier_v1->outputs()[1].type, "tensor(float)");
  EXPECT_EQ(svm_classifier_v1->type_constraints().size(), 2u);
  EXPECT_EQ(svm_classifier_v1->type_constraints()[0].type_param_str, "T1");
  EXPECT_EQ(svm_classifier_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
  EXPECT_EQ(svm_classifier_v1->type_constraints()[1].type_param_str, "T2");
  EXPECT_EQ(svm_classifier_v1->type_constraints()[1].allowed_type_strs.size(), 2u);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsSVMRegressorSchema) {
  const std::vector<core::schema::LightOpSchema> svm_regressor_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("SVMRegressor");

  const core::schema::LightOpSchema *const svm_regressor_v1 =
      FindByVersion(svm_regressor_schemas, 1);
  ASSERT_NE(nullptr, svm_regressor_v1);
  EXPECT_EQ(svm_regressor_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(svm_regressor_v1->inputs().size(), 1u);
  EXPECT_EQ(svm_regressor_v1->outputs().size(), 1u);
  EXPECT_EQ(svm_regressor_v1->inputs()[0].name, "X");
  EXPECT_EQ(svm_regressor_v1->inputs()[0].type, "T");
  EXPECT_EQ(svm_regressor_v1->outputs()[0].name, "Y");
  EXPECT_EQ(svm_regressor_v1->outputs()[0].type, "tensor(float)");
  EXPECT_EQ(svm_regressor_v1->type_constraints().size(), 1u);
  EXPECT_EQ(svm_regressor_v1->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(svm_regressor_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsScalerSchema) {
  const std::vector<core::schema::LightOpSchema> scaler_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("Scaler");

  const core::schema::LightOpSchema *const scaler_v1 = FindByVersion(scaler_schemas, 1);
  ASSERT_NE(nullptr, scaler_v1);
  EXPECT_EQ(scaler_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(scaler_v1->inputs().size(), 1u);
  EXPECT_EQ(scaler_v1->outputs().size(), 1u);
  EXPECT_EQ(scaler_v1->inputs()[0].name, "X");
  EXPECT_EQ(scaler_v1->inputs()[0].type, "T");
  EXPECT_EQ(scaler_v1->outputs()[0].name, "Y");
  EXPECT_EQ(scaler_v1->outputs()[0].type, "tensor(float)");
  EXPECT_EQ(scaler_v1->type_constraints().size(), 1u);
  EXPECT_EQ(scaler_v1->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(scaler_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
  EXPECT_EQ(scaler_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kFloat);
  EXPECT_EQ(scaler_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kDouble);
  EXPECT_EQ(scaler_v1->type_constraints()[0].allowed_type_strs[2],
            core::schema::TensorType::kInt64);
  EXPECT_EQ(scaler_v1->type_constraints()[0].allowed_type_strs[3],
            core::schema::TensorType::kInt32);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsLinearClassifierSchema) {
  const std::vector<core::schema::LightOpSchema> linear_classifier_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("LinearClassifier");

  const core::schema::LightOpSchema *const linear_classifier_v1 =
      FindByVersion(linear_classifier_schemas, 1);
  ASSERT_NE(nullptr, linear_classifier_v1);
  EXPECT_EQ(linear_classifier_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(linear_classifier_v1->inputs().size(), 1u);
  EXPECT_EQ(linear_classifier_v1->outputs().size(), 2u);
  EXPECT_EQ(linear_classifier_v1->inputs()[0].name, "X");
  EXPECT_EQ(linear_classifier_v1->inputs()[0].type, "T1");
  EXPECT_EQ(linear_classifier_v1->outputs()[0].name, "Y");
  EXPECT_EQ(linear_classifier_v1->outputs()[0].type, "T2");
  EXPECT_EQ(linear_classifier_v1->outputs()[1].name, "Z");
  EXPECT_EQ(linear_classifier_v1->outputs()[1].type, "tensor(float)");
  EXPECT_EQ(linear_classifier_v1->type_constraints().size(), 2u);
  EXPECT_EQ(linear_classifier_v1->type_constraints()[0].type_param_str, "T1");
  EXPECT_EQ(linear_classifier_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
  EXPECT_EQ(linear_classifier_v1->type_constraints()[1].type_param_str, "T2");
  EXPECT_EQ(linear_classifier_v1->type_constraints()[1].allowed_type_strs.size(), 2u);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsLinearRegressorSchema) {
  const std::vector<core::schema::LightOpSchema> linear_regressor_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("LinearRegressor");

  const core::schema::LightOpSchema *const linear_regressor_v1 =
      FindByVersion(linear_regressor_schemas, 1);
  ASSERT_NE(nullptr, linear_regressor_v1);
  EXPECT_EQ(linear_regressor_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(linear_regressor_v1->inputs().size(), 1u);
  EXPECT_EQ(linear_regressor_v1->outputs().size(), 1u);
  EXPECT_EQ(linear_regressor_v1->inputs()[0].name, "X");
  EXPECT_EQ(linear_regressor_v1->inputs()[0].type, "T");
  EXPECT_EQ(linear_regressor_v1->outputs()[0].name, "Y");
  EXPECT_EQ(linear_regressor_v1->outputs()[0].type, "tensor(float)");
  EXPECT_EQ(linear_regressor_v1->type_constraints().size(), 1u);
  EXPECT_EQ(linear_regressor_v1->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(linear_regressor_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsImputerSchema) {
  const std::vector<core::schema::LightOpSchema> imputer_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("Imputer");

  const core::schema::LightOpSchema *const imputer_v1 = FindByVersion(imputer_schemas, 1);
  ASSERT_NE(nullptr, imputer_v1);
  EXPECT_EQ(imputer_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(imputer_v1->inputs().size(), 1u);
  EXPECT_EQ(imputer_v1->outputs().size(), 1u);
  EXPECT_EQ(imputer_v1->inputs()[0].name, "X");
  EXPECT_EQ(imputer_v1->inputs()[0].type, "T");
  EXPECT_EQ(imputer_v1->outputs()[0].name, "Y");
  EXPECT_EQ(imputer_v1->outputs()[0].type, "T");
  EXPECT_EQ(imputer_v1->type_constraints().size(), 1u);
  EXPECT_EQ(imputer_v1->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(imputer_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
  EXPECT_EQ(imputer_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kFloat);
  EXPECT_EQ(imputer_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kDouble);
  EXPECT_EQ(imputer_v1->type_constraints()[0].allowed_type_strs[2],
            core::schema::TensorType::kInt64);
  EXPECT_EQ(imputer_v1->type_constraints()[0].allowed_type_strs[3],
            core::schema::TensorType::kInt32);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsCategoryMapperSchema) {
  const std::vector<core::schema::LightOpSchema> category_mapper_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("CategoryMapper");

  const core::schema::LightOpSchema *const category_mapper_v1 =
      FindByVersion(category_mapper_schemas, 1);
  ASSERT_NE(nullptr, category_mapper_v1);
  EXPECT_EQ(category_mapper_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(category_mapper_v1->inputs().size(), 1u);
  EXPECT_EQ(category_mapper_v1->outputs().size(), 1u);
  EXPECT_EQ(category_mapper_v1->inputs()[0].name, "X");
  EXPECT_EQ(category_mapper_v1->inputs()[0].type, "T1");
  EXPECT_EQ(category_mapper_v1->outputs()[0].name, "Y");
  EXPECT_EQ(category_mapper_v1->outputs()[0].type, "T2");
  EXPECT_EQ(category_mapper_v1->type_constraints().size(), 2u);
  EXPECT_EQ(category_mapper_v1->type_constraints()[0].type_param_str, "T1");
  EXPECT_EQ(category_mapper_v1->type_constraints()[0].allowed_type_strs.size(), 2u);
  EXPECT_EQ(category_mapper_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kString);
  EXPECT_EQ(category_mapper_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kInt64);
  EXPECT_EQ(category_mapper_v1->type_constraints()[1].type_param_str, "T2");
  EXPECT_EQ(category_mapper_v1->type_constraints()[1].allowed_type_strs.size(), 2u);
  EXPECT_EQ(category_mapper_v1->type_constraints()[1].allowed_type_strs[0],
            core::schema::TensorType::kString);
  EXPECT_EQ(category_mapper_v1->type_constraints()[1].allowed_type_strs[1],
            core::schema::TensorType::kInt64);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsNormalizerSchema) {
  const std::vector<core::schema::LightOpSchema> normalizer_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("Normalizer");

  const core::schema::LightOpSchema *const normalizer_v1 = FindByVersion(normalizer_schemas, 1);
  ASSERT_NE(nullptr, normalizer_v1);
  EXPECT_EQ(normalizer_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(normalizer_v1->inputs().size(), 1u);
  EXPECT_EQ(normalizer_v1->outputs().size(), 1u);
  EXPECT_EQ(normalizer_v1->inputs()[0].name, "X");
  EXPECT_EQ(normalizer_v1->inputs()[0].type, "T");
  EXPECT_EQ(normalizer_v1->outputs()[0].name, "Y");
  EXPECT_EQ(normalizer_v1->outputs()[0].type, "tensor(float)");
  EXPECT_EQ(normalizer_v1->type_constraints().size(), 1u);
  EXPECT_EQ(normalizer_v1->type_constraints()[0].type_param_str, "T");
  EXPECT_EQ(normalizer_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
  EXPECT_EQ(normalizer_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kFloat);
  EXPECT_EQ(normalizer_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kDouble);
  EXPECT_EQ(normalizer_v1->type_constraints()[0].allowed_type_strs[2],
            core::schema::TensorType::kInt64);
  EXPECT_EQ(normalizer_v1->type_constraints()[0].allowed_type_strs[3],
            core::schema::TensorType::kInt32);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsDictVectorizerSchema) {
  const std::vector<core::schema::LightOpSchema> dv_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("DictVectorizer");
  const core::schema::LightOpSchema *const dv_v1 = FindByVersion(dv_schemas, 1);
  ASSERT_NE(nullptr, dv_v1);
  EXPECT_EQ(dv_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(dv_v1->inputs().size(), 1u);
  EXPECT_EQ(dv_v1->outputs().size(), 1u);
  EXPECT_EQ(dv_v1->inputs()[0].name, "X");
  EXPECT_EQ(dv_v1->inputs()[0].type, "T1");
  EXPECT_EQ(dv_v1->outputs()[0].name, "Y");
  EXPECT_EQ(dv_v1->outputs()[0].type, "T2");
  EXPECT_EQ(dv_v1->type_constraints().size(), 2u);
  EXPECT_EQ(dv_v1->type_constraints()[0].type_param_str, "T1");
  EXPECT_EQ(dv_v1->type_constraints()[1].type_param_str, "T2");
  // T1 lists the 6 supported map (key,value) types.
  EXPECT_EQ(dv_v1->type_constraints()[0].allowed_type_strs.size(), 6u);
  // T2 lists the 4 supported value tensor types.
  EXPECT_EQ(dv_v1->type_constraints()[1].allowed_type_strs.size(), 4u);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsFeatureVectorizerSchema) {
  const std::vector<core::schema::LightOpSchema> fv_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("FeatureVectorizer");
  const core::schema::LightOpSchema *const fv_v1 = FindByVersion(fv_schemas, 1);
  ASSERT_NE(nullptr, fv_v1);
  EXPECT_EQ(fv_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(fv_v1->inputs().size(), 1u);
  EXPECT_EQ(fv_v1->outputs().size(), 1u);
  EXPECT_EQ(fv_v1->inputs()[0].name, "X");
  EXPECT_EQ(fv_v1->inputs()[0].type, "T1");
  EXPECT_EQ(fv_v1->outputs()[0].name, "Y");
  EXPECT_EQ(fv_v1->outputs()[0].type, "tensor(float)");
  EXPECT_EQ(fv_v1->type_constraints().size(), 1u);
  EXPECT_EQ(fv_v1->type_constraints()[0].type_param_str, "T1");
  EXPECT_EQ(fv_v1->type_constraints()[0].allowed_type_strs.size(), 4u);
}

TEST(OnnxOpTraditionalMLRegistrationTest, ReturnsCastMapSchema) {
  const std::vector<core::schema::LightOpSchema> cm_schemas =
      onnx_op::traditionalml::GetAllOnnxOpTraditionalMLSchemasWithHistory("CastMap");
  const core::schema::LightOpSchema *const cm_v1 = FindByVersion(cm_schemas, 1);
  ASSERT_NE(nullptr, cm_v1);
  EXPECT_EQ(cm_v1->domain(), "ai.onnx.ml");
  EXPECT_EQ(cm_v1->inputs().size(), 1u);
  EXPECT_EQ(cm_v1->outputs().size(), 1u);
  EXPECT_EQ(cm_v1->inputs()[0].name, "X");
  EXPECT_EQ(cm_v1->inputs()[0].type, "T1");
  EXPECT_EQ(cm_v1->outputs()[0].name, "Y");
  EXPECT_EQ(cm_v1->outputs()[0].type, "T2");
  EXPECT_EQ(cm_v1->type_constraints().size(), 2u);
  EXPECT_EQ(cm_v1->type_constraints()[0].type_param_str, "T1");
  EXPECT_EQ(cm_v1->type_constraints()[1].type_param_str, "T2");
  // T1 lists the 2 supported map(int64, value) types.
  ASSERT_EQ(cm_v1->type_constraints()[0].allowed_type_strs.size(), 2u);
  EXPECT_EQ(cm_v1->type_constraints()[0].allowed_type_strs[0],
            core::schema::TensorType::kMapInt64String);
  EXPECT_EQ(cm_v1->type_constraints()[0].allowed_type_strs[1],
            core::schema::TensorType::kMapInt64Float);
  // T2 lists the 3 supported output tensor types.
  ASSERT_EQ(cm_v1->type_constraints()[1].allowed_type_strs.size(), 3u);
  EXPECT_EQ(cm_v1->type_constraints()[1].allowed_type_strs[0], core::schema::TensorType::kString);
  EXPECT_EQ(cm_v1->type_constraints()[1].allowed_type_strs[1], core::schema::TensorType::kFloat);
  EXPECT_EQ(cm_v1->type_constraints()[1].allowed_type_strs[2], core::schema::TensorType::kInt64);
}

} // namespace Test
