/**
 * 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 "graph/passes/base_pass.h"

#include <queue>
#include <unordered_set>

#include "common/debug/log.h"
#include "graph/utils/graph_utils.h"

namespace ge {
namespace {
constexpr int kMaxRePassTimes = 10000;
constexpr size_t kMaxOneInNodes = 1000;
// Each iteration, we take about 0.3k memory on the stack, we should change the recursion to loop later
constexpr int kMaxRecursiveDepth = 20;

void GetAllNodesNoInputEdge(const ComputeGraphPtr &graph,
                            GEPass::GraphLevelState &g_state) {
  for (auto &node : graph->GetDirectNode()) {
    if (node == nullptr) {
      continue;
    }
    size_t in_nums = node->GetInNodes().size();
    if (in_nums == 0) {
      g_state.AddNodeToQueueIfNotSeen(node);
    } else if (in_nums > kMaxOneInNodes) {
      g_state.nodes_last.insert(node);
    }
  }
}

bool AllNodesIn(const Node::Vistor<NodePtr> &nodes, const std::unordered_set<NodePtr> &nodes_set) {
  return !std::any_of(nodes.begin(), nodes.end(), [&](const NodePtr &n) {
    return nodes_set.count(n) == 0;
  });
}

bool AnyNodesIn(const Node::Vistor<NodePtr> &nodes, const std::unordered_set<NodePtr> &nodes_set) {
  return std::any_of(nodes.begin(), nodes.end(), [&](const NodePtr &n) {
    return nodes_set.count(n) > 0;
  });
}

bool IsNodeReadyToQueue(const NodePtr &node, GEPass::GraphLevelState &g_state) {
  if (node == nullptr) {
    GELOGW("node is null");
    return false;
  }

  if (g_state.nodes_deleted.count(node) > 0) {
    GELOGD("The node %s was deleted before, skip it.", node->GetName().c_str());
    return false;
  }

  if (g_state.nodes_last.count(node) != 0) {
    return false;
  }

  if (!node->IsAllInNodesSeen(g_state.nodes_seen)) {
    return false;
  }

  // 因为在PassNode之前，会首先将node的输出节点添加queue，因此若在pass node时，suspend了node的输出节点，后续逻辑与上面相同
  // TODO 需要注意的是，这里的保证是一次”尽力而为“，若pass node时，将node之前的节点`A`添加到了suspend，
  //   那么`A`节点的后继和间接后继节点的pass不会受到suspend的影响
  //   理论上来说，如果在pass node之前，首先收集node的输出节点，在pass后，将输出节点做suspend、delete的去除，然后加queue，
  //   这样处理就不需要在这里做额外的确认了
  if (g_state.nodes_suspend.count(node) > 0) {
    GELOGD("The node %s has been added to suspend-iteration nodes list, the iteration of it will be suspend.",
           node->GetName().c_str());
    return false;
  }
  if (AnyNodesIn(node->GetInAllNodes(), g_state.nodes_suspend)) {
    GELOGD("The node %s has been added to suspend-iteration nodes list, the iteration of it will be suspend.",
           node->GetName().c_str());
    return false;
  }
  return true;
}

void CollectOutNodesBeforePass(const NodePtr &node, std::unordered_set<NodePtr> &out_nodes_before_pass) {
  for (const auto &out_node : node->GetOutNodes()) {
    out_nodes_before_pass.insert(out_node);
  }
}

void AddNextIterNodes(const NodePtr &cur_node, std::unordered_set<NodePtr> &out_nodes_before_pass,
                      GEPass::GraphLevelState &g_state) {
  for (auto &node : cur_node->GetOutNodes()) {
    if (node == nullptr) {
      continue;
    }
    if (out_nodes_before_pass.erase(node) == 0) {
      // after pass node , new output node come up
      GELOGD("New output nodes %s come up after pass %s.", node->GetName().c_str(), cur_node->GetName().c_str());
    }

    if (IsNodeReadyToQueue(node, g_state)) {
      g_state.AddNodeToQueueIfNotSeen(node);
    }
  }
  // A-->B-->C
  //      \
  //   D--->E
  // If B has been delete after pass, two case need to consider
  // 1. A & C & E has been repass by B. good choice
  // 2. A & C & E not added to repass, C will not pass because no one trigger it.
  //    while E will pass because D will trigger it.
  // So here we need add node which has no input_node to queue.
  for (const auto &node : out_nodes_before_pass) {
    if (!node->GetInAllNodes().empty()) {
      GELOGD("Node %s used to be output of node %s, but after pass it doesnt. "
          "It may triggered by other node, so no need add to queue now.");
      continue;
    }
    if (IsNodeReadyToQueue(node, g_state)) {
      // unlink edge may happen, add these node to queue otherwise they can not pass
      GELOGI("Node %s may lost from cur node, add to queue if not seen.",
             node->GetName().c_str(), cur_node->GetName().c_str());
      g_state.AddNodeToQueueIfNotSeen(node);
    }
  }
}

void AddImmediateRepassNodesToQueue(NodePtr &cur_node,
                                    const std::unordered_map<NodePtr, std::string> re_pass_imm_nodes_to_pass_names,
                                    GEPass::GraphLevelState &g_state) {
  for (const auto &node_2_pass_names : re_pass_imm_nodes_to_pass_names) {
    auto repass_imm_node = node_2_pass_names.first;
    if (repass_imm_node == nullptr) {
      GELOGW("Found null immediately re-pass node when executing pass %s on node %s type %s",
             node_2_pass_names.second.c_str(),
             cur_node->GetName().c_str(), cur_node->GetType().c_str());
      continue;
    }
    if (g_state.nodes_passed.count(repass_imm_node) > 0) {
      GELOGD("The node %s specified by pass %s has been passed, it will repass immediately",
             repass_imm_node->GetName().c_str(), node_2_pass_names.second.c_str());
      g_state.AddNodeToQueueFront(repass_imm_node);
      continue;
    }
    GELOGW("The node %s specified by pass %s has un-passed in_nodes, it will not repass immediately",
           repass_imm_node->GetName().c_str(), node_2_pass_names.second.c_str());
  }
}

void AddLastNodesToQueue(GEPass::GraphLevelState &g_state) {
  for (auto &node : g_state.nodes_last) {
    // todo 为什么会在node_seen中看到node_last，blame一下看看历史合入记录
    if (node->IsAllInNodesSeen(g_state.nodes_seen)) {
      g_state.AddNodeToQueueIfNotSeen(node);
    }
  }
  g_state.nodes_last.clear();
}

void AddResumeNodesToQueue(const std::unordered_map<NodePtr, std::string> resume_nodes_to_pass_names,
                           GEPass::GraphLevelState &g_state) {
  // Currently we dont keep the order of suspend nodes and resume nodes, so its hard to know
  // which one comes first. Simple way : if a node both have suspend & resume state, we will resume it.
  // Better way: keep the order when suspend/resume a node, and in this func suspend/resume in order.
  for (const auto &node_2_pass_names : resume_nodes_to_pass_names) {
    auto node = node_2_pass_names.first;
    if (g_state.nodes_suspend.erase(node) > 0) {
      if (g_state.nodes_seen.count(node.get()) > 0 || node->IsAllInNodesSeen(g_state.nodes_seen)) {
        g_state.nodes.push_back(node);
        GELOGD("Node %s has been resumed by pass %s, add to queue.",
               node->GetName().c_str(), node_2_pass_names.second.c_str());
      }
    }
  }
}

void PushToRePassIfSeen(NodePtr &node, const std::pair<std::string, BaseNodePass *> &name_to_pass,
                        std::unordered_set<Node *> &nodes_seen, const std::unordered_set<NodePtr> &nodes_to_re_pass,
                        std::unordered_set<NodePtr> &nodes_re_pass) {
  for (const auto &node_to_re_pass : nodes_to_re_pass) {
    if (node_to_re_pass == nullptr) {
      GELOGW("Found null re-pass node when executing %s on node %s type %s", name_to_pass.first.c_str(),
             node->GetName().c_str(), node->GetType().c_str());
      continue;
    }
    if (nodes_seen.count(node_to_re_pass.get()) > 0 || node_to_re_pass->IsAllInNodesSeen(nodes_seen)) {
      GELOGD("The node %s will be re-pass.", node_to_re_pass->GetName().c_str());
      nodes_re_pass.insert(node_to_re_pass);
    } else {
      GELOGD("The node %s are not all seen, don't set repass this time", node_to_re_pass->GetName().c_str());
    }
  }
}

void SetFlagOption(NodePassOption option, NamesToPass names_to_pass) {
  for (auto &name_to_pass : names_to_pass) {
    name_to_pass.second->SetOption(option, "");
  }
}

void ClearOption(NamesToPass names_to_pass) {
  for (auto &name_to_pass : names_to_pass) {
    name_to_pass.second->ClearOptions();
  }
}
}  // namespace

Status BaseNodePass::IsolateAndDeleteNode(NodePtr &node, const std::vector<int> &io_map,
                                          bool is_repass_io_immediately) {
  if (node == nullptr) {
    REPORT_INNER_ERROR("E19999", "Param node is nullptr, check invalid.");
    GELOGE(FAILED, "[Check][Param] parameter node is nullptr.");
    return FAILED;
  }
  GELOGI("Prepare to isolate and delete node, name:%s, type:%s.", node->GetName().c_str(),
         node->GetType().c_str());
  ComputeGraphPtr graph = node->GetOwnerComputeGraph();
  if (graph == nullptr) {
    REPORT_INNER_ERROR("E19999", "The owner graph of node:%s must not be null.", node->GetName().c_str());
    GELOGE(FAILED, "[Get][OwnerComputeGraph] failed, The owner graph of node:%s must not be null.",
           node->GetName().c_str());
    return FAILED;
  }

  is_repass_io_immediately ? AddImmediateRePassNodesWithInOut(node) : AddRePassNodesWithInOut(node);

  if (GraphUtils::IsolateNode(node, io_map) != GRAPH_SUCCESS) {
    REPORT_CALL_ERROR("E19999", "Isolate Node:%s failed", node->GetName().c_str());
    GELOGE(FAILED, "[Isolate][Node] %s failed.", node->GetName().c_str());
    return FAILED;
  }

  if (GraphUtils::RemoveNodeWithoutRelink(graph, node) != SUCCESS) {
    REPORT_CALL_ERROR("E19999", "call RemoveNodeWithoutRelink for node:%s failed.", node->GetName().c_str());
    GELOGE(FAILED, "[Call][RemoveNodeWithoutRelink] for node:%s failed.", node->GetName().c_str());
    return FAILED;
  }

  AddNodeDeleted(node);
  return SUCCESS;
}

Status GEPass::Run(const NamesToPass &names_to_passes) {
  if (graph_ == nullptr) {
    REPORT_INNER_ERROR("E19999", "graph_ is nullptr, check invalid.");
    GELOGE(INTERNAL_ERROR, "[Check][Param] The graph is nullptr");
    return INTERNAL_ERROR;
  }
  if (names_to_passes.empty()) {
    GELOGW("No passes input, the GEPass will do nothing");
    return INTERNAL_ERROR;
  }
  for (const auto &name_to_pass : names_to_passes) {
    if (name_to_pass.second == nullptr) {
      GELOGE(INTERNAL_ERROR, "[Check][Param] There is null pointer in passes(%s)", name_to_pass.first.c_str());
      return INTERNAL_ERROR;
    }
  }

  if (depth_ > kMaxRecursiveDepth) {
    GELOGE(PARAM_INVALID,
        "[Check][Param] The pass for root graph %s will be terminated because too many nesting"
        " levels(%d) of subgraphs, last subgraph is %s",
        root_graph_->GetName().c_str(), depth_, graph_->GetName().c_str());
    return PARAM_INVALID;
  }

  return RunPassesOneGraph(names_to_passes);
}

void NotifyPassGraphStart(const ComputeGraphPtr &graph, const NamesToPass &names_to_pass) {
  for (auto &name_to_pass : names_to_pass) {
    name_to_pass.second->OnStartPassGraph(graph);
  }
}

Status GEPass::HandleLeakedSuspendNodes(const NamesToPass &names_to_passes, GraphLevelState &g_state) {
  std::unordered_map<NodePtr, std::string> resume_nodes_to_pass_names;
  for (auto &name_to_pass : names_to_passes) {
    name_to_pass.second->init();
    auto ret = name_to_pass.second->OnSuspendNodesLeaked();
    if (ret != SUCCESS) {
      GELOGE(ret, "Internal Error happened when pass %s handle on suspend nodes leaked.",
             name_to_pass.first.c_str());
      return ret;
    }
    for (const auto &resume_node : name_to_pass.second->GetNodesResume()){
      resume_nodes_to_pass_names[resume_node].append(name_to_pass.first + ",");
    }
  }
  AddResumeNodesToQueue(resume_nodes_to_pass_names, g_state);
  return SUCCESS;
}

Status GEPass::RunPassesOneGraph(const NamesToPass &names_to_passes) {
  GELOGD("Begin to run pass on graph, passes count %zu", names_to_passes.size());
  NotifyPassGraphStart(graph_, names_to_passes);
  GraphLevelState g_state;
  g_state.re_pass_times = 0;
  GetAllNodesNoInputEdge(graph_, g_state);
  GELOGD("Start points count %zu", g_state.nodes.size());

  do {
    if (!g_state.nodes_suspend.empty()) {
      auto ret = HandleLeakedSuspendNodes(names_to_passes, g_state);
      if (ret != SUCCESS) {
        GELOGE(ret, "Failed to handle leaked suspend nodes, break base pass.");
        return ret;
      }
      if (g_state.nodes.empty()) {
        // There are suspend nodes leaked, but no pass resume it
        GELOGE(INTERNAL_ERROR, "There are suspend nodes but no pass resume, which means"
               "some nodes in this graph never pass.");
        return INTERNAL_ERROR;
      }
    }
    auto ret = RunPassesGraphRepass(names_to_passes, g_state);
    if (ret != SUCCESS) {
      return ret;
    }
  } while (!g_state.nodes_suspend.empty());

  return SUCCESS;
}


Status GEPass::RunPassesGraphRepass(const NamesToPass &names_to_passes, GraphLevelState &g_state) {
  RepassLevelState rp_state;
  do {
    for (auto &node : rp_state.nodes_re_pass) {
      GELOGD("Add node %s to queue for re-pass.", node->GetName().c_str());
      g_state.AddNodeToQueue(node);
    }
    rp_state.nodes_re_pass.clear();

    while (!g_state.nodes.empty()) {
      auto node = g_state.PopFront();

      if (g_state.nodes_deleted.count(node) > 0) {
        GELOGD("The node %s was deleted before, skip it.", node->GetName().c_str());
      }
      (void)rp_state.nodes_re_pass.erase(node);// todo why
      g_state.nodes_seen.insert(node.get()); // todo 为什么这里seen

      std::unordered_set<NodePtr> out_nodes_before_pass;
      CollectOutNodesBeforePass(node, out_nodes_before_pass);

      auto ret = RunPassesNodeOnce(node, names_to_passes, g_state, rp_state);
      if (ret != SUCCESS) {
        GELOGE(ret, "[Process][Passes] on node %s type %s failed, error code:%u", node->GetName().c_str(),
               node->GetType().c_str(), ret);
        return ret;
      }
      AddNextIterNodes(node, out_nodes_before_pass, g_state);
    }
    AddLastNodesToQueue(g_state);
  } while ((!rp_state.nodes_re_pass.empty() || !g_state.nodes.empty()) && ++g_state.re_pass_times < kMaxRePassTimes);

  if (g_state.re_pass_times == kMaxRePassTimes) {
    GELOGW("re_pass_times should not come to %d", kMaxRePassTimes);
  }
  GELOGD("All passes runs end");
  return SUCCESS;
}

Status GEPass::RunPassesOnSubGraph(const NodePtr &node, const NamesToPass &names_to_passes, bool &has_sub_graph) {
  auto sub_graph_names = node->GetOpDesc()->GetSubgraphInstanceNames();
  has_sub_graph = false;
  for (const auto &name : sub_graph_names) {
    auto graph = root_graph_->GetSubgraph(name);
    if (graph == nullptr) {
      GELOGW("Can not find the sub graph %s from node %s, the pass-process will skip it",
          name.c_str(), node->GetName().c_str());
      continue;
    }
    has_sub_graph = true;
    GELOGI("Begin to run passes on the sub graph %s of node %s", name.c_str(), node->GetName().c_str());
    GEPass pass(graph, root_graph_, depth_ + 1);
    auto ret = pass.Run(names_to_passes);
    if (ret != SUCCESS) {
      GELOGE(ret, "[Run][Passes] for sub graph:%s from node:%s failed", name.c_str(), node->GetName().c_str());
      return ret;
    }
  }
  return SUCCESS;
}

Status GEPass::RunPassesNodeOnce(NodePtr &node, const NamesToPass &names_to_passes,
                                 GraphLevelState &g_state, RepassLevelState &rp_state) {
  auto ret = RunPassesOnNode(node, names_to_passes, g_state, rp_state);
  if (ret != SUCCESS) {
    GELOGE(ret, "[Process][Passes] on node %s type %s failed, error code:%u", node->GetName().c_str(),
           node->GetType().c_str(), ret);
    return ret;
  }

  bool has_sub_graph = false;
  ret = RunPassesOnSubGraph(node, names_to_passes, has_sub_graph);
  if (ret != SUCCESS) {
    GELOGE(ret, "[Run][Passes] on the sub graph of node %s failed", node->GetName().c_str());
    return ret;
  }

  if (has_sub_graph) {
    GELOGD("There are subgraphs on node %s, run passes for for the second time", node->GetName().c_str());
    SetFlagOption(kOptimizeAfterSubGraph, names_to_passes);
    ret = RunPassesOnNode(node, names_to_passes, g_state, rp_state);
    if (ret != SUCCESS) {
      GELOGE(ret, "[Process][Passes] on node %s type %s failed, error code: %u", node->GetName().c_str(),
             node->GetType().c_str(), ret);
      return ret;
    }

    // There is only one option scene, so set and clear options around the `RunPasses` func.
    // if there are more than one scene to set options, the `ClearOption` function
    // should be called each time at the begin of the iteration
    ClearOption(names_to_passes);
  }
  return SUCCESS;
}

Status GEPass::RunPassesOnNode(NodePtr &node, const NamesToPass &names_to_passes, GraphLevelState &g_state,
                               RepassLevelState &rp_state) {
  if (node == nullptr) {
    REPORT_INNER_ERROR("E19999", "Param node is nullptr, check invalid.");
    GELOGE(FAILED, "[Check][Param] parameter node is nullptr.");
    return FAILED;
  }
  GELOGD("Begin to run pass for node %s", node->GetName().c_str());
  for (const auto &name_to_pass : names_to_passes) {
    GELOGD("Begin to run pass %s for node %s", name_to_pass.first.c_str(), node->GetName().c_str());
    name_to_pass.second->init();
    auto result = name_to_pass.second->Run(node);
    if (result != SUCCESS) {
      REPORT_CALL_ERROR("E19999", "process pass %s on node:%s failed, ret:%u", name_to_pass.first.c_str(),
                        node->GetName().c_str(), result);
      GELOGE(INTERNAL_ERROR,
             "[Process][Pass] %s on node %s failed, result "
             "%u, the passes will be terminated immediately.",
             name_to_pass.first.c_str(), node->GetName().c_str(), result);
      return result;
    }
    if (name_to_pass.second->GetNodesDeleted().count(node) > 0) {
      GELOGD("The node %s was deleted by pass %s, stop the remain passes", node->GetName().c_str(),
             name_to_pass.first.c_str());
      break;
    }
  }

  g_state.nodes_passed.insert(node);

  std::unordered_map<NodePtr, std::string> repass_imm_nodes_to_pass_names;
  std::unordered_map<NodePtr, std::string> resume_nodes_to_pass_names;
  // if multi pass add one node to repass immediately, here need to remove duplication
  for (const auto &name_to_pass : names_to_passes) {
    PushToRePassIfSeen(node, name_to_pass, g_state.nodes_seen, name_to_pass.second->GetNodesNeedRePass(),
                       rp_state.nodes_re_pass);
    // collect imm_node && resume_node among these passes
    for (const auto &imm_node : name_to_pass.second->GetNodesNeedRePassImmediately()) {
      repass_imm_nodes_to_pass_names[imm_node].append(name_to_pass.first + ",");
    }
    for (const auto &resume_node : name_to_pass.second->GetNodesResume()) {
      resume_nodes_to_pass_names[resume_node].append(name_to_pass.first + ",");
    }

    for (const auto &suspend_node : name_to_pass.second->GetNodesSuspend()) {
      GELOGD("The iteration suspend of node %s has been set by pass %s", suspend_node->GetName().c_str(),
             name_to_pass.first.c_str());
      g_state.nodes_suspend.insert(suspend_node);
    }
    const auto &nodes_deleted_by_pass = name_to_pass.second->GetNodesDeleted();
    g_state.nodes_deleted.insert(nodes_deleted_by_pass.begin(), nodes_deleted_by_pass.end());
  }
  AddImmediateRepassNodesToQueue(node, repass_imm_nodes_to_pass_names, g_state);
  AddResumeNodesToQueue(resume_nodes_to_pass_names, g_state);
  return SUCCESS;
}
}  // namespace ge