#ifndef FRIGIONI_H_
#define FRIGIONI_H_

#include "algorithm/decremental_reachability.h"
#include "algorithm/roditty_zwick.h"

#include <utility>

using namespace graph;

namespace algo {

template<typename T>
class Frigioni : public DecrementalReachability<T> {
public:
  Frigioni() = default;

  Frigioni(Digraph<T> G) { this->G = G; }

  // Initialize the decremental maintenance data structure for general graphs
  void init() override;

  // Execute reachability query q(u, v) from vertex u to vertex v
  // in O(1) using the transitive closure matrix
  bool query(const T& u, const T& v) override;

  // Delete set of edges and explicitely maintain the transitive closure
  void remove(const std::vector<std::pair<T, T>>& edges) override;
private:
  // Transitive closure matrix, used to answer reachability queries in O(1)
  std::unordered_map<T, std::unordered_map<T, bool>> TC;

  // For each SCC, store a reachability tree created as a BFS tree
  std::unordered_map<T, BreadthFirstTree<T>> RT;

  // For each SCC, store collections of incoming, outgoing and internal edges
  struct Edges {
    std::set<std::pair<T, T>> in;
    std::set<std::pair<T, T>> inc;
    std::set<std::pair<T, T>> out;
  };
  std::unordered_map<T, Edges> E;
  
  // Decremental maintenance of strongly connected components
  RodittyZwick<T> rodittyZwick;
};

template<typename T>
void Frigioni<T>::init() {
  auto SCCs = Tarjan<T>(this->G.adjList).execute();
  rodittyZwick = RodittyZwick<T>(this->G);
  rodittyZwick.init();

  for (auto& scc : SCCs) {
    RT[scc.id] = BreadthFirstTree<T>(this->G, scc.id);
    for (const auto& u : this->G.vertices()) {
      for (const auto& v : this->G.adjList[u]) {
        if (scc.member(u)) {
          if (scc.member(v))
            E[scc.id].in.insert(std::make_pair(u, v));  
          else
            E[scc.id].out.insert(std::make_pair(u, v));
        } else if (scc.member(v)) {
          E[scc.id].inc.insert(std::make_pair(u, v));
        }
        TC[u][v] = false;
      }
    }
  }
  for (auto& scc : SCCs) {
    for (const auto& u : scc.vertices()) {
      for (const auto& v : RT[scc.id].vertices()) {
        TC[u][v] = true;
      }
    }
  }
}

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

template<typename T>
void Frigioni<T>::remove(const std::vector<std::pair<T, T>>& edges) {
  std::vector<std::pair<T, T>> Eint;
  std::vector<std::pair<T, T>> Eext;

  for (const auto& [u, v] : edges) {
    if (rodittyZwick.query(u, v))
      Eint.push_back({ u, v });
    else
      Eext.push_back({ u, v });
  }

  std::unordered_map<T, std::stack<T>> H;
  for (const auto& [u, v] : Eint) {
    this->G.remove(u, v);
    rodittyZwick.remove(u, v);

    if (!rodittyZwick.query(u, v)) {
      TC[u][v] = false;
      auto C = rodittyZwick.getSCCs();

      E[C[u].id].inc.erase({ u, v });
      E[C[u].id].out.erase({ u, v });
      E[C[u].id].in.erase({ u, v });
      auto D = E[C[u].id];
      E.erase(C[u].id);

      for (const auto& id : std::views::keys(C)) {
        if (RT[id].contains(id, v)) {
          if (E[C[v].id].inc.size() > 1)
            H[id].push(C[v].id);
          else {
            for (const auto& w : C[id].vertices())
              TC[w][v] = false;
            for (const auto& c : E[v].out)
              H[id].push(C[c.second].id);
          }
          RT[id].adjList[u].erase(v);
        }
      }

      for (const auto& [w, z] : D.inc) {
        E[C[z].id].inc.insert({ w, z });
      }

      for (const auto& [w, z] : D.out) {
        E[C[w].id].out.insert({ w, z });
      }

      for (const auto& [w, z] : D.in) {
        if (C[w] == C[z])
          E[C[w].id].in.insert({ w, z });
        else {
          E[C[w].id].out.insert({ w, z });
          E[C[z].id].in.insert({ w, z });
        }
      }
    } else {
      E[u].in.erase({ u, v });
    }
  }

  for (const auto& [u, v] : Eext) {
    this->G.remove(u, v);
    auto C = rodittyZwick.getSCCs();

    for (const auto& id : std::views::keys(C)) {
      if (RT[id].contains(id, v)) {
        if (E[C[v].id].inc.size() > 1)
          H[id].push(C[v].id);
        else {
          for (const auto& w : C[id].vertices())
            TC[w][v] = false;
          for (const auto& c : E[v].out)
            H[id].push(C[c.second].id);
        }
        RT[id].adjList[u].erase(v);
      }
    }
  }

  auto C = rodittyZwick.getSCCs();
  for (const auto& id : std::views::keys(rodittyZwick.getSCCs())) {
    while (H[id].size() > 0) {
      const auto& h = H[id].top();
      bool found = false;
      
      for (const auto& [u, v] : E[h].inc) {
        if (RT[id].contains(id, u)) {
          RT[id].adjList[u].insert(h);
          found = true;
          break;
        }
      }
      H[id].pop();
      if (!found) {
        auto C = rodittyZwick.getSCCs();
        for (const auto& w : C[id].vertices())
          TC[w][h] = false;

        for (const auto& [u, v] : E[h].out) {
          H[id].push(v);
        }
      }
    }
  }
}

} // namespace algo

#endif