#ifndef BREADTH_FIRST_SEARCH_H_
#define BREADTH_FIRST_SEARCH_H_

#include <map>
#include <set>
#include <queue>

using namespace graph;

namespace algo {

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

  BreadthFirstSearch(std::map<T, std::set<T>> adjMatrix)
    : adjMatrix(adjMatrix) {}

  // Traverse whole graph using the BFS search, and save the tree graph
  // which is created when visiting new vertices (Breadth First Tree)
  std::map<T, std::set<T>> execute(const T& root);

  // Search if target vertex exists in graph
  bool query(const T& root, const T& target);

  void setGraph(std::map<T, std::set<T>> adjMatrix);
private:
  // Represents the graph on which the algorithm will be executed
  std::map<T, std::set<T>> adjMatrix;
}; 

template<typename T>
std::map<T, std::set<T>> BreadthFirstSearch<T>::execute(const T& root) {
  std::map<T, std::set<T>> tree;
  std::map<T, bool> visited;
  std::queue<T> Q;
  Q.push(root);
  visited[root] = true;

  while (!Q.empty()) {
    const auto v = Q.front();
    Q.pop();

    for (const auto& u : adjMatrix[v]) {
      if (!visited[u]) {
        visited[u] = true;
        tree[v].insert(u);
        Q.push(u);
      }
    }
  }
  return tree;
}

template<typename T>
bool BreadthFirstSearch<T>::query(const T& root, const T& target) {
  std::map<T, bool> visited;
  std::queue<T> Q;
  Q.push(root);
  visited[root] = true;

  while (!Q.empty()) {
    const auto v = Q.front();
    Q.pop();

    if (v == target) return true;

    for (const auto& u : adjMatrix[v]) {
      if (!visited[u]) {
        visited[u] = true;
        Q.push(u);
      }
    }
  }
  return false;
}

template<typename T>
void BreadthFirstSearch<T>::setGraph(std::map<T, std::set<T>> adjMatrix) {
  this->adjMatrix = adjMatrix;
}

} // namespace algo

#endif