/**
 * 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 <operator_factory_impl.h>

#define protected public
#define private public
#include "graph/passes/infershape_pass.h"

#include "graph/utils/tensor_utils.h"
#include "graph/utils/graph_utils.h"
#include "graph_builder_utils.h"
#include "inc/external/graph/operator_reg.h"
#include "inc/external/graph/operator.h"
#include "inc/external/graph/operator_factory.h"
#include "inc/graph/operator_factory_impl.h"


using namespace std;
using namespace testing;
namespace ge {
class UtestGraphInfershapePass : public testing::Test {
 protected:
  void SetUp() {}
  void TearDown() {}
};

/*
 *   data1  const1
 *     \     /
 *      case1
 *        |
 *      relu10
 *        |
 *    netoutput
 */
ut::GraphBuilder ParentGraphBuilder() {
  ut::GraphBuilder builder = ut::GraphBuilder("g1");
  auto data1 = builder.AddNode("data1", "Data", 0, 1);
  std::vector<int64_t> const_shape = {1};
  auto const1 = builder.AddNode("const1", "Const", 0, 1, FORMAT_NCHW, DT_INT32, const_shape);
  auto case1 = builder.AddNode("case1", CASE, 2, 1);
  auto relu1 = builder.AddNode("relu10", "Relu", 1, 1);
  auto netoutput = builder.AddNode("netoutput", NETOUTPUT, 1, 0);

  int32_t weight[1] = {1};
  GeTensorDesc weight_desc(GeShape({1}), FORMAT_NHWC, DT_INT32);
  GeTensorPtr tensor = std::make_shared<GeTensor>(weight_desc, (uint8_t *)weight, sizeof(weight));
  OpDescUtils::SetWeights(const1, {tensor});

  builder.AddDataEdge(data1, 0, case1, 0);
  builder.AddDataEdge(const1, 0, case1, 1);
  builder.AddDataEdge(case1, 0, relu1, 0);
  builder.AddDataEdge(relu1, 0, netoutput, 0);
  return builder;
}

/*
 *   data1   data2
 *     \      /
 *      switch
 *     /      \
 *   relu1   relu2
 *     \      /
 *       merge
 *         |
 *     netoutput
 */
ut::GraphBuilder SwitchSubgraphBuilder(string graph_name, uint32_t num) {
  ut::GraphBuilder builder = ut::GraphBuilder(graph_name);

  std::vector<int64_t> shape1 = {2,2};
  string data1_name = "data1_" + std::to_string(num);
  auto data1 = builder.AddNode(data1_name, "Data", 1, 1, FORMAT_NCHW, DT_INT32, shape1);
  auto data1_desc = data1->GetOpDesc();
  EXPECT_NE(data1_desc, nullptr);
  AttrUtils::SetInt(data1_desc, "_parent_node_index", 0);

  std::vector<int64_t> shape2 = {3,3};
  string data2_name = "data2_" + std::to_string(num);
  auto data2 = builder.AddNode(data2_name, "Data", 1, 1, FORMAT_NCHW, DT_INT32, shape2);
  auto data2_desc = data2->GetOpDesc();
  EXPECT_NE(data2_desc, nullptr);
  AttrUtils::SetInt(data2_desc, "_parent_node_index", 1);

  string switch_name = "switch_" + std::to_string(num);
  auto switch1 = builder.AddNode(switch_name, "Switch", 2, 2);

  string relu1_name = "relu1_" + std::to_string(num);
  auto relu1 = builder.AddNode(relu1_name, "Relu", 1, 1);

  string relu2_name = "relu2_" + std::to_string(num);
  auto relu2 = builder.AddNode(relu2_name, "Relu", 1, 1);

  string merge_name = "merge_" + std::to_string(num);
  auto merge = builder.AddNode(merge_name, "Merge", 2, 1);

  std::vector<int64_t> shape7 = {8,8};
  string output_name = "output_" + std::to_string(num);
  auto netoutput = builder.AddNode(output_name, NETOUTPUT, 1, 0, FORMAT_NCHW, DT_INT32, shape7);
  auto input0_desc = netoutput->GetOpDesc()->MutableInputDesc(0);
  EXPECT_NE(input0_desc, nullptr);
  AttrUtils::SetInt(input0_desc, "_parent_node_index", 0);

  builder.AddDataEdge(data1, 0, switch1, 0);
  builder.AddDataEdge(data2, 0, switch1, 1);
  builder.AddDataEdge(switch1, 0, relu1, 0);
  builder.AddDataEdge(switch1, 1, relu2, 0);
  builder.AddDataEdge(relu1, 0, merge, 0);
  builder.AddDataEdge(relu2, 0, merge, 1);
  builder.AddDataEdge(merge, 0, netoutput, 0);

  return builder;
}

void AddCaseSubgraph(ComputeGraphPtr &parent_graph, uint32_t branch_num) {
  auto case_node = parent_graph->FindNode("case1");
  EXPECT_NE(case_node, nullptr);

  for (uint32_t i = 0; i < branch_num; ++i) {
    string name = "Branch_Graph_" + std::to_string(i);

    auto builder_subgraph = SwitchSubgraphBuilder(name, i);
    auto switch_subgraph = builder_subgraph.GetGraph();

    case_node->GetOpDesc()->AddSubgraphName(switch_subgraph->GetName());
    case_node->GetOpDesc()->SetSubgraphInstanceName(i, switch_subgraph->GetName());

    switch_subgraph->SetParentNode(case_node);
    switch_subgraph->SetParentGraph(parent_graph);
    EXPECT_EQ(parent_graph->AddSubgraph(switch_subgraph->GetName(), switch_subgraph), GRAPH_SUCCESS);
  }
}

static NodePtr CreateNode(ComputeGraph &graph, const string &name, const string &type, int in_num, int out_num) {
  OpDescPtr op_desc = std::make_shared<OpDesc>(name, type);
  op_desc->SetStreamId(0);
  static int32_t index = 0;
  op_desc->SetId(index++);

  GeTensorDesc tensor(GeShape(), FORMAT_NCHW, DT_FLOAT);
  TensorUtils::SetSize(tensor, 512);
  vector<int64_t> input_offset;
  for (int i = 0; i < in_num; i++) {
    op_desc->AddInputDesc(tensor);
    input_offset.emplace_back(1024);
  }
  op_desc->SetInputOffset(input_offset);

  vector<int64_t> output_offset;
  for (int i = 0; i < out_num; i++) {
    op_desc->AddOutputDesc(tensor);
    output_offset.emplace_back(1024);
  }
  op_desc->SetOutputOffset(output_offset);

  op_desc->SetWorkspace({});
  op_desc->SetWorkspaceBytes({});
  op_desc->SetOpKernelLibName("DNN_VM_RTS_OP_STORE");

  const auto stub_func = [](Operator &op) { return GRAPH_SUCCESS; };
  op_desc->AddInferFunc(stub_func);
  op_desc->AddInferFormatFunc(stub_func);
  op_desc->AddVerifierFunc(stub_func);

  return graph.AddNode(op_desc);
}

TEST_F(UtestGraphInfershapePass, infershape_pass_failed) {
  GeTensorDesc ge_tensor_desc(GeShape({-2, 2, 3, 4}), ge::FORMAT_NCHW, DT_FLOAT16);
  string type = "AddN";
  auto addn_op_desc = std::make_shared<OpDesc>("AddN", type);
  addn_op_desc->AddInputDesc(ge_tensor_desc);
  addn_op_desc->AddOutputDesc(ge_tensor_desc);
  auto graph = std::make_shared<ComputeGraph>("test");
  auto addn_node = std::make_shared<Node>(addn_op_desc, graph);
  addn_node->Init();

  InferShapePass infershape_pass;
  EXPECT_EQ(infershape_pass.Run(addn_node), GE_GRAPH_INFERSHAPE_FAILED);
}

TEST_F(UtestGraphInfershapePass, delete_need_infer_again) {
  auto graph = std::make_shared<ComputeGraph>("test");

  auto no_op_desc = std::make_shared<OpDesc>("No", "NoOp");
  auto no_op_node = graph->AddNode(no_op_desc);
  AttrUtils::SetBool(no_op_desc, "_need_infer_again", false);

  InferShapePass infershape_pass;
  infershape_pass.options_[kOptimizeAfterSubGraph] = "yes";
  EXPECT_EQ(infershape_pass.Run(no_op_node), SUCCESS);
}

TEST_F(UtestGraphInfershapePass, stop_node_for_while_loop) {
/*******************************************************************************
 *      Exit         Identify
 *        \         /       \.
 *         \       /         \.
 *          Switch           Add
 *         /     |            |
 *        /      |            |
 *       /       |            |
 *  LoopCond     |            |
 *      \        |            |
 *       \       |            |
 *        \      |            |
 *       Less    |            |
 *          \    |       NextIteration
 *           \   |            |
 *            \  |            |
 *            Merge <---------|
 *              |
 *              |
 *            Enter
 ******************************************************************************/
  auto graph = std::make_shared<ComputeGraph>("test_infer_shape");
  auto data1 = CreateNode(*graph, "data", DATA, 1, 1);
  auto enter1 = CreateNode(*graph, "enter", ENTER, 1, 1);
  auto merge1 = CreateNode(*graph, "merge", MERGE, 2, 2);
  auto less1 = CreateNode(*graph, "less", LESS, 2, 1);
  auto loop1 = CreateNode(*graph, "loopcond", LOOPCOND, 1, 1);
  auto switch1 = CreateNode(*graph, "switch", SWITCH, 2, 2);
  auto ident1 = CreateNode(*graph, "identity", IDENTITY, 1, 1);
  auto add1 = CreateNode(*graph, "add", ADD, 2, 1);
  auto next1 = CreateNode(*graph, "next", NEXTITERATION, 1, 1);
  auto exit1 = CreateNode(*graph, "exit", EXIT, 1, 1);
  auto value0 = CreateNode(*graph, "const", CONSTANT, 0, 1);
  auto value1 = CreateNode(*graph, "const", CONSTANT, 0, 1);
  auto output1 = CreateNode(*graph, "net_output", NETOUTPUT, 1, 1);

  GraphUtils::AddEdge(data1->GetOutDataAnchor(0), enter1->GetInDataAnchor(0));
  GraphUtils::AddEdge(enter1->GetOutDataAnchor(0), merge1->GetInDataAnchor(0));
  GraphUtils::AddEdge(merge1->GetOutDataAnchor(0), less1->GetInDataAnchor(0));
  GraphUtils::AddEdge(value1->GetOutDataAnchor(0), less1->GetInDataAnchor(1));
  GraphUtils::AddEdge(less1->GetOutDataAnchor(0), loop1->GetInDataAnchor(0));

  GraphUtils::AddEdge(loop1->GetOutDataAnchor(0), switch1->GetInDataAnchor(0));
  GraphUtils::AddEdge(merge1->GetOutDataAnchor(0), switch1->GetInDataAnchor(1));

  GraphUtils::AddEdge(switch1->GetOutDataAnchor(0), exit1->GetInDataAnchor(0));
  GraphUtils::AddEdge(switch1->GetOutDataAnchor(1), ident1->GetInDataAnchor(0));

  GraphUtils::AddEdge(ident1->GetOutDataAnchor(0), add1->GetInDataAnchor(0));
  GraphUtils::AddEdge(value1->GetOutDataAnchor(0), add1->GetInDataAnchor(1));
  GraphUtils::AddEdge(add1->GetOutDataAnchor(0), next1->GetInDataAnchor(0));

  GraphUtils::AddEdge(next1->GetOutDataAnchor(0), merge1->GetInDataAnchor(1));
  GraphUtils::AddEdge(exit1->GetOutDataAnchor(0), output1->GetInDataAnchor(0));

  GEPass ge_passes(graph);
  NamesToPass names_to_passes;
  InferShapePass infer_shape_pass;
  names_to_passes.emplace_back("InferShapePass", &infer_shape_pass);

  EXPECT_EQ(ge_passes.Run(names_to_passes), SUCCESS);
}

TEST_F(UtestGraphInfershapePass, infer_with_case_subgraph) {
  auto builder = ParentGraphBuilder();
  auto parent_graph = builder.GetGraph();
  AddCaseSubgraph(parent_graph, 2);
  auto subgraphs = parent_graph->GetAllSubgraphs();
  EXPECT_EQ(subgraphs.size(), 2);

  auto case_node = parent_graph->FindNode("case1");
  EXPECT_NE(case_node, nullptr);
  InferShapePass infershape_pass;
  EXPECT_EQ(infershape_pass.Run(case_node), SUCCESS);

  std::vector<int64_t> target_dims_0 = {1, 1, 224, 224};
  std::vector<int64_t> target_dims_1 = {1};
  {
    auto data_node = subgraphs[0]->FindNode("data1_0");
    auto dims = data_node->GetOpDesc()->GetInputDescPtr(0)->GetShape().GetDims();
    EXPECT_EQ(dims, target_dims_0);
    data_node = subgraphs[0]->FindNode("data2_0");
    dims = data_node->GetOpDesc()->GetInputDescPtr(0)->GetShape().GetDims();
    EXPECT_EQ(dims, target_dims_1);
  }

  infershape_pass.options_[kOptimizeAfterSubGraph] = "yes";
  EXPECT_EQ(infershape_pass.Run(case_node), SUCCESS);
  {
    auto dims = case_node->GetOpDesc()->GetOutputDescPtr(0)->GetShape().GetDims();
    std::vector<int64_t> out_target_dims = {8, 8};
    EXPECT_EQ(out_target_dims, dims);
  }
}

/*
 *   data1  const1
 *     \     /
 *      while
 *     /    \
 * relu1  netoutput
 */
ut::GraphBuilder ParentWhileGraphBuilder() {
  ut::GraphBuilder builder = ut::GraphBuilder("g1");
  auto data1 = builder.AddNode("data1", "Data", 0, 1);
  std::vector<int64_t> const_shape = {1};
  auto const1 = builder.AddNode("const1", "Const", 0, 1, FORMAT_NCHW, DT_FLOAT, const_shape);
  auto case1 = builder.AddNode("case1", WHILE, 2, 2);
  auto relu1 = builder.AddNode("relu1", "Relu", 1, 1);
  auto netoutput = builder.AddNode("netoutput", NETOUTPUT, 1, 0);

  int32_t weight[1] = {1};
  GeTensorDesc weight_desc(GeShape({1}), FORMAT_NHWC, DT_FLOAT);
  GeTensorPtr tensor = std::make_shared<GeTensor>(weight_desc, (uint8_t *)weight, sizeof(weight));
  OpDescUtils::SetWeights(const1, {tensor});

  builder.AddDataEdge(data1, 0, case1, 0);
  builder.AddDataEdge(const1, 0, case1, 1);
  builder.AddDataEdge(case1, 0, relu1, 0);
  builder.AddDataEdge(case1, 1, netoutput, 0);
  return builder;
}

/*
 *   data1   data2
 *     \      /
 *      switch
 *      |    |
 *      \    /
 *     netoutput
 */
ut::GraphBuilder WhileSubgraphBuilder(string graph_name, uint32_t num) {
  ut::GraphBuilder builder = ut::GraphBuilder(graph_name);

  std::vector<int64_t> shape1 = {2,2};
  string data1_name = "data1_" + std::to_string(num);
  auto data1 = builder.AddNode(data1_name, "Data", 1, 1, FORMAT_NCHW, DT_FLOAT, shape1);
  auto data1_desc = data1->GetOpDesc();
  EXPECT_NE(data1_desc, nullptr);
  AttrUtils::SetInt(data1_desc, "_parent_node_index", 0);

  std::vector<int64_t> shape2 = {3,3};
  string data2_name = "data2_" + std::to_string(num);
  auto data2 = builder.AddNode(data2_name, "Data", 1, 1, FORMAT_NCHW, DT_FLOAT, shape2);
  auto data2_desc = data2->GetOpDesc();
  EXPECT_NE(data2_desc, nullptr);
  AttrUtils::SetInt(data2_desc, "_parent_node_index", 1);

  string switch_name = "switch_" + std::to_string(num);
  auto switch1 = builder.AddNode(switch_name, "Switch", 2, 2);

  std::vector<int64_t> shape7 = {8,8,8,8};
  string output_name = "output_" + std::to_string(num);
  auto netoutput = builder.AddNode(output_name, NETOUTPUT, 2, 0, FORMAT_NCHW, DT_FLOAT, shape7);
  auto input0_desc = netoutput->GetOpDesc()->MutableInputDesc(0);
  EXPECT_NE(input0_desc, nullptr);
  AttrUtils::SetInt(input0_desc, "_parent_node_index", 0);
  auto input1_desc = netoutput->GetOpDesc()->MutableInputDesc(1);
  EXPECT_NE(input1_desc, nullptr);
  AttrUtils::SetInt(input1_desc, "_parent_node_index", 1);

  builder.AddDataEdge(data1, 0, switch1, 0);
  builder.AddDataEdge(data2, 0, switch1, 1);
  builder.AddDataEdge(switch1, 0, netoutput, 0);
  builder.AddDataEdge(switch1, 1, netoutput, 1);
  return builder;
}

void AddWhileSubgraph(ComputeGraphPtr &parent_graph, uint32_t branch_num) {
  auto case_node = parent_graph->FindNode("case1");
  EXPECT_NE(case_node, nullptr);

  for (uint32_t i = 0; i < branch_num; ++i) {
    string name = "Branch_Graph_" + std::to_string(i);

    auto builder_subgraph = WhileSubgraphBuilder(name, i);
    auto switch_subgraph = builder_subgraph.GetGraph();

    case_node->GetOpDesc()->AddSubgraphName(switch_subgraph->GetName());
    case_node->GetOpDesc()->SetSubgraphInstanceName(i, switch_subgraph->GetName());

    switch_subgraph->SetParentNode(case_node);
    switch_subgraph->SetParentGraph(parent_graph);
    EXPECT_EQ(parent_graph->AddSubgraph(switch_subgraph->GetName(), switch_subgraph), GRAPH_SUCCESS);
  }
}

TEST_F(UtestGraphInfershapePass, infer_with_while_subgraph) {
  auto builder = ParentWhileGraphBuilder();
  auto parent_graph = builder.GetGraph();
  AddWhileSubgraph(parent_graph, 1);
  auto subgraphs = parent_graph->GetAllSubgraphs();
  EXPECT_EQ(subgraphs.size(), 1);

  auto case_node = parent_graph->FindNode("case1");
  EXPECT_NE(case_node, nullptr);
  InferShapePass infershape_pass;
  EXPECT_EQ(infershape_pass.Run(case_node), SUCCESS);

  std::vector<int64_t> target_dims_0 = {1, 1, 224, 224};
  std::vector<int64_t> target_dims_1 = {1};
  {
    auto data_node = subgraphs[0]->FindNode("data1_0");
    auto dims = data_node->GetOpDesc()->GetInputDescPtr(0)->GetShape().GetDims();
    EXPECT_EQ(dims, target_dims_0);
    data_node = subgraphs[0]->FindNode("data2_0");
    dims = data_node->GetOpDesc()->GetInputDescPtr(0)->GetShape().GetDims();
    EXPECT_EQ(dims, target_dims_1);
  }

  infershape_pass.options_[kOptimizeAfterSubGraph] = "yes";
  EXPECT_EQ(infershape_pass.Run(case_node), SUCCESS);
  {
    auto dims = case_node->GetOpDesc()->GetOutputDescPtr(0)->GetShape().GetDims();
    std::vector<int64_t> out_target_dims = {-1, -1, -1, -1};
    EXPECT_EQ(out_target_dims, dims);
  }
}

TEST_F(UtestGraphInfershapePass, infer_with_while_subgraph_failed) {
  auto builder = ParentWhileGraphBuilder();
  auto parent_graph = builder.GetGraph();
  AddWhileSubgraph(parent_graph, 2);
  auto subgraphs = parent_graph->GetAllSubgraphs();
  EXPECT_EQ(subgraphs.size(), 2);

  auto case_node = parent_graph->FindNode("case1");
  EXPECT_NE(case_node, nullptr);
  InferShapePass infershape_pass;
  infershape_pass.options_[kOptimizeAfterSubGraph] = "yes";
  EXPECT_EQ(infershape_pass.Run(case_node), GE_GRAPH_INFERSHAPE_FAILED);
}

auto InferFunc = [&](Operator &op) {
  return GRAPH_SUCCESS;
};
TEST_F(UtestGraphInfershapePass, infer_forrunning_with_while_subgraph) {
  auto builder = ParentWhileGraphBuilder();
  auto parent_graph = builder.GetGraph();
  AddWhileSubgraph(parent_graph, 1);
  auto subgraphs = parent_graph->GetAllSubgraphs();
  EXPECT_EQ(subgraphs.size(), 1);

  OperatorFactoryImpl::RegisterInferShapeFunc("Relu", InferFunc);

  auto relu_node = parent_graph->FindNode("relu1");
  EXPECT_NE(relu_node, nullptr);
  InferShapeForRunning infershape_for_running;
  EXPECT_EQ(infershape_for_running.Run(relu_node), SUCCESS);
}

TEST_F(UtestGraphInfershapePass, infer_static_func) {
  auto builder = ut::GraphBuilder("test_graph");
  auto data_1 = builder.AddNode("data_1", DATA, 0, 1);
  auto data_2 = builder.AddNode("data_2", DATA, 0, 1);
  auto add = builder.AddNode("Add", "Add", 2, 1);

  builder.AddDataEdge(data_1, 0, add, 0);
  builder.AddDataEdge(data_2, 0, add, 1);
  auto test_graph = builder.GetGraph();

//  OperatorFactoryImpl::CreateOperator("Add", "Flatten");
  auto test_node = test_graph->FindNode("Add");
  auto ret = InferShapePass::InferShapeAndType(test_node);
  EXPECT_EQ(ret, GRAPH_SUCCESS);

  OperatorFactoryImpl::RegisterInferShapeFunc("Add", InferFunc);
  ret = InferShapePass::InferShapeAndType(test_node);
  EXPECT_EQ(ret, GRAPH_SUCCESS);

  ret = InferShapePass::InferShapeAndType(test_node, true);
  EXPECT_EQ(ret, GRAPH_SUCCESS);

  ret = InferShapeForRunning::InferShapeAndTypeForRunning(test_node, true);
  EXPECT_EQ(ret, GRAPH_SUCCESS);
}
}  // namespace ge