/**
 * Copyright 2019-2020 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 <gtest/gtest.h>
#include <iostream>
#include "parser/common/op_parser_factory.h"
#include "graph/operator_reg.h"
#include "external/graph/types.h"
#include "register/op_registry.h"
#include "parser/common/register_tbe.h"
#include "external/parser/onnx_parser.h"
#include "ut/parser/parser_ut_utils.h"
#include "external/ge/ge_api_types.h"
#include "tests/depends/ops_stub/ops_stub.h"

#define protected public
#define private public
#include "parser/onnx/onnx_constant_parser.h"
#include "parser/onnx/onnx_util.h"
#include "parser/onnx/onnx_parser.h"
#undef protected
#undef private

namespace ge {
class UtestOnnxParser : public testing::Test {
 protected:
  void SetUp() {
    ParerUTestsUtils::ClearParserInnerCtx();
    RegisterCustomOp();
  }

  void TearDown() {}

 public:
  void RegisterCustomOp();
};

static Status ParseParams(const google::protobuf::Message* op_src, ge::Operator& op_dest) {
  return SUCCESS;
}

static Status ParseParamByOpFunc(const ge::Operator &op_src, ge::Operator& op_dest) {
  return SUCCESS;
}

Status ParseSubgraphPostFnIf(const std::string& subgraph_name, const ge::Graph& graph) {
  domi::AutoMappingSubgraphIOIndexFunc auto_mapping_subgraph_index_func =
      domi::FrameworkRegistry::Instance().GetAutoMappingSubgraphIOIndexFunc(domi::ONNX);
  if (auto_mapping_subgraph_index_func == nullptr) {
    std::cout<<"auto mapping if subgraph func is nullptr!"<<std::endl;
    return FAILED;
  }
  return auto_mapping_subgraph_index_func(graph,
                                          [&](int data_index, int &parent_index) -> Status {
                                            parent_index = data_index + 1;
                                            return SUCCESS;
                                          },
                                          [&](int output_index, int &parent_index) -> Status {
                                            parent_index = output_index;
                                            return SUCCESS;
                                          });
}

void UtestOnnxParser::RegisterCustomOp() {
  REGISTER_CUSTOM_OP("Conv2D")
  .FrameworkType(domi::ONNX)
  .OriginOpType("ai.onnx::11::Conv")
  .ParseParamsFn(ParseParams);

  // register if op info to GE
  REGISTER_CUSTOM_OP("If")
  .FrameworkType(domi::ONNX)
  .OriginOpType({"ai.onnx::9::If",
                 "ai.onnx::10::If",
                 "ai.onnx::11::If",
                 "ai.onnx::12::If",
                 "ai.onnx::13::If"})
  .ParseParamsFn(ParseParams)
  .ParseParamsByOperatorFn(ParseParamByOpFunc)
  .ParseSubgraphPostFn(ParseSubgraphPostFnIf);

  REGISTER_CUSTOM_OP("Add")
  .FrameworkType(domi::ONNX)
      .OriginOpType("ai.onnx::11::Add")
      .ParseParamsFn(ParseParams);

  REGISTER_CUSTOM_OP("Identity")
  .FrameworkType(domi::ONNX)
      .OriginOpType("ai.onnx::11::Identity")
      .ParseParamsFn(ParseParams);
  std::vector<OpRegistrationData> reg_datas = domi::OpRegistry::Instance()->registrationDatas;
  for (auto reg_data : reg_datas) {
    OpRegistrationTbe::Instance()->Finalize(reg_data);
    domi::OpRegistry::Instance()->Register(reg_data);
  }
  domi::OpRegistry::Instance()->registrationDatas.clear();
}

TEST_F(UtestOnnxParser, onnx_parser_if_node) {
  std::string case_dir = __FILE__;
  case_dir = case_dir.substr(0, case_dir.find_last_of("/"));
  std::string model_file = case_dir + "/onnx_model/if.onnx";
  std::map<ge::AscendString, ge::AscendString> parser_params;
  ge::Graph graph;
  auto ret = ge::aclgrphParseONNX(model_file.c_str(), parser_params, graph);
  EXPECT_EQ(ret, GRAPH_SUCCESS);
}

TEST_F(UtestOnnxParser, onnx_parser_user_output_with_name_and_index) {
  std::string case_dir = __FILE__;
  case_dir = case_dir.substr(0, case_dir.find_last_of("/"));
  std::string model_file = case_dir + "/onnx_model/conv2d.onnx";
  std::map<ge::AscendString, ge::AscendString> parser_params;
  parser_params.insert({AscendString(ge::ir_option::OUT_NODES), AscendString("Conv_0:0")});
  ge::Graph graph;
  auto ret = ge::aclgrphParseONNX(model_file.c_str(), parser_params, graph);
  ASSERT_EQ(ret, GRAPH_SUCCESS);
  ge::ComputeGraphPtr compute_graph = ge::GraphUtils::GetComputeGraph(graph);
  auto output_nodes_info = compute_graph->GetGraphOutNodesInfo();
  ASSERT_EQ(output_nodes_info.size(), 1);
  EXPECT_EQ((output_nodes_info.at(0).first->GetName()), "Conv_0");
  EXPECT_EQ((output_nodes_info.at(0).second), 0);
  auto &net_out_name = ge::GetParserContext().net_out_nodes;
  ASSERT_EQ(net_out_name.size(), 1);
  EXPECT_EQ(net_out_name.at(0), "Conv_0:0");
}

TEST_F(UtestOnnxParser, onnx_parser_user_output_with_tensor) {
  std::string case_dir = __FILE__;
  case_dir = case_dir.substr(0, case_dir.find_last_of("/"));
  std::string model_file = case_dir + "/onnx_model/conv2d.onnx";
  std::map<ge::AscendString, ge::AscendString> parser_params;
  parser_params.insert({AscendString(ge::ir_option::OUT_NODES), AscendString("y")});
  ge::Graph graph;
  auto ret = ge::aclgrphParseONNX(model_file.c_str(), parser_params, graph);
  ASSERT_EQ(ret, GRAPH_SUCCESS);
  ge::ComputeGraphPtr compute_graph = ge::GraphUtils::GetComputeGraph(graph);
  auto output_nodes_info = compute_graph->GetGraphOutNodesInfo();
  ASSERT_EQ(output_nodes_info.size(), 1);
  EXPECT_EQ((output_nodes_info.at(0).first->GetName()), "Conv_0");
  EXPECT_EQ((output_nodes_info.at(0).second), 0);
  auto &net_out_name = ge::GetParserContext().net_out_nodes;
  ASSERT_EQ(net_out_name.size(), 1);
  EXPECT_EQ(net_out_name.at(0), "Conv_0:0:y");
}

TEST_F(UtestOnnxParser, onnx_parser_user_output_with_default) {
  std::string case_dir = __FILE__;
  case_dir = case_dir.substr(0, case_dir.find_last_of("/"));
  std::string model_file = case_dir + "/onnx_model/conv2d.onnx";
  std::map<ge::AscendString, ge::AscendString> parser_params;
  ge::Graph graph;
  auto ret = ge::aclgrphParseONNX(model_file.c_str(), parser_params, graph);
  ASSERT_EQ(ret, GRAPH_SUCCESS);
  ge::ComputeGraphPtr compute_graph = ge::GraphUtils::GetComputeGraph(graph);
  auto output_nodes_info = compute_graph->GetGraphOutNodesInfo();
  ASSERT_EQ(output_nodes_info.size(), 1);
  EXPECT_EQ((output_nodes_info.at(0).first->GetName()), "Conv_0");
  EXPECT_EQ((output_nodes_info.at(0).second), 0);
  auto &net_out_name = ge::GetParserContext().net_out_nodes;
  ASSERT_EQ(net_out_name.size(), 1);
  EXPECT_EQ(net_out_name.at(0), "Conv_0:0:y");
}

TEST_F(UtestOnnxParser, onnx_parser_user_output_with_tensor_failed) {
  std::string case_dir = __FILE__;
  case_dir = case_dir.substr(0, case_dir.find_last_of("/"));
  std::string model_file = case_dir + "/onnx_model/conv2d.onnx";
  std::map<ge::AscendString, ge::AscendString> parser_params;
  parser_params.insert({AscendString(ge::ir_option::OUT_NODES), AscendString("not_exist_output")});
  ge::Graph graph;
  auto ret = ge::aclgrphParseONNX(model_file.c_str(), parser_params, graph);
  EXPECT_EQ(ret, FAILED);
}

TEST_F(UtestOnnxParser, onnx_parser_expand_one_to_many) {
  std::string case_dir = __FILE__;
  case_dir = case_dir.substr(0, case_dir.find_last_of("/"));
  std::string model_file = case_dir + "/onnx_model/onnx_clip_v9.onnx";
  std::map<ge::AscendString, ge::AscendString> parser_params;
  ge::Graph graph;
  auto ret = ge::aclgrphParseONNX(model_file.c_str(), parser_params, graph);
  EXPECT_EQ(ret, GRAPH_SUCCESS);

  MemBuffer *buffer = ParerUTestsUtils::MemBufferFromFile(model_file.c_str());
  ret = ge::aclgrphParseONNXFromMem(reinterpret_cast<char *>(buffer->data), buffer->size, parser_params, graph);
  EXPECT_EQ(ret, GRAPH_SUCCESS);
}

TEST_F(UtestOnnxParser, onnx_parser_to_json) {
  std::string case_dir = __FILE__;
  case_dir = case_dir.substr(0, case_dir.find_last_of("/"));
  std::string model_file = case_dir + "/onnx_model/onnx_clip_v9.onnx";
  std::map<ge::AscendString, ge::AscendString> parser_params;
  OnnxModelParser onnx_parser;

  const char *json_file = "tmp.json";
  auto ret = onnx_parser.ToJson(model_file.c_str(), json_file);
  EXPECT_EQ(ret, SUCCESS);

  const char *json_null = nullptr;
  ret = onnx_parser.ToJson(model_file.c_str(), json_null);
  EXPECT_EQ(ret, FAILED);
  const char *model_null = nullptr;
  ret = onnx_parser.ToJson(model_null, json_null);
  EXPECT_EQ(ret, FAILED);
}

TEST_F(UtestOnnxParser, onnx_parser_const_data_type) {
  std::string case_dir = __FILE__;
  case_dir = case_dir.substr(0, case_dir.find_last_of("/"));
  std::string model_file = case_dir + "/onnx_model/onnx_const_type.onnx";
  std::map<ge::AscendString, ge::AscendString> parser_params;
  ge::Graph graph;
  auto ret = ge::aclgrphParseONNX(model_file.c_str(), parser_params, graph);
  EXPECT_EQ(ret, GRAPH_SUCCESS);
}

TEST_F(UtestOnnxParser, OnnxModelParser_ConvertToGeDataType_test)
{
  OnnxModelParser model_parser;
  uint32_t type = OnnxDataType::FLOAT;

  Status ret = model_parser.ConvertToGeDataType(type);
  EXPECT_EQ(ret, SUCCESS);

  type = 20;
  ret = model_parser.ConvertToGeDataType(type);
  EXPECT_EQ(ret, ge::DataType::DT_UNDEFINED);
}

TEST_F(UtestOnnxParser, OnnxModelParser_ParseConvertData_test)
{
  OnnxConstantParser constant_parser;
  ge::onnx::TensorProto tensor_proto;
  tensor_proto.set_data_type(ge::DataType::DT_UNDEFINED);

  ge::Tensor tensor;
  int count = 1;

  Status ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
  EXPECT_EQ(ret, FAILED);

  tensor_proto.set_data_type(OnnxDataType::INT32);
  count = 0;
  ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
  EXPECT_EQ(ret, SUCCESS);

  tensor_proto.set_data_type(OnnxDataType::BFLOAT16);
  count = 1;
  ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
  EXPECT_EQ(ret, FAILED);

  tensor_proto.set_data_type(OnnxDataType::STRING);
  ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
  EXPECT_EQ(ret, FAILED);

  tensor_proto.set_raw_data("Test");
  ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
  EXPECT_EQ(ret, SUCCESS);

  tensor_proto.set_data_type(OnnxDataType::FLOAT);
  ret = constant_parser.ParseConvertData(tensor_proto, tensor, count);
  EXPECT_EQ(ret, SUCCESS);
}

TEST_F(UtestOnnxParser, OnnxConstantParser_ParseConvertTensor_test)
{
  OnnxConstantParser constant_parser;
  ge::onnx::NodeProto input_node;
  ge::onnx::AttributeProto *attribute = input_node.add_attribute();
  attribute->set_name("attribute");
  attribute->set_type(onnx::AttributeProto::AttributeType(1));
  attribute->set_f(1.0);
  ge::onnx::TensorProto *attribute_tensor = attribute->mutable_t();
  attribute_tensor->set_data_type(1);
  attribute_tensor->add_dims(4);

  ge::Tensor tensor;
  Status ret = constant_parser.ParseConvertTensor(*attribute_tensor, tensor);
  EXPECT_EQ(ret, FAILED);

  attribute_tensor->add_dims(-1);
  ret = constant_parser.ParseConvertTensor(*attribute_tensor, tensor);
  EXPECT_EQ(ret, FAILED);
}

TEST_F(UtestOnnxParser, OnnxConstantParser_ParseConvertDataType_test)
{
  OnnxConstantParser constant_parser;
  ge::onnx::NodeProto input_node;
  ge::onnx::AttributeProto *attribute = input_node.add_attribute();
  attribute->set_name("attribute");
  attribute->set_type(onnx::AttributeProto::AttributeType(1));
  attribute->set_f(1.0);
  ge::onnx::TensorProto *attribute_tensor = attribute->mutable_t();
  attribute_tensor->set_data_type(OnnxDataType::BFLOAT16);
  attribute_tensor->add_dims(4);

  ge::Tensor tensor;
  Status ret = constant_parser.ParseConvertDataType(*attribute_tensor, tensor);
  EXPECT_EQ(ret, FAILED);
}

TEST_F(UtestOnnxParser, OnnxModelParser_ParseInput_test)
{
  OnnxModelParser model_parser;
  ge::onnx::ModelProto model_proto;
  ge::onnx::GraphProto graph = model_proto.graph();
  std::map<std::string, ge::onnx::TensorProto> initializer_name_tensor;
  bool is_subgraph = false;

  Status ret = model_parser.ParseInput(initializer_name_tensor, is_subgraph, graph);
  EXPECT_EQ(ret, domi::FAILED);

  ret = model_parser.ParseOutput(graph);
  EXPECT_EQ(ret, domi::FAILED);
}

TEST_F(UtestOnnxParser, onnx_test_ConstructOriType)
{
  ge::onnx::ModelProto model_proto;
  ge::onnx::GraphProto* graph = model_proto.mutable_graph();
  ge::onnx::NodeProto* add_node = graph->add_node();
  add_node->set_op_type("Add");
  add_node->set_domain("ai.onnx");

  OnnxModelParser onnx_parser ;
  onnx_parser.domain_verseion_["ai.onnx"] = 11;
  string ori_type;
  Status ret = onnx_parser.ConstructOriType(add_node, ori_type);
  EXPECT_EQ(ret, domi::SUCCESS);

  ge::onnx::NodeProto* add_node1 = graph->add_node();
  add_node1->set_op_type("Add1");
  add_node1->set_domain("add.onnx");
  string op_type;
  ret = onnx_parser.AdapterOpType(add_node1, ori_type, op_type);
  EXPECT_EQ(ret, ge::PARAM_INVALID);

  add_node->set_op_type("Add1");
  ret = onnx_parser.AdapterOpType(add_node, ori_type, op_type);
  EXPECT_EQ(ret, PARAM_INVALID);
}

TEST_F(UtestOnnxParser, onnx_test_TransNodeToOperator)
{
  ge::onnx::ModelProto model_proto;
  ge::onnx::GraphProto* graph = model_proto.mutable_graph();
  ge::onnx::NodeProto *node_proto = graph->add_node();
  node_proto->set_op_type("Add1");
  node_proto->set_domain("add.onnx");
  node_proto->set_name("Conv2D");
  ge::OpDescPtr op_desc_src = std::make_shared<ge::OpDesc>("Add", "add.onnx");
  ge::Operator op = ge::OpDescUtils::CreateOperatorFromOpDesc(op_desc_src);
  std::string op_type = "Add";

  OnnxModelParser onnx_parser;
  Status ret = onnx_parser.TransNodeToOperator(node_proto, op, op_type);
  EXPECT_EQ(ret, SUCCESS);
}

} // namespace ge
