#ifndef DECREMENTAL_SCC_H_
#define DECREMENTAL_SCC_H_

#include "graph/digraph.h"
#include "graph/scc.h"
#include "algorithm/tarjan.h"
#include "algorithm/bfs.h"
#include "algorithm/roditty_zwick.h"

using namespace graph;

namespace algo {

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

  //
  void init();

  //
  void findSCC();

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

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

  // Array used to answer strong connectivity queries in O(1) time
  std::map<T, T> A;

  // Connect SCC/SCC representative with 2 SPTs
  // pair.first: Shortest-paths out-tree
  // pair.second: Shortest-paths in-tree
  std::map<T, std::pair<Digraph<T>, Digraph<T>>> SPT;

  // Connect each representative with its SCC
  std::map<T, SCC<T>> connection;
};

template<typename T>
void DecrementalSCC<T>::init() {
  findSCC();
}

template<typename T>
void DecrementalSCC<T>::findSCC() {
  auto SCCs = Tarjan<T>(G).findSCC();

  for (auto& C : SCCs) {
    const auto& w = C.representative();

    // Create shortest-paths out-tree/in-tree
    auto outTree = Digraph<T>(BFS<T>(C).run(w));
    SPT[w] = std::make_pair(outTree, outTree.reverse());

    // Update A with current SCCs vertices
    for (const auto& v : C.vertices)
      A[v] = w;

    // Link SCC with its representative w
    connection[w] = C;
  }
}

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

template<typename T>
void DecrementalSCC<T>::remove(const T& u, const T& v) {
  // If u and v are not in the same SCC, do nothing
  if (A[u] != A[v]) return;

  const auto& w = A[u];

  // If edge e(u, v) is not contained in In(w) and Out(w), do nothing
  if (!SPT[w].first.vertices.contains(w) &&
      !SPT[w].second.vertices.contains(w)){
    connection[w].adjMatrix[u].erase(v);
    G.adjMatrix[u].erase(v);
    return;
  }
  // Update In(w) and Out(w)
  connection[w].adjMatrix[u].erase(v);
  G.adjMatrix[u].erase(v);
  auto inTree = Digraph<T>(BFS<T>(connection[w]).run(w));
  SPT[w] = std::make_pair(inTree, inTree.reverse());

  // If a SCC is broken, compute all SCCs again
  if (!SPT[w].second.vertices.contains(u) ||
      !SPT[w].first.vertices.contains(v)) {
    findSCC();
  }

}

}; // namespace algo

#endif