#ifndef ITALIANO_H_
#define ITALIANO_H_

#include "algorithm/decremental_reachability.h"
#include "graph/breadth_first_tree.h"
#include <stack>

namespace algo {

template <typename T> class Italiano : public DecrementalReachability<T> {
public:
  Italiano() = default;

  explicit Italiano(graph::Digraph<T> G) { this->G = G; }

  // Initialize the decremental maintenance data structure for DAGs
  void init() override;

  // Execute reachability query q(u, v) from vertex u to vertex v
  // in O(1) using the transitive closure matrix
  bool query(const T &u, const T &v) override;

  // Delete collection of edges
  void remove(const std::vector<std::pair<T, T>> &edges) override;

  // Delete edge e(u, v) and explicitly maintain the transitive closure
  void remove(const T &u, const T &v);

private:
  // Transitive closure matrix
  std::unordered_map<T, std::unordered_map<T, bool>> TC;

  // For each vertex, store a reachability tree created as a BFS tree
  std::unordered_map<T, BreadthFirstTree<T>> RT;

  // For each vertex, store collections of its incoming and outgoing edges
  struct Edges {
    std::set<T> inc;
    std::set<T> out;
  };
  std::unordered_map<T, Edges> E;

  //
  std::unordered_map<T, std::stack<T>> H;

  // Repair reachability trees after edge deletions
  void repairTrees();
};

template <typename T> void Italiano<T>::init() {
  for (const auto &u : this->G.vertices()) {
    for (const auto &v : this->G.adjList[u]) {
      E[v].inc.insert(u);
      E[u].out.insert(v);
      TC[u][v] = false;
    }
    RT[u] = BreadthFirstTree<T>(this->G, u);
    TC[u][u] = true;
    for (const auto &v : RT[u].vertices())
      TC[u][v] = true;
  }
}

template <typename T> bool Italiano<T>::query(const T &u, const T &v) {
  return TC[u][v];
}

template <typename T>
void Italiano<T>::remove(const std::vector<std::pair<T, T>> &edges) {
  for (const auto &[u, v] : edges) {
    if (!this->G.adjList[u].contains(v))
      continue;
    remove(u, v);
  }
}

template <typename T> void Italiano<T>::remove(const T &u, const T &v) {
  this->G.remove(u, v);
  E[u].out.erase(v);
  E[v].inc.erase(u);

  for (const auto &w : this->G.vertices()) {
    if (!RT[w].adjList[u].contains(v))
      continue;

    RT[w].adjList[u].erase(v);
    if (E[v].inc.size() > 0) {
      H[w].push(v);
      continue;
    }
    TC[w][v] = false;
    for (const auto &c : E[v].out)
      H[w].push(c);
  }
  repairTrees();
}

template <typename T> void Italiano<T>::repairTrees() {
  for (const auto &w : this->G.vertices()) {
    while (H[w].size() > 0) {
      const auto &h = H[w].top();
      H[w].pop();
      bool foundHook = false;

      for (const auto &i : E[h].inc) {
        if (RT[w].adjList[w].contains(i)) {
          RT[w].adjList[i].insert(h);
          foundHook = true;
          break;
        }
      }
      if (foundHook)
        continue;

      TC[w][h] = false;
      for (const auto &o : E[h].out) {
        if (RT[w].adjList[h].contains(o)) {
          H[w].push(o);
        }
      }
    }
  }
}

} // namespace algo

#endif