#ifndef ITALIANO_H_
#define ITALIANO_H_

#include "algorithm/decremental_reachability.h"
#include "graph/breadth_first_tree.h"

#include <stack>

using namespace graph;

namespace algo {

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

  Italiano(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 edge e(u, v) and explicitely maintain the transitive closure
  void remove(const T& u, const T& v) override;
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;

  // Repair reachability trees after edge deletions
  void repairTrees(std::unordered_map<T, std::stack<T>>& H);
};

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 T& u, const T& v) {
  if (!this->G.adjList[u].contains(v)) return;

  std::unordered_map<T, std::stack<T>> H;
  for (const auto& w : this->G.vertices()) {
    if (RT[w].contains(u, v)) {
      if (E[v].inc.size() > 1)
        H[w].push(v);
      else {
        TC[w][v] = false;
        for (const auto& c : E[v].out)
          H[w].push(c);
      }
      RT[w].adjList[u].erase(v);
    }
  }
  E[u].out.erase(v);
  E[v].inc.erase(u);
  this->G.remove(u, v);

  repairTrees(H);
}

template<typename T>
void Italiano<T>::repairTrees(std::unordered_map<T, std::stack<T>>& H) {
  for (const auto& z : this->G.vertices()) {
    while (H[z].size() > 0) {
      const auto& h = H[z].top();
      bool found = false;

      for (const auto& i : E[h].inc) {
        if (RT[z].contains(z, i)) {
          RT[z].adjList[i].insert(h);
          found = true;
          break;
        }
      }
      H[z].pop();
      if (!found) {
        TC[z][h] = false;
        for (const auto& o : E[h].out) {
          H[z].push(o);
        }
      }
    }
  }
}

} // namespace algo

#endif