#ifndef TARJAN_H_
#define TARJAN_H_

#include "graph/graph.h"
using namespace graph;

#include <algorithm>
#include <iostream>
#include <ranges>
#include <stack>

namespace algo {

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

  //
  std::vector<std::vector<T>> run();

  //
  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;
  };

  Graph<T> G;
  std::stack<T> S;
  std::int16_t index = 0;
  std::map<T, Payload> vp;
  std::vector<std::vector<T>> SCCs;
};

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

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

template<typename T>
std::vector<std::vector<T>> Tarjan<T>::run() {

  for (const auto& v : G.vertices) {
    if (vp[v].index == -1) {
      strongConnect(v);
    }
  }

  return SCCs;
}

} // namespace algo

#endif