#ifndef DECREMENTAL_SCC_H_
#define DECREMENTAL_SCC_H_

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

using namespace graph;

namespace algo {

template<typename T>
class DecrementalSCC : public RodittyZwick<T> {
public:
  DecrementalSCC() = default;

  DecrementalSCC(Digraph<T> G) : G(G) {}

  void init();

  void findSCC();

  // Return true if u and v are in the same SCC
  bool query(const T& u, const T& v);

  // Remove edge (u,v) and update A accordingly for fast checking query
  void remove(const T& u, const T& v);

  void setGraph(Digraph<T> G);
private:
  Digraph<T> G;

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

  // Maintain in-out bfs trees
  std::map<T, BreadthFirstTree<T>> inTree;
  std::map<T, BreadthFirstTree<T>> outTree;

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

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

template<typename T>
void DecrementalSCC<T>::findSCC() {
  auto SCCs = Tarjan<T>(G.adjMatrix).execute();
  
  for (auto& C : SCCs) {
    const auto& w = C.id;

    for (const auto& v : std::views::keys(C.adjMatrix))
      A[v] = w;
    
    outTree[w] = BreadthFirstTree<T>(C, w);
    inTree[w] = BreadthFirstTree<T>(C.reverse(), w);

    SCC[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) {
  const auto& w = A[u];
  SCC[w].remove(u, v);
  G.remove(u, v);

  // If u and v are not in the same SCC, do nothing
  if (A[u] != A[v]) return;

  // If edge (u,v) is not contained in both inTree and outTree do nothing
  if (!inTree[w].adjMatrix[u].contains(v) &&
      !outTree[w].adjMatrix[u].contains(v))
    return;

  // Update In(w) and Out(w)
  outTree[w] = BreadthFirstTree<T>(SCC[w], w);
  inTree[w] = BreadthFirstTree<T>(SCC[w].reverse(), w);

  // If a SCC is broken, compute all SCCs again
  if (!inTree[w].adjMatrix.count(u) || !outTree[w].adjMatrix.count(v))
    findSCC();
}

template<typename T>
void DecrementalSCC<T>::setGraph(Digraph<T> G) {
  this->G = G;
}

}; // namespace algo

#endif