#include "bnf.h"

std::unordered_map<std::string, std::unordered_set<std::string>> Reverse(const BNF& bnf)
{
 std::unordered_map<std::string, std::unordered_set<std::string>> result;

 for (const auto& [from, to] : bnf) {
  for (const auto& list : to) {
   for (const auto& element : list) {
    auto i{result.find(element)};
    if (i != result.end()) // already present
     i->second.insert(from);
    else // new element
     result.emplace(element, std::unordered_set{from});
   }
  }
 }

 return result;
}

std::unordered_map<std::string, std::unordered_set<std::string>> reverseFirst(const BNF& bnf)
{
 std::unordered_map<std::string, std::unordered_set<std::string>> result;

 for (const auto& [from, to] : bnf) {
  for (const auto& list : to) {
   if (list.size() > 0) {
    const auto& element{list[0]};
    auto i{result.find(element)};
    if (i != result.end()) // already present
     i->second.insert(from);
    else // new element
     result.emplace(element, std::unordered_set{from});
   }
  }
 }

 return result;
}

BNF SubBNF(const BNF& bnf, const std::string& top)
{
 BNF result;

 std::vector<std::string> todo{top};

 while (!todo.empty()) {
  std::string current{todo.back()};
  todo.pop_back();

  auto it = bnf.find(current);
  if (it != bnf.end()) {
   // add current value
   result[it->first] = it->second;

   // add sub-tree values if not present yet, but available in original bnf
   for (auto& variant: it->second) {
    for (auto& child: variant) {
     if (result.find(child) == result.end() && bnf.find(child) != bnf.end()) {
      todo.push_back(child);
     }
    }
   }
  }
 }

 return result;
}

namespace {

 bool isHeadRecursive(const std::vector<std::string>& list, const std::string& symbol) {
  if (list.size() > 0 && list[0] == symbol)
   return true;
  return false;
 }

 bool isHeadRecursive(const std::vector<std::vector<std::string>>& lists, const std::string& symbol) {
  for (const auto& list: lists) {
   if (isHeadRecursive(list, symbol))
    return true;
  }
  return false;
 }

} // anonymous namespace

// Change head recursion to tail recursion
BNF removeHeadRecursion(const BNF& bnf)
{
 std::unordered_map<std::string, std::vector<std::vector<std::string>>> result;

 for (const auto& [from, to]: bnf) {
  if (isHeadRecursive(to, from)) { // modify rule by adding additional one
   std::string from_ext = from + "-ext";
   if (bnf.find(from_ext) != bnf.end())
    throw std::runtime_error("ICE: Symbol "s + from_ext + " already exists in original BNF");

   std::vector<std::vector<std::string>> to_new;
   std::vector<std::vector<std::string>> to_ext{{}}; // starts with empty list as first element
   for (const auto& list: to) {
    if (isHeadRecursive(list, from)) {
     std::vector<std::string> list_new{list.begin() + 1, list.end()};
     list_new.push_back(from_ext);
     to_ext.push_back(list_new);
    } else {
     std::vector<std::string> list_new{list.begin(), list.end()};
     list_new.push_back(from_ext);
     to_new.push_back(list_new);
    }
   }

   result[from] = to_new;
   result[from_ext] = to_ext;
  } else { // just add
   result[from] = to;
  }
 }

 return result;
}

bool isTerminal(const BNF& bnf, const std::string& symbol)
{
 return bnf.find(symbol) == bnf.end();
}

std::unordered_set<std::string> getTerminals(const BNF& bnf)
{
 std::unordered_set<std::string> result;

 for (const auto& [from, to] : bnf) {
  for (const auto& list : to) {
   for (const auto& element : list) {
    if (isTerminal(bnf, element))
     result.insert(element);
   }
  }
 }

 return result;
}

std::unordered_set<std::pair<std::string, size_t>, PairHash> getEmptyPositions(const BNF& bnf)
{
 std::unordered_set<std::pair<std::string, size_t>, PairHash> result;

 for (const auto& [from, to] : bnf) {
  for (size_t i = 0; i < to.size(); i++) {
   const auto& list{to[i]};
   if (list.size() == 0)
    result.insert({from, i});
  }
 }

 return result;
}

std::unordered_map<std::string, std::unordered_set<std::pair<std::string, index_t>, PairHash>> reversePosFirst(const BNF& bnf)
{
 std::unordered_map<std::string, std::unordered_set<std::pair<std::string, index_t>, PairHash>> result;

 for (const auto& [from, to] : bnf) {
  for (size_t i = 0; i < to.size(); i++) {
   const auto& list{to[i]};
   if (list.size() > 0) {
    const auto& element{list[0]};
    auto it{result.find(element)};
    if (it != result.end()) // already present
     it->second.insert({from, i});
    else // new element
     result.emplace(element, std::unordered_set<std::pair<std::string, index_t>, PairHash>{{from, i}});
   }
  }
 }

 return result;
}