#include <iostream>
#include <map>
#include <set>
#include <string>
#include <vector>

using namespace std;

map<char, string> mapGrammar(vector <string> grammar)       // returns a map from input strings. Variables = keys, Productions = values
{
    map<char, string> tempMap;      // creates a temporary map to be returned
    for (int i = 0; i < grammar.size(); i++)        // parse through the vector of strings
    {
        string tmp = "";        // creates an initial string for the production
        for (int j = 3; j < (grammar.at(i).size()); j++)    // looks at everything after "->"
        {
            tmp += grammar[i][j];   // adds production to string, character by character
        }
        cout << "map @ " << i << ": " << grammar[i][0] << " -> " << tmp << endl;
        tempMap.insert({ grammar[i][0], tmp });     // inserts the variable and production to the map
    }
    return tempMap;     // returns completed map
}

map<char, set<char>> mapSet(vector <string> grammar)        // returs a map from input strings. Variables = keys. Sets = productions
{
    map<char, set<char>> tempMap;       // creates a temporary map to be returned
    for (int i = 0; i < grammar.size(); i++)        // parse through the vector of strings
    {
        tempMap.insert({ grammar[i][0], {} });      // creates empty sets for each variable
    }
    return tempMap;     // returns completed map
}

vector<string> getSubProd(string Prod)      // returns a vector of all the individual sub-productions of a variable. ( D->Aa|d|# returns 0:Aa, 1:d, 2:# )
{
    vector<string> temp;    // create a temp vector to be returned
    string str;     // creates a tmp string that takes in each individual sub-production
    for (int i = 0; i < Prod.size(); i++)   // parse through the entire production
    {
        if (Prod.at(i) == '|')      // if OR statement is reached
        {
            temp.push_back(str);    // add the string before the OR statement to the vector list.
            str = "";       // resets the temp string
        }
        else
        {
            str += Prod.at(i);      // add the current character to the sub-Production string
        }
    }
    if (str.size() > 0)     // no OR statement found
    {
        temp.push_back(str);    // add the string to the vector list
    }
    return temp;    // returns the vector list of strings.
}

set<char> getUnion(const set<char>& a, const set<char>& b)      // union two sets
{
    set<char> result = a;
    result.insert(b.begin(), b.end());
    return result;
}

void printSet(set<char> set)        // print the contents of a set
{
    cout << "{ ";
    for (auto it = set.begin(); it != set.end(); it++)
    {
        cout << *it << ", ";
    }
    cout << "}";
}

set<char> First(map<char, string> &map, char S)     // returns the First() set of a Variable. Takes in the map to be used as reference, and a character you want to return the First() set of.
{
    set<char> finalSet;     // set that will be returned to the function
    string production;      // entire string value that is associated with each variable key. Will be broken down later by getSubProd()

    auto it = map.find(S);  // Look for the key associated with input character
    if (it != map.end())    // if foudn
    {
        production = it->second;    // assign production to the string value mapped to Variable 'S'
    }
    else    // 'S' not found in map keys
    {
        finalSet.insert(S);     // insert unidentified Terminal into finalSet
        return finalSet;
    }

    vector<string> initSet;
    initSet = getSubProd(production);       // a vector of the sub-Productions of char 'S'

    for (int i = 0; i < initSet.size(); i++)    // check each sub-Production. ( Aa | d | # )
    {
        int j = 0;
        char firstChar = initSet[i][j];     // firstChar = first character of each sub-production ( A | d | # )
        if (map.find(firstChar) == map.end())       // if first character is a Terminal
        {
            finalSet.insert(firstChar);     // add character to finalSet
        }
        else        // first character is a Non-Terminal
        {
            finalSet = getUnion(finalSet, First(map, firstChar));   // unions finalSet with finalSet of the Non-Terminal
            while (finalSet.find('#') != finalSet.end() && (j + 1) < initSet.at(i).size())      // while epsilon is returned from First() and it's not the last Variable in the production
            {
                finalSet.erase('#');        // erase epsilon
                char nextChar = initSet[i][++j];        // look at the next character in the production
                finalSet = getUnion(finalSet, First(map, nextChar));    // find First() of the next character
            }
        }
    }
    return finalSet;
}

set<char> Follow(map<char, string> &Map, char S)
{
    set<char> finalSet;     // set that will be returned to the function
    string production;

    if (S == 'S')       // S is the first Key in the map. AKA starting production
    {
        finalSet.insert('$');
    }

    auto itr = Map.begin();

    for (itr = Map.begin(); itr != Map.end(); itr++)        // for each key in the map
    {
        production = itr->second;           // production = production string in map
        vector<string> initSet;
        initSet = getSubProd(production);       // create a vector of sub-Productions givin the production

        for (int j = 0; j < initSet.size(); j++)        // for each sub-Production in the vector
        {
            for (int k = 0; k < initSet.at(j).size(); k++)      // for each character in the sub-Production
            {
                if (initSet[j][k] == S)     // found Variable in sub-Production
                {
                    int l = k + 1;      // points to next character in the sub-Production
                    if (l >= initSet.at(j).size())      // if Variable is last character of sub-Production
                    {
                        finalSet = getUnion(finalSet, Follow(Map, itr->first));     // Follow() of Variable = Follow() of Variable on left side of the production
                    }
                    else
                    {
                        char nextChar = initSet[j][l];  // next character in the sub-Production
                        finalSet = getUnion(finalSet, First(Map, nextChar));  // Follow() = First(nextChar)
                        while (finalSet.find('#') != finalSet.end() && (l+1) < initSet.at(j).size())        // next character is a Non-Terminal that contains '#' && Not at end of sub-Production
                        {
                            finalSet.erase('#');
                            finalSet = getUnion(finalSet, First(Map, initSet[j][++l]));     // Moves onto the next character in the sub-Production
                        }
                        if (finalSet.find('#') != finalSet.end())       // if the last character returned an epsilon
                        {
                            finalSet.erase('#');
                            finalSet = getUnion(finalSet, Follow(Map, itr->first)); // Follow() of Variable = Follow() of Variable on left side of the production
                        }
                    }
                }
            }
        }
    }
    return finalSet;
}

int main()
{
    vector <string> grammar;

    grammar.push_back("S->ABCD");
    grammar.push_back("A->a|#");
    grammar.push_back("B->CD|b");
    grammar.push_back("C->c|#");
    grammar.push_back("D->Aa|d|#");

    cout << endl;

    map<char, string> gramMap;
    gramMap = mapGrammar(grammar);      // creates a map of the grammar

    map<char, set<char>> firstMap, followMap;
    firstMap = mapSet(grammar);     // creates a map for the first set
    followMap = mapSet(grammar);    // creates a map for the follow set

    cout << endl;

    for (int i = 0; i < grammar.size(); i++)    // assigns First() to each Variable in the grammar
    {
        set<char> temp = First(gramMap, grammar[i][0]);
        firstMap.find(grammar[i][0])->second = temp;
        cout << "First(" << grammar[i][0] << ") = ";
        printSet(temp);
        cout << endl;
    }

    cout << endl;

    for (int i = 0; i < grammar.size(); i++)    // assigns First() to each Variable in the grammar
    {
        set<char> temp = Follow(gramMap, grammar[i][0]);
        followMap.find(grammar[i][0])->second = temp;
        cout << "Follow(" << grammar[i][0] << ") = ";
        printSet(temp);
        cout << endl;
    }

    cout << endl;

    system("pause");
    return 0;
}