#ifndef ITALIANO_H_
#define ITALIANO_H_

#include "algorithm/roditty_zwick.h"
#include "algorithm/tarjan.h"
#include "graph/breadth_first_tree.h"

#include <forward_list>
#include <iostream>

using namespace graph;

namespace algo {

template<typename T>
class Italiano : public RodittyZwick<T> {
public:
  Italiano(Digraph<T> G) : G(G) {}

  void init();

  bool query(const T& u, const T& v);

  void remove(const T& u, const T& v);
private:
  Digraph<T> G;

  // Transitive closure matrix, used to answer reachability queries in O(1)
  std::map<T, std::map<T, bool>> TC;

  // Each vertex's reachability tree
  std::map<T, BreadthFirstTree<T>> RT;

  // Incoming / Outgoing edges
  struct Edges {
    std::set<T> inc;
    std::set<T> out;
  };
  std::map<T, Edges> E;
};

template<typename T>
void Italiano<T>::init() {
  for (const auto& u : G.vertices()) {
    for (const auto& v : G.adjMatrix[u]) {
      E[v].inc.insert(u);
      E[u].out.insert(v);
      TC[u][v] = false;
    }
    RT[u] = BreadthFirstTree<T>(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) {
  std::map<T, std::stack<T>> H;

  for (const auto& w : 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);
      }
    }
  }

  E[u].out.erase(v);
  E[v].inc.erase(u);
  G.remove(u, v);

  for (const auto& z : G.vertices()) {
    RT[z].adjMatrix[u].erase(v);
    while (H[z].size() > 0) {
      auto& h = H[z].top();
      bool found = false;
      for (const auto& i : E[h].inc) {
        if (RT[z].contains(z, i)) {
          RT[z].adjMatrix[i].insert(h);
          H[z].pop();
          found = true;
          break;
        }
      }
      if (!found) {
        TC[u][v] = false;
        for (const auto& o : E[h].out) {
          if (RT[z].contains(z, o)) {
            H[z].push(o);
          }
        }
      }
    }
  }
}

} // namespace algo

#endif