/**
 * Copyright 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 "common/backend_common_test.h"
#include "debug/anf_ir_dump.h"
#include "common/py_func_graph_fetcher.h"
#include "backend/optimizer/ascend/format_type/remove_internal_output.h"

#define private public
#define protected public
#include "backend/optimizer/ascend/format_type/insert_trans_op.h"
#undef private
#undef protected

namespace mindspore {
namespace opt {
using KernelBuildInfoBuilder = kernel::KernelBuildInfo::KernelBuildInfoBuilder;

class TestHWRemoveInternalOutput : public BackendCommon {
 public:
  TestHWRemoveInternalOutput() : getPyFun_("gtest_input.pre_activate.remove_internal_output_test", true) {}
  ~TestHWRemoveInternalOutput() override = default;

  AnfNodePtr GetMakeTuple(const KernelGraphPtr &kg) {
    auto ret = kg->get_return();
    MS_EXCEPTION_IF_NULL(ret);
    auto make_tuple = ret->input(1);
    return make_tuple;
  }

  KernelGraphPtr GetSingleOutputGraph(const std::string &func_name, const std::string &sub_func_name) {
    FuncGraphPtr g = getPyFun_.CallAndParseRet(func_name, sub_func_name);
    std::vector<int> shp{2, 32, 224, 224};
    auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
    AbstractBasePtrList args_spec_list{x_abstract, x_abstract};
    auto kg = GetKernelGraph(g, args_spec_list);
    auto make_tuple = GetMakeTuple(kg);
    auto add = make_tuple->cast<CNodePtr>()->input(1);
    MS_EXCEPTION_IF_NULL(add);
    kg->AddInternalOutput(add, add);
    KernelBuildInfoBuilder builder;
    builder.SetInputsFormat({kOpFormat_DEFAULT, kOpFormat_DEFAULT});
    builder.SetInputsDeviceType({kFloat32->type_id(), kFloat32->type_id()});
    builder.SetInputsReshapeType({{}, {}});
    builder.SetOutputsReshapeType({{}});
    builder.SetOutputsFormat({kOpFormat_NC1HWC0});
    builder.SetOutputsDeviceType({kFloat16->type_id()});
    add->set_kernel_info(std::make_shared<device::KernelInfo>());
    AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), add.get());
    return kg;
  }

  KernelGraphPtr GetMutilpleOutputGraph(const std::string &func_name, const std::string &sub_func_name) {
    FuncGraphPtr g = getPyFun_.CallAndParseRet(func_name, sub_func_name);
    std::vector<int> shp{2, 32, 224, 224};
    auto x_abstract = std::make_shared<abstract::AbstractTensor>(kFloat32, shp);
    AbstractBasePtrList args_spec_list{x_abstract};
    auto kg = GetKernelGraph(g, args_spec_list);
    auto output_make_tuple = GetMakeTuple(kg);
    auto make_tuple = output_make_tuple->cast<CNodePtr>()->input(1);
    MS_EXCEPTION_IF_NULL(make_tuple);
    auto tuple_getitem1 = make_tuple->cast<CNodePtr>()->input(1);
    MS_EXCEPTION_IF_NULL(tuple_getitem1);
    auto tuple_getitem2 = make_tuple->cast<CNodePtr>()->input(2);
    MS_EXCEPTION_IF_NULL(tuple_getitem2);
    auto max_pool = tuple_getitem1->cast<CNodePtr>()->input(1);
    MS_EXCEPTION_IF_NULL(max_pool);
    kg->AddInternalOutput(tuple_getitem1, max_pool);
    kg->AddInternalOutput(tuple_getitem2, max_pool);
    KernelBuildInfoBuilder builder;
    builder.SetInputsReshapeType({{}});
    builder.SetOutputsReshapeType({{}, {}});
    builder.SetInputsFormat({kOpFormat_DEFAULT});
    builder.SetInputsDeviceType({kFloat32->type_id()});
    builder.SetOutputsFormat({kOpFormat_NC1HWC0, kOpFormat_NC1HWC0});
    builder.SetOutputsDeviceType({kFloat16->type_id(), kFloat16->type_id()});
    max_pool->set_kernel_info(std::make_shared<device::KernelInfo>());
    AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), max_pool.get());
    return kg;
  }
  UT::PyFuncGraphFetcher getPyFun_;
};

class MockRemoveInternalOutputTransOpKernelSelect : public KernelSelect {
 public:
  MockRemoveInternalOutputTransOpKernelSelect() = default;
  ~MockRemoveInternalOutputTransOpKernelSelect() override = default;
  void SelectKernel(const CNodePtr &cnode) override {
    KernelBuildInfoBuilder builder;
    builder.SetInputsFormat({kOpFormat_NC1HWC0});
    builder.SetInputsDeviceType({kFloat16->type_id()});
    builder.SetOutputsFormat({kOpFormat_DEFAULT});
    builder.SetOutputsDeviceType({kFloat32->type_id()});
    builder.SetInputsReshapeType({{}});
    builder.SetOutputsReshapeType({{}});
    AnfAlgo::SetSelectKernelBuildInfo(builder.Build(), cnode.get());
  }
};

TEST_F(TestHWRemoveInternalOutput, test_remove_internal_output_trans_op_for_single_output) {
  auto ms_context = MsContext::GetInstance();
  MS_EXCEPTION_IF_NULL(ms_context);
  ms_context->set_execution_mode(kGraphMode);
  auto kg = GetSingleOutputGraph("test_remove_internal_output_trans_op_for_single_output", "before");
  // insert trans op for output
  auto graph_optimizer = std::make_shared<opt::GraphOptimizer>();
  auto pass_manager = std::make_shared<opt::PassManager>();
  auto insert_trans_op_pass = std::make_shared<opt::InsertTransOp>();
  insert_trans_op_pass->kernel_select_ = std::make_shared<MockRemoveInternalOutputTransOpKernelSelect>();
  pass_manager->AddPass(insert_trans_op_pass);
  graph_optimizer->AddPassManager(pass_manager);
  auto new_g = graph_optimizer->Optimize(kg);
  FuncGraphPtr g_after =
    getPyFun_.CallAndParseRet("test_remove_internal_output_trans_op_for_single_output", "after_insert_trans_op");
  EXPECT_TRUE(CheckEqualGraph(g_after, new_g));

  auto make_tuple = GetMakeTuple(kg);
  auto trans_data = make_tuple->cast<CNodePtr>()->input(1);
  EXPECT_TRUE(kg->IsInternalOutput(trans_data, 0));

  // remove trans op for internal output
  auto graph_optimizer1 = std::make_shared<opt::GraphOptimizer>();
  auto pass_manager1 = std::make_shared<opt::PassManager>();
  auto remove_internal_output_trans_op_pass = std::make_shared<opt::RemoveInternalOutputTransOp>();
  pass_manager1->AddPass(remove_internal_output_trans_op_pass);
  graph_optimizer1->AddPassManager(pass_manager1);
  auto new_g1 = graph_optimizer1->Optimize(new_g);
  FuncGraphPtr g_after1 = getPyFun_.CallAndParseRet("test_remove_internal_output_trans_op_for_single_output",
                                                    "after_remove_internal_output_trans_op");
  EXPECT_TRUE(CheckEqualGraph(g_after1, new_g1));
}

TEST_F(TestHWRemoveInternalOutput, test_remove_internal_output_trans_op_for_multiple_output) {
  auto kg = GetMutilpleOutputGraph("test_remove_internal_output_trans_op_for_multiple_output", "before");
  // insert trans op for output
  auto graph_optimizer = std::make_shared<opt::GraphOptimizer>();
  auto pass_manager = std::make_shared<opt::PassManager>();
  auto insert_trans_op_pass = std::make_shared<opt::InsertTransOp>();
  insert_trans_op_pass->kernel_select_ = std::make_shared<MockRemoveInternalOutputTransOpKernelSelect>();
  pass_manager->AddPass(insert_trans_op_pass);
  graph_optimizer->AddPassManager(pass_manager);
  auto new_g = graph_optimizer->Optimize(kg);
  FuncGraphPtr g_after =
    getPyFun_.CallAndParseRet("test_remove_internal_output_trans_op_for_multiple_output", "after_insert_trans_op");
  EXPECT_TRUE(CheckEqualGraph(g_after, new_g));

  auto output_make_tuple = GetMakeTuple(kg);
  auto make_tuple = output_make_tuple->cast<CNodePtr>()->input(1);
  auto tuple_getitem = make_tuple->cast<CNodePtr>()->input(1);
  auto make_tuple1 = tuple_getitem->cast<CNodePtr>()->input(1);
  auto trans_data1 = make_tuple1->cast<CNodePtr>()->input(1);
  auto trans_data2 = make_tuple1->cast<CNodePtr>()->input(2);
  EXPECT_TRUE(kg->IsInternalOutput(trans_data1, 0));
  EXPECT_TRUE(kg->IsInternalOutput(trans_data2, 0));

  // remove trans op for internal output
  auto graph_optimizer1 = std::make_shared<opt::GraphOptimizer>();
  auto pass_manager1 = std::make_shared<opt::PassManager>();
  auto remove_internal_output_trans_op_pass = std::make_shared<opt::RemoveInternalOutputTransOp>();
  pass_manager1->AddPass(remove_internal_output_trans_op_pass);
  graph_optimizer1->AddPassManager(pass_manager1);
  auto new_g1 = graph_optimizer1->Optimize(new_g);
  FuncGraphPtr g_after1 = getPyFun_.CallAndParseRet("test_remove_internal_output_trans_op_for_multiple_output",
                                                    "after_remove_internal_output_trans_op");
  EXPECT_TRUE(CheckEqualGraph(g_after1, new_g1));
}
}  // namespace opt
}  // namespace mindspore
