/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of 
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

#include "graph/preprocess/insert_op/insert_aipp_op_util.h"

#include "base/err_msg.h"

#include <fstream>
#include <utility>
#include "common/checker.h"
#include "common/dynamic_aipp.h"
#include "formats/utils/formats_trans_utils.h"
#include "common/plugin/ge_make_unique_util.h"
#include "framework/common/op/ge_op_utils.h"
#include "framework/common/util.h"
#include "base/err_msg.h"
#include "framework/common/debug/ge_log.h"
#include "framework/common/debug/log.h"
#include "framework/common/ge_inner_error_codes.h"
#include "framework/omg/omg_inner_types.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/preprocess/insert_op/aipp_op.h"
#include "graph/utils/attr_utils.h"
#include "graph/utils/graph_utils.h"
#include "graph/utils/op_desc_utils.h"
#include "graph/utils/tensor_utils.h"
#include "graph/utils/type_utils.h"
#include "common/proto_util/proto_util.h"
#include "graph/utils/op_type_utils.h"

using domi::AippOpParams;

namespace ge {
static void ConvertShape2Nhwc(Format &format, std::vector<int64_t> &shape_vec) {
  if ((format == FORMAT_NHWC) || (shape_vec.size() != static_cast<size_t>(NORMAL_TENSOR_SIZE))) {
    return;
  }
  if (format != FORMAT_NCHW) {
    GELOGW("The format is not NCHW, current format is %s.", TypeUtils::FormatToSerialString(format).c_str());
    return;
  }
  std::vector<int64_t> shape_vec_tmp;
  shape_vec.swap(shape_vec_tmp);
  shape_vec.push_back(shape_vec_tmp[NCHW_DIM_N]);
  shape_vec.push_back(shape_vec_tmp[NCHW_DIM_H]);
  shape_vec.push_back(shape_vec_tmp[NCHW_DIM_W]);
  shape_vec.push_back(shape_vec_tmp[NCHW_DIM_C]);
  return;
}

Status InsertAippOpUtil::Init() {
  insert_op_conf_.reset((new (std::nothrow) domi::InsertNewOps()));
  GE_CHECK_NOTNULL(insert_op_conf_);
  return SUCCESS;
}

Status InsertAippOpUtil::CheckAndCopyAippOpParams(const GraphManagerOptions& options) {
  const bool dynamic_image_size_is_set = !options.dynamic_image_size.empty();
  for (int32_t i = 0; i < insert_op_conf_->aipp_op_size(); i++) {
    domi::AippOpParams *aipp_op_params = insert_op_conf_->mutable_aipp_op(i);
    if (dynamic_image_size_is_set) {
      const bool crop_padding = aipp_op_params->crop() || aipp_op_params->padding();
      const bool size_w_h = (aipp_op_params->aipp_mode() == domi::AippOpParams::static_) &&
                            ((aipp_op_params->src_image_size_w() != 0) || (aipp_op_params->src_image_size_h() != 0));
      if (crop_padding || size_w_h) {
        std::string reason("When --dynamic_image_size is set, ");
        reason += (crop_padding ? "crop and padding cannot be set to 'true'"
                                : "src_image_size_w and src_image_size_h must be set to '0'");
        REPORT_PREDEFINED_ERR_MSG("E10052", std::vector<const char *>({"reason"}), std::vector<const char *>({reason.c_str()}));
        GELOGE(PARAM_INVALID, "%s", reason.c_str());
        return PARAM_INVALID;
      }
    }
    auto aipp_op = MakeUnique<AippOp>();
    GE_CHECK_NOTNULL(aipp_op);
    GE_CHK_STATUS_RET(aipp_op->Init(aipp_op_params), "[Call][Init] Aipp op init failed.");
    insert_ops_.push_back(std::move(aipp_op));
  }
  return SUCCESS;
}

Status InsertAippOpUtil::Parse(const GraphManagerOptions& options) {
  if (options.insert_op_file.empty()) {
    return SUCCESS;
  }
  GE_CHK_BOOL_RET_STATUS(ReadProtoFromText(options.insert_op_file.c_str(), insert_op_conf_.get()), FAILED,
                         "[Read][Proto] from file:%s failed", options.insert_op_file.c_str());

  GE_CHK_STATUS_RET(CheckPositionNotRepeat(), "[Check][InsertPosition] of op failed");
  Status ret = CheckAndCopyAippOpParams(options);
  if (ret != SUCCESS) {
    GELOGE(PARAM_INVALID, "Check file info of [%s] failed!", options.insert_op_file.c_str());
    return ret;
  }
  for (auto &dynamic_op : insert_ops_) {
    GE_CHECK_NOTNULL(dynamic_op);
    
    GE_CHK_STATUS_RET(dynamic_op->ValidateParams(), "[Call][ValidateParams] Validate insert_op config file failed");
    GE_CHK_STATUS_RET(dynamic_op->SetDefaultParams(), "[Call][SetDefaultParams] Set default value of insert_op failed");
  }

  return SUCCESS;
}

Status InsertAippOpUtil::InsertAippOps(ComputeGraphPtr &graph, std::string &aippConfigPath) {
  GE_CHECK_NOTNULL(graph);
  for (uint32_t index = 0; index < insert_ops_.size(); ++index) {
    GE_CHK_STATUS_RET(insert_ops_[index]->InsertAippToGraph(graph, aippConfigPath, index),
                      "[Insert][Op] to graph failed");
  }

  GE_CHK_STATUS_RET(CheckGraph(graph), "[Check][Graph] failed after inserting all ops");

  GE_CHK_GRAPH_STATUS_RET(graph->TopologicalSorting(), "[Sort][Graph] failed after insert dynamic op");

  ClearNewOps();

  return SUCCESS;
}

void InsertAippOpUtil::ClearNewOps() {
  if (insert_op_conf_ != nullptr) {
    insert_op_conf_->Clear();
    insert_ops_.clear();
  }
}

Status InsertAippOpUtil::CheckInputNamePositionNotRepeat() const {
  for (int32_t i = 0; i < insert_op_conf_->aipp_op_size(); i++) {
    const domi::AippOpParams *item = insert_op_conf_->mutable_aipp_op(i);
    GE_CHECK_NOTNULL(item);

    for (int32_t j = i + 1; j < insert_op_conf_->aipp_op_size(); j++) {
      const domi::AippOpParams *another_item = insert_op_conf_->mutable_aipp_op(j);
      GE_CHECK_NOTNULL(another_item);
      if (another_item->related_input_name().empty()) {
        std::string error_msg =
            "Cannot both set related_input_name and related_input_rank. Please ensure param is the same as the first "
            "aipp config(related_input_name)";
        GELOGE(PARAM_INVALID, "[Check][InputParam]%s", error_msg.c_str());
        REPORT_PREDEFINED_ERR_MSG("E10052", std::vector<const char_t *>({"reason"}),
                                  std::vector<const char_t *>({error_msg.c_str()}));
        return PARAM_INVALID;
      }
      if (item->related_input_name() == another_item->related_input_name()) {
        std::string error_msg =
            "Cannot insert aipp to the same position! Please ensure related_input_name"
            " param is different in different aipp config";
        GELOGE(PARAM_INVALID, "[Check][InputParam]%s", error_msg.c_str());
        REPORT_PREDEFINED_ERR_MSG("E10052", std::vector<const char_t *>({"reason"}),
                                  std::vector<const char_t *>({error_msg.c_str()}));
        return PARAM_INVALID;
      }
    }
  }

  return SUCCESS;
}

Status InsertAippOpUtil::CheckInputRankPositionNoRepeat() const {
  for (int32_t i = 0; i < insert_op_conf_->aipp_op_size(); i++) {
    const domi::AippOpParams *item = insert_op_conf_->mutable_aipp_op(i);
    GE_CHECK_NOTNULL(item);

    for (int32_t j = i + 1; j < insert_op_conf_->aipp_op_size(); j++) {
      const domi::AippOpParams *another_item = insert_op_conf_->mutable_aipp_op(j);
      GE_CHECK_NOTNULL(another_item);
      if (!another_item->related_input_name().empty()) {
        std::string error_msg =
            "Cannot both set related_input_rank and related_input_name!"
            " Please ensure param is the same with the first aipp config(related_input_rank)";
        GELOGE(PARAM_INVALID, "[Check][InputParam]%s", error_msg.c_str());
        REPORT_PREDEFINED_ERR_MSG("E10052", std::vector<const char_t *>({"reason"}),
                                  std::vector<const char_t *>({error_msg.c_str()}));
        return PARAM_INVALID;
      }
      if (item->related_input_rank() == another_item->related_input_rank()) {
        std::string error_msg =
            "Cannot insert aipp to the same position! Please ensure related_input_rank"
            " param is different in different aipp config";
        GELOGE(PARAM_INVALID, "[Check][InputParam]%s", error_msg.c_str());
        REPORT_PREDEFINED_ERR_MSG("E10052", std::vector<const char_t *>({"reason"}),
                                  std::vector<const char_t *>({error_msg.c_str()}));
        return PARAM_INVALID;
      }
    }
  }
  return SUCCESS;
}

Status InsertAippOpUtil::CheckPositionNotRepeat() {
  GE_CHECK_NOTNULL(insert_op_conf_);

  if (insert_op_conf_->aipp_op_size() <= 1) {
    GELOGI("Aipp op size[%d] less than 2, no need to check position repeat.", insert_op_conf_->aipp_op_size());
    return SUCCESS;
  }

  const domi::AippOpParams *item = insert_op_conf_->mutable_aipp_op(0);
  GE_CHECK_NOTNULL(item);

  std::string related_input_name = item->related_input_name();
  Status ret = FAILED;
  if (related_input_name.empty()) {
    ret = CheckInputRankPositionNoRepeat();
  } else {
    ret = CheckInputNamePositionNotRepeat();
  }
  if (ret != SUCCESS) {
    GELOGE(FAILED, "[Check][Position] not repeat failed.");
    return FAILED;
  }

  return SUCCESS;
}

Status InsertAippOpUtil::CheckGraph(const ComputeGraphPtr &graph) const {
  GE_CHECK_NOTNULL(graph);
  domi::AippOpParams::AippMode aippMode = domi::AippOpParams::undefined;

  for (const auto &node : graph->GetDirectNode()) {
    if (node->GetType() != DATA) {
      continue;
    }
    size_t next_nodes_cnt = 0;
    std::vector<NodePtr> aippNodes;
    for (const auto &anchor : node->GetAllOutDataAnchors()) {
      for (const auto &inAnchor : anchor->GetPeerInDataAnchors()) {
        const std::string &nodeType = inAnchor->GetOwnerNode()->GetType();
        next_nodes_cnt++;
        if (nodeType == AIPP) {
          aippNodes.push_back(inAnchor->GetOwnerNode());
          continue;
        }
      }
    }
    GE_CHK_LOG_AND_ERRORMSG((aippNodes.size() == 0) || (aippNodes.size() == next_nodes_cnt),
        PARAM_INVALID,
        "Cannot config part of outputs of Data node to support AIPP, config all "
        "of the outputs of Data to support AIPP, or config none of them");

    auto aippParams = MakeUnique<domi::AippOpParams>();
    GE_CHECK_NOTNULL(aippParams);
    GE_IF_BOOL_EXEC(
        aippNodes.size() > 1, for (decltype(aippNodes)::size_type i = 1; i < aippNodes.size(); i++) {
          auto currAippParam = MakeUnique<domi::AippOpParams>();
          GE_CHECK_NOTNULL(currAippParam);
          GE_CHK_STATUS(GetAippParams(currAippParam, aippNodes[i]));

          if (aippMode == domi::AippOpParams::static_) {
            GE_CHK_LOG_AND_ERRORMSG(
                aippParams->input_format() == currAippParam->input_format(),
                PARAM_INVALID, "The input_format of all aipp_ops after one Data should be the same");
            GE_CHK_LOG_AND_ERRORMSG(
                aippParams->src_image_size_w() == currAippParam->src_image_size_w(),
                PARAM_INVALID, "The src_image_size_w of all aipp_ops after one Data should be the same");
            GE_CHK_LOG_AND_ERRORMSG(
                    aippParams->src_image_size_h() == currAippParam->src_image_size_h(),
                PARAM_INVALID, "The src_image_size_h of all aipp_ops after one Data should be the same");
          } else {
            GE_CHK_LOG_AND_ERRORMSG(
                aippParams->max_src_image_size() == currAippParam->max_src_image_size(),
                PARAM_INVALID, "The max_src_image_size of all aipp_ops after one Data should be the same");
          }
        });
  }

  return SUCCESS;
}

Status InsertAippOpUtil::GetAippParams(const std::unique_ptr<domi::AippOpParams> &aippParams,
    const NodePtr &aipp_node) const {
  GE_CHECK_NOTNULL(aipp_node);
  ge::NamedAttrs aipp_attr;
  const OpDescPtr tmpOpPtr = aipp_node->GetOpDesc();
  GE_CHECK_NOTNULL(tmpOpPtr);
  GE_CHK_BOOL_RET_STATUS(AttrUtils::GetNamedAttrs(tmpOpPtr, ATTR_NAME_AIPP, aipp_attr), FAILED,
                         "[Get][Attr] %s from Aipp node:%s failed",
                         ATTR_NAME_AIPP.c_str(), tmpOpPtr->GetName().c_str());
  GE_CHK_STATUS_RET(OpUtils::ConvertAippParams(aipp_attr, *aippParams),
                    "[Convert][AippParams] get aipp params failed");

  return SUCCESS;
}

Status InsertAippOpUtil::UpdateDataNodeByAipp(const ComputeGraphPtr &graph) const {
  NodePtr multbatch_case;
  for (auto &node : graph->GetDirectNode()) {
    if (node->GetType() == AIPP) {
      GE_RETURN_IF_ERROR(UpdatePrevNodeByAipp(node));
    }
    if (node->GetType() == CASE && node->GetOpDesc()->HasAttr(ATTR_NAME_BATCH_NUM)) {
      multbatch_case = node;
    }
  }

  if (multbatch_case != nullptr) {
    GE_RETURN_IF_ERROR(UpdateCaseNode(graph, multbatch_case));
  }
  return SUCCESS;
}

Status InsertAippOpUtil::FindMaxSizeNode(const ComputeGraphPtr &graph, const NodePtr &case_node,
                                         std::map<uint32_t, int64_t> &max_sizes,
                                         std::map<uint32_t, GeTensorDescPtr> &aipp_inputs) const {
  const auto &func_desc = case_node->GetOpDesc();
  for (const auto &name : func_desc->GetSubgraphInstanceNames()) {
    const auto &subgraph = graph->GetSubgraph(name);
    if (subgraph == nullptr) {
      REPORT_INNER_ERR_MSG("E19999", "Subgraph:%s of op:%s(%s) not find in graph:%s, check invalid",
                         name.c_str(), func_desc->GetName().c_str(),
                         func_desc->GetType().c_str(), graph->GetName().c_str());
      GELOGE(GE_GRAPH_EMPTY_SUBGRAPH, "[Get][SubGraph] failed, Subgraph:%s of op:%s(%s) not find in graph:%s",
             name.c_str(), func_desc->GetName().c_str(), func_desc->GetType().c_str(), graph->GetName().c_str());
      return GE_GRAPH_EMPTY_SUBGRAPH;
    }

    for (auto &node : subgraph->GetDirectNode()) {
      if (node->GetType() == AIPP) {
        GE_RETURN_IF_ERROR(UpdatePrevNodeByAipp(node));
        int64_t size = 0;
        auto in_data_anchor = node->GetInDataAnchor(0);
        GE_CHECK_NOTNULL(in_data_anchor);
        auto peer_out_anchor = in_data_anchor->GetPeerOutAnchor();
        GE_CHECK_NOTNULL(peer_out_anchor);
        const auto &src_node = peer_out_anchor->GetOwnerNode();
        const auto &src_op = src_node->GetOpDesc();
        GE_CHECK_NOTNULL(src_op);

        uint32_t parent_index = 0;
        if (!AttrUtils::GetInt(src_op, ATTR_NAME_PARENT_NODE_INDEX, parent_index)) {
          REPORT_INNER_ERR_MSG("E19999", "Get Attr:%s of op:%s(%s) failed",
                             ATTR_NAME_PARENT_NODE_INDEX.c_str(),
                             src_op->GetName().c_str(), src_op->GetType().c_str());
          GELOGE(FAILED, "[Get][Attr] %s of op:%s(%s) failed", ATTR_NAME_PARENT_NODE_INDEX.c_str(),
                 src_op->GetName().c_str(), src_op->GetType().c_str());
          return FAILED;
        }

        auto aipp_op_desc = node->GetOpDesc();
        GE_CHECK_NOTNULL(aipp_op_desc);
        auto input = aipp_op_desc->MutableInputDesc(0);
        GE_CHECK_NOTNULL(input);
        if (TensorUtils::GetSize(*input, size) == GRAPH_SUCCESS) {
          if (max_sizes[parent_index] < size) {
            max_sizes[parent_index] = size;
            aipp_inputs[parent_index] = input;
          }
        }
      }
    }
  }

  return SUCCESS;
}

Status InsertAippOpUtil::UpdateCaseNode(const ComputeGraphPtr &graph, const NodePtr &case_node) const {
  const auto &func_desc = case_node->GetOpDesc();
  std::map<uint32_t, int64_t> max_sizes;
  std::map<uint32_t, GeTensorDescPtr> aipp_inputs;

  GE_RETURN_IF_ERROR(FindMaxSizeNode(graph, case_node, max_sizes, aipp_inputs));
  for (const auto &item : aipp_inputs) {
    uint32_t parent_index = item.first;
    const GeTensorDescPtr &aipp_input = item.second;
    GE_CHECK_NOTNULL(aipp_input);

    const GeTensorDescPtr &input_desc = func_desc->MutableInputDesc(parent_index);
    GE_CHECK_NOTNULL(input_desc);
    input_desc->SetDataType(aipp_input->GetDataType());
    input_desc->SetOriginDataType(aipp_input->GetOriginDataType());
    input_desc->SetShape(aipp_input->GetShape());
    input_desc->SetOriginShape(aipp_input->GetShape());
    input_desc->SetFormat(aipp_input->GetFormat());
    input_desc->SetOriginFormat(aipp_input->GetFormat());
    ge::TensorUtils::SetSize(*input_desc, max_sizes[item.first]);

    const auto &in_anchor = case_node->GetInDataAnchor(parent_index);
    GE_ASSERT_NOTNULL(in_anchor);
    const auto &out_anchor = in_anchor->GetPeerOutAnchor();
    GE_ASSERT_NOTNULL(out_anchor);
    const auto &data = out_anchor->GetOwnerNode();
    auto data_opdesc = data->GetOpDesc();
    GE_CHECK_NOTNULL(data_opdesc);
    GE_CHECK_NOTNULL(data_opdesc->MutableOutputDesc(0));
    Format old_format = data_opdesc->MutableOutputDesc(0)->GetFormat();

    auto ret = data_opdesc->UpdateOutputDesc(0, *input_desc);
    if (ret != GRAPH_SUCCESS) {
      REPORT_INNER_ERR_MSG("E19999", "Update OutputDesc to op:%s(%s) failed, index:0",
                        data_opdesc->GetName().c_str(), data_opdesc->GetType().c_str());
      GELOGE(INTERNAL_ERROR, "[Update][OutputDesc] to op:%s(%s) failed, index:0",
             data_opdesc->GetName().c_str(), data_opdesc->GetType().c_str());
      return INTERNAL_ERROR;
    }
    ret = data_opdesc->UpdateInputDesc(0, *input_desc);
    if (ret != GRAPH_SUCCESS) {
      REPORT_INNER_ERR_MSG("E19999", "Update InputDesc to op:%s(%s) failed, index:0",
                        data_opdesc->GetName().c_str(), data_opdesc->GetType().c_str());
      GELOGE(INTERNAL_ERROR, "[Update][InputDesc] to op:%s(%s) failed, index:0",
             data_opdesc->GetName().c_str(), data_opdesc->GetType().c_str());
      return INTERNAL_ERROR;
    }

    // Update attr _mbatch_origin_input_dims for data when it is linked to aipp
    UpdateMultiBatchInputDims(data_opdesc, old_format);
  }

  return SUCCESS;
}

Status InsertAippOpUtil::UpdatePrevNodeByAipp(const NodePtr &node) const {
  GELOGI("Start to update prev node size by aipp %s.", node->GetName().c_str());
  auto aipp_op_desc = node->GetOpDesc();
  GE_CHECK_NOTNULL(aipp_op_desc);
  auto aipp_input = aipp_op_desc->MutableInputDesc(0);
  GE_CHECK_NOTNULL(aipp_input);

  int64_t size = 0;
  graphStatus graph_ret = ge::TensorUtils::GetSize(*aipp_input, size);
  if (graph_ret != GRAPH_SUCCESS) {
    REPORT_INNER_ERR_MSG("E19999", "Get input size of op:%s(%s), index:0, failed",
                      aipp_op_desc->GetName().c_str(), aipp_op_desc->GetType().c_str());
    GELOGE(FAILED, "[Get][InputSize] of op:%s(%s), index:0, failed",
           aipp_op_desc->GetName().c_str(), aipp_op_desc->GetType().c_str());
    return FAILED;
  }
  GELOGI("Get input size [%ld] from aipp [%s].", size, aipp_op_desc->GetName().c_str());
  if (size == 0) {
    REPORT_INNER_ERR_MSG("E19999", "Tensor size of op:%s(%s) is 0, input_index:0, check invalid",
                      aipp_op_desc->GetName().c_str(), aipp_op_desc->GetType().c_str());
    GELOGE(FAILED, "[Check][Param] Tensor size of op:%s(%s) is 0, input_index:0",
           aipp_op_desc->GetName().c_str(), aipp_op_desc->GetType().c_str());
    return FAILED;
  }
  (void)AttrUtils::SetInt(aipp_input, ATTR_NAME_INPUT_ORIGIN_SIZE, size);

  auto in_data_anchor = node->GetInDataAnchor(0);
  GE_CHECK_NOTNULL(in_data_anchor);
  auto peer_out_anchor = in_data_anchor->GetPeerOutAnchor();
  GE_CHECK_NOTNULL(peer_out_anchor);
  const auto &src_node = peer_out_anchor->GetOwnerNode();
  const auto &src_op = src_node->GetOpDesc();
  GE_CHECK_NOTNULL(src_op);

  DataType aipp_dt = aipp_input->GetDataType();
  aipp_input->SetOriginDataType(aipp_dt);
  DataType aipp_origni_dt = aipp_input->GetOriginDataType();
  GeShape aipp_shape = aipp_input->GetShape();
  Format aipp_format = aipp_input->GetFormat();
  std::vector<std::pair<int64_t, int64_t>> shape_ranges;
  (void)aipp_input->GetShapeRange(shape_ranges);
  GELOGI("Aipp [%s] input datatype is %s, origin datatype is %s, input shape is %s", aipp_op_desc->GetName().c_str(),
         TypeUtils::DataTypeToSerialString(aipp_dt).c_str(), TypeUtils::DataTypeToSerialString(aipp_origni_dt).c_str(),
         ge::formats::ShapeToString(aipp_shape.GetDims()).c_str());

  const GeTensorDescPtr &input = src_op->MutableInputDesc(0);
  GE_CHECK_NOTNULL(input);
  input->SetDataType(aipp_dt);
  input->SetOriginDataType(aipp_origni_dt);
  input->SetShape(aipp_shape);
  input->SetOriginShape(aipp_shape);
  input->SetFormat(aipp_format);
  input->SetOriginFormat(aipp_format);
  input->SetShapeRange(shape_ranges);
  input->SetOriginShapeRange(shape_ranges);
  ge::TensorUtils::SetSize(*input, size);

  const GeTensorDescPtr &output = src_op->MutableOutputDesc(peer_out_anchor->GetIdx());
  GE_CHECK_NOTNULL(output);
  output->SetDataType(aipp_dt);
  output->SetOriginDataType(aipp_origni_dt);
  output->SetShape(aipp_shape);
  output->SetOriginShape(aipp_shape);
  output->SetFormat(aipp_format);
  output->SetOriginFormat(aipp_format);
  output->SetShapeRange(shape_ranges);
  output->SetOriginShapeRange(shape_ranges);
  ge::TensorUtils::SetSize(*output, size);
  GELOGI("Set node %s output %d size %ld by aipp.", src_node->GetName().c_str(), peer_out_anchor->GetIdx(), size);

  return SUCCESS;
}

void InsertAippOpUtil::UpdateMultiBatchInputDims(const OpDescPtr &data_opdesc, Format &old_format) const {
  if (!data_opdesc->HasAttr(ATTR_MBATCH_ORIGIN_INPUT_DIMS)) {
    GELOGW("Failed to acquire _mbatch_origin_input_dims attr from node [%s]", data_opdesc->GetName().c_str());
    return;
  }
  auto new_data_dims = data_opdesc->GetOutputDesc(0).GetShape().GetDims();
  std::vector<int64_t> origin_input_dims;
  (void)AttrUtils::GetListInt(data_opdesc, ATTR_MBATCH_ORIGIN_INPUT_DIMS, origin_input_dims);
  // Convert origin_input_dims to NHWC because data format is set to NHWC when it is linked to aipp.
  ConvertShape2Nhwc(old_format, origin_input_dims);
  if (new_data_dims.size() != origin_input_dims.size()) {
    return;
  }
  for (size_t i = 0; i < origin_input_dims.size(); ++i) {
    // Need to update shape when aipp has crop function because H,W is different, ignore -1.
    if (origin_input_dims[i] > 0) {
      origin_input_dims[i] = new_data_dims[i];
    }
  }
  (void)AttrUtils::SetListInt(data_opdesc, ATTR_MBATCH_ORIGIN_INPUT_DIMS, origin_input_dims);
  return;
}

Status InsertAippOpUtil::GetDataRelatedNode(const NodePtr &node,
                                            std::map<NodePtr, std::set<NodePtr>> &data_next_node_map) const {
  GELOGI("Start to get data and next node %s.", node->GetName().c_str());
  const OpDescPtr data_op = node->GetOpDesc();
  GE_CHECK_NOTNULL(data_op);
  if (!data_op->HasAttr(ATTR_NAME_AIPP)) {
    GELOGI("there is not AIPP info for Data: %s.", data_op->GetName().c_str());
    return SUCCESS;
  }

  auto aipp_params = MakeUnique<domi::AippOpParams>();
  GE_CHECK_NOTNULL(aipp_params);
  ge::NamedAttrs aipp_attr;
  GE_CHK_BOOL_RET_STATUS(AttrUtils::GetNamedAttrs(data_op, ATTR_NAME_AIPP, aipp_attr), ACL_ERROR_GE_AIPP_NOT_EXIST,
                         "[Get][Attr] %s from op:%s failed", ATTR_NAME_AIPP.c_str(), data_op->GetName().c_str());
  GE_CHK_STATUS_RET(OpUtils::ConvertAippParams(aipp_attr, *aipp_params), "[Get][AippParams] failed");

  for (const auto &out_data_anchor : node->GetAllOutDataAnchors()) {
    GE_CHECK_NOTNULL(out_data_anchor);
    auto peer_in_anchors = out_data_anchor->GetPeerInDataAnchors();
    for (const auto &peer_in_data_anchor : peer_in_anchors) {
      GE_CHECK_NOTNULL(peer_in_data_anchor);
      const auto &dst_node = peer_in_data_anchor->GetOwnerNode();
      const auto &dst_op = dst_node->GetOpDesc();
      GE_CHECK_NOTNULL(dst_op);

      if ((dst_op->GetType() == AIPP) || (dst_op->GetType() == CASE)) {
        std::map<NodePtr, std::set<NodePtr>>::const_iterator data_iter = data_next_node_map.find(node);
        if (data_iter == data_next_node_map.end()) {
          std::set<NodePtr> next_node_set;
          next_node_set.insert(dst_node);
          data_next_node_map[node] = next_node_set;
        } else {
          if (data_next_node_map[node].find(dst_node) == data_next_node_map[node].end()) {
            data_next_node_map[node].insert(dst_node);
          }
        }
      }
    }
  }

  return SUCCESS;
}

Status InsertAippOpUtil::GetAllAipps(const NodePtr &data_node, const NodePtr &node, std::vector<NodePtr> &aipps) const {
  GE_CHECK_NOTNULL(node);
  const OpDescPtr op = node->GetOpDesc();
  GE_CHECK_NOTNULL(op);
  GELOGI("Get all aipp node from this node %s.", op->GetName().c_str());
  if (op->GetType() == AIPP) {
    aipps.emplace_back(node);
  } else if (op->GetType() == CASE) {
    const ComputeGraphPtr &graph = node->GetOwnerComputeGraph();
    for (const auto &name : op->GetSubgraphInstanceNames()) {
      const auto &subgraph = graph->GetSubgraph(name);
      if (subgraph == nullptr) {
        REPORT_INNER_ERR_MSG("E19999", "Subgraph:%s of op:%s(%s) not find in graph:%s, check invalid",
                           name.c_str(), op->GetName().c_str(),
                           op->GetType().c_str(), graph->GetName().c_str());
        GELOGE(GE_GRAPH_EMPTY_SUBGRAPH, "[Get][SubGraph] Subgraph:%s of op:%s(%s) not find in graph:%s",
               name.c_str(), op->GetName().c_str(), op->GetType().c_str(), graph->GetName().c_str());
        return GE_GRAPH_EMPTY_SUBGRAPH;
      }

      for (auto &subgraph_node : subgraph->GetDirectNode()) {
        if (subgraph_node->GetType() == AIPP) {
          auto src_node = subgraph_node->GetInDataNodes().at(0);
          const auto &src_op = src_node->GetOpDesc();
          GE_CHECK_NOTNULL(src_op);
          uint32_t parent_index = 0;
          if (!AttrUtils::GetInt(src_op, ATTR_NAME_PARENT_NODE_INDEX, parent_index)) {
            REPORT_INNER_ERR_MSG("E19999", "Get Attr:%s of op:%s(%s) failed",
                               ATTR_NAME_PARENT_NODE_INDEX.c_str(),
                               src_op->GetName().c_str(), src_op->GetType().c_str());
            GELOGE(FAILED, "[Get][Attr] %s of op:%s(%s) failed",
                   ATTR_NAME_PARENT_NODE_INDEX.c_str(), src_op->GetName().c_str(), src_op->GetType().c_str());
            return FAILED;
          }
          auto data = node->GetInDataNodes().at(parent_index);
          if (data->GetName() == data_node->GetName()) {
            aipps.emplace_back(subgraph_node);
          }
        }
      }
    }
  }
  return SUCCESS;
}

Status InsertAippOpUtil::RecordAIPPInfoToData(const ComputeGraphPtr &graph) const {
  GELOGI("Start to record aipp info to Data.");
  std::map<NodePtr, std::set<NodePtr>> data_next_node_map;
  for (auto &node : graph->GetDirectNode()) {
    if (OpTypeUtils::IsDataNode(node->GetType())) {
      GE_RETURN_IF_ERROR(GetDataRelatedNode(node, data_next_node_map));
    }
  }

  for (auto it : data_next_node_map) {
    std::vector<std::string> input_dims;
    std::vector<std::string> output_dims;
    auto data_node = it.first;
    auto data_op_desc = data_node->GetOpDesc();
    GE_CHECK_NOTNULL(data_op_desc);
    std::set<NodePtr> aipps_or_switchs_or_case = it.second;
    if (aipps_or_switchs_or_case.size() != 1) {
      GELOGW("The number of successors swith or aipp of data is more than 1");
      continue;
    }

    std::vector<NodePtr> aipps;
    GE_RETURN_IF_ERROR(GetAllAipps(data_node, *aipps_or_switchs_or_case.begin(), aipps));
    GELOGI("RecordAIPPInfoToData: Data: name[%s], type[%s], batch size[%zu]", data_node->GetName().c_str(),
           data_node->GetType().c_str(), aipps.size());

    for (auto aipp_it : aipps) {
      std::string input;
      std::string output;
      GetInputOutputInfo(data_node, aipp_it, input, output);
      input_dims.emplace_back(input);
      output_dims.emplace_back(output);

      // When static aipp is set, need to get the model input dims which processed by aipp
      GE_RETURN_IF_ERROR(SetModelInputDims(data_node, aipp_it));
    }
    // if _all_origin_gears_inputs is set, use its value directly.
    if (data_op_desc->HasAttr("_all_origin_gears_inputs")) {
      std::vector<std::string> input_dims_str;
      (void)AttrUtils::GetListStr(data_op_desc, "_all_origin_gears_inputs", input_dims_str);
      (void)AttrUtils::SetListStr(data_op_desc, ATTR_NAME_AIPP_INPUTS, input_dims_str);
      if ((input_dims_str.size() > output_dims.size()) && !output_dims.empty()) {
        // make sure output and input counts is equal, appears in dynamic aipp and dynamic shape/batch scene.
        std::vector<std::string> output_dims_str{input_dims_str.size(), output_dims[0]};
        (void)AttrUtils::SetListStr(data_op_desc, ATTR_NAME_AIPP_OUTPUTS, output_dims_str);
      } else {
        (void)AttrUtils::SetListStr(data_op_desc, ATTR_NAME_AIPP_OUTPUTS, output_dims);
      }
    } else {
      (void)AttrUtils::SetListStr(data_op_desc, ATTR_NAME_AIPP_INPUTS, input_dims);
      (void)AttrUtils::SetListStr(data_op_desc, ATTR_NAME_AIPP_OUTPUTS, output_dims);
    }
  }

  return SUCCESS;
}

Status InsertAippOpUtil::GetInputOutputInfo(NodePtr &data_node, NodePtr &aipp_node, std::string &input,
                                            std::string &output) const {
  GE_CHECK_NOTNULL(data_node);
  GE_CHECK_NOTNULL(aipp_node);
  const OpDescPtr data_op = data_node->GetOpDesc();
  GE_CHECK_NOTNULL(data_op);
  const OpDescPtr aipp_op = aipp_node->GetOpDesc();
  GE_CHECK_NOTNULL(aipp_op);

  // aipp node's original output shape equals to original model data's shape
  const ConstGeTensorDescPtr output_desc = aipp_op->GetOutputDescPtr(0);
  GE_CHECK_NOTNULL(output_desc);
  const Format orig_format = output_desc->GetOriginFormat();
  const DataType orig_data_type = output_desc->GetOriginDataType();
  const std::string tensor_name = data_op->GetName();
  const size_t dim_num = output_desc->GetOriginShape().GetDimNum();
  int64_t tensor_size = 0;
  (void)TensorUtils::CalcTensorMemSize(output_desc->GetOriginShape(), orig_format, orig_data_type, tensor_size);
  int64_t input_size = tensor_size;
  input = TypeUtils::FormatToSerialString(orig_format) + ":" + TypeUtils::DataTypeToSerialString(orig_data_type) + ":" +
          tensor_name + ":" + std::to_string(input_size) + ":" + std::to_string(dim_num) + ":" +
          formats::JoinToString(output_desc->GetOriginShape().GetDims());

  const Format format = output_desc->GetFormat();
  const DataType data_type = output_desc->GetDataType();
  const std::string output_name = aipp_op->GetOutputNameByIndex(0);
  const size_t output_dim_num = output_desc->GetShape().GetDimNum();
  (void)TensorUtils::CalcTensorMemSize(output_desc->GetShape(), output_desc->GetFormat(), output_desc->GetDataType(),
                                       tensor_size);
  int64_t output_size = tensor_size;
  output = TypeUtils::FormatToSerialString(format) + ":" + TypeUtils::DataTypeToSerialString(data_type) + ":" +
           output_name + ":" + std::to_string(output_size) + ":" + std::to_string(output_dim_num) + ":" +
           formats::JoinToString(output_desc->GetShape().GetDims());

  GELOGI("GetInputOutputInfo: get data[%s] node related aipp[%s] node info, input[%s], output[%s].",
         data_node->GetName().c_str(), aipp_node->GetName().c_str(), input.c_str(), output.c_str());
  return SUCCESS;
}

Status InsertAippOpUtil::SetModelInputDims(NodePtr &data_node, NodePtr &aipp_node) const {
  GE_CHECK_NOTNULL(data_node);
  GE_CHECK_NOTNULL(aipp_node);
  const OpDescPtr data_opdesc = data_node->GetOpDesc();
  GE_CHECK_NOTNULL(data_opdesc);
  const OpDescPtr aipp_opdesc = aipp_node->GetOpDesc();
  GE_CHECK_NOTNULL(aipp_opdesc);

  // In dynamic bacth/hw scenario, the new model input dims only need be set once
  if (data_node->GetOpDesc()->HasAttr(ATTR_NAME_INPUT_DIMS)) {
    GELOGD("Data %s already has attribute %s", data_node->GetOpDesc()->GetName().c_str(), ATTR_NAME_INPUT_DIMS.c_str());
    return SUCCESS;
  }
  std::vector<int64_t> model_input_dims;
  std::vector<int64_t> origin_input_dims;
  if (AttrUtils::GetListInt(aipp_opdesc, ATTR_NAME_INPUT_DIMS, model_input_dims) && !model_input_dims.empty()) {
    // When dynamic bacth/hw is set, N or HW need to be set to -1
    if (AttrUtils::GetListInt(data_opdesc, ATTR_MBATCH_ORIGIN_INPUT_DIMS, origin_input_dims) &&
        !origin_input_dims.empty()) {
      GELOGI("In dynamic bacth/hw scenario, N or HW need to be set to -1. model_input_dims: %s, origin_input_dims: %s",
             ToString(model_input_dims).c_str(), ToString(origin_input_dims).c_str());
      for (size_t i = 0UL; i < origin_input_dims.size(); ++i) {
        // N or HW need to be set to -1
        if (origin_input_dims[i] < 0L) {
          model_input_dims[i] = origin_input_dims[i];
        }
      }
    }
    GELOGD("After set N or H/W to -1, the model input dims: %s.", ToString(model_input_dims).c_str());
    if (!AttrUtils::SetListInt(data_opdesc, ATTR_NAME_INPUT_DIMS, model_input_dims)) {
      REPORT_INNER_ERR_MSG("E19999", "Set Attr:%s on op:%s(%s) failed",
                         ATTR_NAME_INPUT_DIMS.c_str(),
                         data_opdesc->GetName().c_str(), data_opdesc->GetType().c_str());
      GELOGE(FAILED, "[Set][Attr] %s on op:%s(%s) failed",
             ATTR_NAME_INPUT_DIMS.c_str(), data_opdesc->GetName().c_str(), data_opdesc->GetType().c_str());
      return FAILED;
    }
  }
  return SUCCESS;
}
}  // namespace ge