#ifndef TARJAN_H_
#define TARJAN_H_

#include "graph/scc.h"

#include <stack>
#include <vector>

using namespace graph;

namespace algo {

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

  std::vector<SCC<T>> execute();

  void strongConnect(const T& v);
private:
  // Necessary info about vertices when running Tarjan's algorithm
  struct Payload {
    std::int16_t index = -1;
    std::int16_t lowlink = -1;
    bool onStack = false;
  };

  Digraph<T> G;
  std::stack<T> S;
  std::int16_t index = 0;
  std::map<T, Payload> p;
  std::vector<SCC<T>> SCCs;
};

template<typename T>
void Tarjan<T>::strongConnect(const T& v) {
  p[v].index = p[v].lowlink = index++;
  p[v].onStack = true;
  S.push(v);
  
  for (const auto& w : G.adjMatrix[v]) {
    if (p[w].index == -1) {
      strongConnect(w);
      p[v].lowlink = std::min(p[v].lowlink, p[w].lowlink);
    } else if (p[w].onStack) {
      p[v].lowlink = std::min(p[v].lowlink, p[w].index);
    }
  }

  // If v is a root node, pop the stack and generate an SCC
  if (p[v].lowlink == p[v].index) {
    //std::vector<T> scc;
    std::map<T, std::set<T>> scc;
    bool finished = false;
    do {
      const auto w = S.top();
      S.pop();
      p[w].onStack = false;
      scc[w] = G.adjMatrix[w];
      finished = p[w].index == p[v].index;
    } while (!finished);
    SCCs.push_back(scc);
  }
}

template<typename T>
std::vector<SCC<T>> Tarjan<T>::execute() {
  for (auto& v : G.adjMatrix) {
    if (p[v.first].index == -1) {
      strongConnect(v.first);
    }
  }

  return SCCs;
}

} // namespace algo

#endif