import copy
from collections import defaultdict
from tabulate import tabulate  # Ensure you have tabulate installed: pip install tabulate

# Perform grammar augmentation
def grammarAugmentation(rules, nonterm_userdef, start_symbol):
    newRules = []
    newChar = start_symbol + "'"
    while newChar in nonterm_userdef:
        newChar += "'"
    # Adding rule to bring start symbol to RHS with a dot at the beginning
    newRules.append([newChar, ['.', start_symbol]])
    # Process existing rules
    for rule in rules:
        lhs, rhs = rule.split("->")
        lhs = lhs.strip()
        rhs = rhs.strip()
        multirhs = rhs.split('|')
        for rhs1 in multirhs:
            rhs1 = rhs1.strip().split()
            # ADD dot pointer at the start of RHS
            rhs1.insert(0, '.')  # Dot at position 0
            newRules.append([lhs, rhs1])
    return newRules

# Find closure
def findClosure(items, separatedRulesList, nonterm_userdef):
    closureSet = list(items)  # Start with the current items
    added = True
    while added:
        added = False
        new_items = []
        for rule in closureSet:
            if '.' in rule[1]:
                dot_pos = rule[1].index('.')
                # Check if dot is not at the end
                if dot_pos < len(rule[1]) - 1:
                    symbol_after_dot = rule[1][dot_pos + 1]
                    if symbol_after_dot in nonterm_userdef:
                        # Add all rules where lhs == symbol_after_dot
                        for prod in separatedRulesList:
                            if prod[0] == symbol_after_dot and prod not in closureSet and prod not in new_items:
                                new_items.append(prod)
        # Add any new items found
        for new_item in new_items:
            if new_item not in closureSet:
                closureSet.append(new_item)
                added = True
    return closureSet

# Compute GOTO
def compute_GOTO(current_items, symbol, separatedRulesList, nonterm_userdef):
    moved = []
    for rule in current_items:
        if '.' in rule[1]:
            dot_pos = rule[1].index('.')
            if dot_pos < len(rule[1]) - 1 and rule[1][dot_pos + 1] == symbol:
                # Move the dot past the symbol
                new_rule = copy.deepcopy(rule)
                # Swap the dot with the symbol
                new_rule[1][dot_pos], new_rule[1][dot_pos + 1] = new_rule[1][dot_pos + 1], new_rule[1][dot_pos]
                moved.append(new_rule)
    # Get the closure of the moved items
    return findClosure(moved, separatedRulesList, nonterm_userdef) if moved else []

# Calculation of FIRST
def first(symbol, diction, term_userdef, firsts, visited=None):
    if visited is None:
        visited = set()
    if symbol in term_userdef:
        return {symbol}
    if symbol in visited:
        return set()
    visited.add(symbol)
    result = set()
    for rule in diction.get(symbol, []):
        if rule[0] == '#':  # Epsilon
            result.add('#')
        else:
            for sym in rule:
                first_sym = first(sym, diction, term_userdef, firsts, visited)
                result.update(first_sym - {'#'})
                if '#' not in first_sym:
                    break
            else:
                result.add('#')
    return result

# Calculation of FOLLOW
def follow_symbol(symbol, diction, nonterm_userdef, term_userdef, firsts, follows):
    if not follows[symbol]:
        if symbol == start_symbol_augmented:
            follows[symbol].add('$')
    # Iterate over all productions to find occurrences of the symbol
    for lhs in diction:
        for rhs in diction[lhs]:
            if symbol in rhs:
                indices = [i for i, sym in enumerate(rhs) if sym == symbol]
                for index in indices:
                    if index + 1 < len(rhs):
                        next_sym = rhs[index + 1]
                        first_next = first(next_sym, diction, term_userdef, firsts)
                        follows[symbol].update(first_next - {'#'})
                        if '#' in first_next:
                            follows[symbol].update(follow_symbol(lhs, diction, nonterm_userdef, term_userdef, firsts, follows))
                    else:
                        if lhs != symbol:
                            follows[symbol].update(follow_symbol(lhs, diction, nonterm_userdef, term_userdef, firsts, follows))
    return follows[symbol]

# Create FIRST and FOLLOW sets
def compute_first_follow(diction, nonterm_userdef, term_userdef):
    firsts = {}
    for nt in nonterm_userdef:
        firsts[nt] = first(nt, diction, term_userdef, firsts)
    
    follows = defaultdict(set)
    for nt in nonterm_userdef:
        follow_symbol(nt, diction, nonterm_userdef, term_userdef, firsts, follows)
    
    return firsts, follows

# Create parse table
def createParseTable(statesDict, stateMap, separatedRulesList, diction, firsts, follows, term_userdef, nonterm_userdef):
    action_table = defaultdict(dict)
    goto_table = defaultdict(dict)
    prod_num = {}
    num = 1
    # Number the productions (skip augmented start rule at 0)
    for rule in separatedRulesList:
        if rule[0] == start_symbol_augmented:
            prod_num[0] = tuple((rule[0], tuple(rule[1])))
        else:
            prod_num[num] = tuple((rule[0], tuple(rule[1])))
            num += 1

    for state in statesDict:
        for rule in statesDict[state]:
            lhs, rhs = rule[0], rule[1]
            if rhs[-1] == '.':
                if lhs == start_symbol_augmented:
                    action_table[state]['$'] = 'ACC'
                else:
                    # Find production number
                    for num_p, prod in prod_num.items():
                        if prod == tuple((lhs, tuple(rhs))):
                            break
                    # Add reduce actions for all symbols in FOLLOW(lhs)
                    for terminal in follows[lhs]:
                        if terminal in action_table[state]:
                            action_table[state][terminal] += f'/R{num_p}'
                        else:
                            action_table[state][terminal] = f'R{num_p}'
            else:
                dot_pos = rhs.index('.')
                if dot_pos < len(rhs) - 1:
                    symbol = rhs[dot_pos + 1]
                    if symbol in term_userdef:
                        # Shift action
                        goto_result = compute_GOTO(statesDict[state], symbol, separatedRulesList, nonterm_userdef)
                        if not goto_result:
                            continue
                        goto_key = tuple((r[0], tuple(r[1])) for r in goto_result)
                        if goto_key in stateMap:
                            next_state = stateMap[goto_key]
                            if symbol in action_table[state]:
                                action_table[state][symbol] += f'/S{next_state}'
                            else:
                                action_table[state][symbol] = f'S{next_state}'
        # Populate GOTO table
        for symbol in nonterm_userdef:
            goto_result = compute_GOTO(statesDict[state], symbol, separatedRulesList, nonterm_userdef)
            if not goto_result:
                continue
            goto_key = tuple((r[0], tuple(r[1])) for r in goto_result)
            if goto_key in stateMap:
                goto_state = stateMap[goto_key]
                goto_table[state][symbol] = goto_state
    return action_table, goto_table, prod_num

# Print parse table
def printParseTable(action_table, goto_table, term_userdef, nonterm_userdef):
    headers = ['State'] + term_userdef + ['$'] + nonterm_userdef
    table = []
    for state in sorted(statesDict.keys()):
        row = [f'I{state}']
        for symbol in term_userdef + ['$']:
            row.append(action_table[state].get(symbol, ''))
        for symbol in nonterm_userdef:
            row.append(goto_table[state].get(symbol, ''))
        table.append(row)
    print("\nSLR(1) Parsing Table:\n")
    print(tabulate(table, headers, tablefmt="pretty"))

# Parsing the string
def parseString(string, action_table, goto_table, prod_num, diction):
    stack = [0]
    pointer = 0
    # Split the string into tokens; ensure proper tokenization
    symbol_list = string.strip().split() + ['$']
    headers = ['Stack', 'Input', 'Action']  # Ensure 'headers' is defined
    data = []
    accepted = False
    while True:
        current_state = stack[-1]
        if pointer >= len(symbol_list):
            current_symbol = '$'
        else:
            current_symbol = symbol_list[pointer]
        action = action_table.get(current_state, {}).get(current_symbol, '')
        stack_content = ' '.join(map(str, stack))
        input_remaining = ' '.join(symbol_list[pointer:])
        if '/' in action:
            # Handle multiple actions (conflicts), which shouldn't occur in SLR(1)
            actions = action.split('/')
            action = actions[0]  # Choose the first action; ideally report conflict
        if action.startswith('S'):
            try:
                next_state = int(action[1:])
            except ValueError:
                data.append([stack_content, input_remaining, f'Invalid Shift State: {action}'])
                break
            data.append([stack_content, input_remaining, f'Shift to state {next_state}'])
            stack.append(current_symbol)
            stack.append(next_state)
            pointer += 1
        elif action.startswith('R'):
            try:
                prod_number = int(action[1:])
            except ValueError:
                data.append([stack_content, input_remaining, f'Invalid Reduce Action: {action}'])
                break
            if prod_number not in prod_num:
                data.append([stack_content, input_remaining, f'Unknown Production Number: R{prod_number}'])
                break
            lhs, rhs = prod_num[prod_number]
            if rhs != ('#',):  # Assuming '#' denotes epsilon
                pop_length = 2 * len(rhs)  # Each symbol on stack has state
                if pop_length > len(stack):
                    data.append([stack_content, input_remaining, f'Error: Stack underflow during Reduce R{prod_number}'])
                    break
                stack = stack[:-pop_length]
            else:
                # Epsilon production, no symbols to pop
                pass
            current_state = stack[-1]
            stack.append(lhs)
            if lhs in goto_table[current_state]:
                goto_state = goto_table[current_state][lhs]
                stack.append(goto_state)
                data.append([stack_content, input_remaining, f'Reduce using production {prod_number}: {lhs} -> {" ".join(rhs)}'])
            else:
                data.append([stack_content, input_remaining, f'Error: No GOTO state for {lhs} from state {current_state}'])
                break
        elif action == 'ACC':
            data.append([stack_content, input_remaining, 'Accept'])
            accepted = True
            break
        else:
            data.append([stack_content, input_remaining, 'Error'])
            break
    # Print parsing steps
    print("\nParsing Steps:\n")
    print(tabulate(data, headers=headers, tablefmt="pretty"))
    return accepted

# Utility function to print rules
def printResult(rules):
    for rule in rules:
        print(f"{rule[0]} -> {' '.join(rule[1])}")

# *** MAIN *** Driver Code

# Example sample set 01
rules = [
    "E -> E + T | T",
    "T -> T * F | F",
    "F -> ( E ) | id"
]

nonterm_userdef = ['E', 'T', 'F']
term_userdef = ['id', '+', '*', '(', ')']
start_symbol = nonterm_userdef[0]

# Example sample set 02 (Uncomment to test)
# rules = [
#     "S -> a Xd | byda Ye | bXe",
#     "X -> c",
#     "Y -> c"
# ]
# nonterm_userdef = ['S', 'X', 'Y']
# term_userdef = ['a', 'b', 'c', 'd', 'e']
# start_symbol = nonterm_userdef[0]

print("\nOriginal Grammar Input:\n")
for y in rules:
    print(y)

# Processed rules after augmentation
separatedRulesList = grammarAugmentation(rules, nonterm_userdef, start_symbol)
print("\nGrammar after Augmentation:\n")
printResult(separatedRulesList)

# Initialize statesDict and stateMap
statesDict = {}
stateMap = {}
# Find the augmented start symbol
start_symbol_augmented = separatedRulesList[0][0]
# Find closure for the initial state
initial_items = [separatedRulesList[0]]
closure_result = findClosure(initial_items, separatedRulesList, nonterm_userdef)
statesDict[0] = closure_result
# Correctly convert to tuples for hashing
closure_key = tuple((rule[0], tuple(rule[1])) for rule in closure_result)
stateMap[closure_key] = 0
stateCount = 1

# Generate states
while True:
    added_new_state = False
    # Iterate over existing states
    for state in list(statesDict.keys()):
        current_items = statesDict[state]
        symbols = set()
        for rule in current_items:
            if '.' in rule[1]:
                dot_pos = rule[1].index('.')
                if dot_pos < len(rule[1]) - 1:
                    symbols.add(rule[1][dot_pos + 1])
        for symbol in symbols:
            goto_result = compute_GOTO(current_items, symbol, separatedRulesList, nonterm_userdef)
            if not goto_result:
                continue
            goto_key = tuple((r[0], tuple(r[1])) for r in goto_result)
            if goto_key not in stateMap:
                statesDict[stateCount] = goto_result
                stateMap[goto_key] = stateCount
                stateCount += 1
                added_new_state = True
    if not added_new_state:
        break

# Print generated states
print("\nStates Generated:\n")
for st in sorted(statesDict.keys()):
    print(f"State I{st}:")
    printResult(statesDict[st])
    print()

# Create dictionary for productions
diction = defaultdict(list)
for rule in separatedRulesList:
    lhs, rhs = rule
    diction[lhs].append(rhs)

# Compute FIRST and FOLLOW sets
firsts, follows = compute_first_follow(diction, nonterm_userdef, term_userdef)

# Create parse table
action_table, goto_table, prod_num = createParseTable(statesDict, stateMap, separatedRulesList, diction, firsts, follows, term_userdef, nonterm_userdef)

# Print parse table
printParseTable(action_table, goto_table, term_userdef, nonterm_userdef)

# Input string to parse
string = input("\nEnter the string to be parsed (tokens separated by space, use '$' for end): ")
string = string.strip() + ' $'  # Ensure end symbol

# Parsing the string
accepted = parseString(string, action_table, goto_table, prod_num, diction)

if accepted:
    print("\nThe string is parsable and accepted!")
else:
    print("\nThe string is not parsable and is rejected.")


output i am getting:
Original Grammar Input:

E -> E + T | T
T -> T * F | F
F -> ( E ) | id

Grammar after Augmentation:

E' -> . E
E -> . E + T
E -> . T
T -> . T * F
T -> . F
F -> . ( E )
F -> . id

States Generated:

State I0:
E' -> . E
E -> . E + T
E -> . T
T -> . T * F
T -> . F
F -> . ( E )
F -> . id

State I1:
E' -> E .
E -> E . + T

State I2:
E -> T .
T -> T . * F

State I3:
F -> ( . E )
E -> . E + T
E -> . T
T -> . T * F
T -> . F
F -> . ( E )
F -> . id

State I4:
F -> id .

State I5:
T -> F .

State I6:
E -> E + . T
T -> . T * F
T -> . F
F -> . ( E )
F -> . id

State I7:
T -> T * . F
F -> . ( E )
F -> . id

State I8:
F -> ( E . )
E -> E . + T

State I9:
E -> E + T .
T -> T . * F

State I10:
T -> T * F .

State I11:
F -> ( E ) .


SLR(1) Parsing Table:

+-------+----+----+----+----+-----+-----+---+---+----+
| State | id | +  | *  | (  |  )  |  $  | E | T | F  |
+-------+----+----+----+----+-----+-----+---+---+----+
|  I0   | S4 |    |    | S3 |     |     | 1 | 2 | 5  |
|  I1   |    | S6 |    |    |     | ACC |   |   |    |
|  I2   |    | R6 | S7 |    | R6  | R6  |   |   |    |
|  I3   | S4 |    |    | S3 |     |     | 8 | 2 | 5  |
|  I4   |    | R6 | R6 |    | R6  | R6  |   |   |    |
|  I5   |    | R6 | R6 |    | R6  | R6  |   |   |    |
|  I6   | S4 |    |    | S3 |     |     |   | 9 | 5  |
|  I7   | S4 |    |    | S3 |     |     |   |   | 10 |
|  I8   |    | S6 |    |    | S11 |     |   |   |    |
|  I9   |    | R6 | S7 |    | R6  | R6  |   |   |    |
|  I10  |    | R6 | R6 |    | R6  | R6  |   |   |    |
|  I11  |    | R6 | R6 |    | R6  | R6  |   |   |    |
+-------+----+----+----+----+-----+-----+---+---+----+

Enter the string to be parsed (tokens separated by space, use '$' for end): id + id * id$

Parsing Steps:

+--------+-------------------+-----------------------------------------+
| Stack  |       Input       |                 Action                  |
+--------+-------------------+-----------------------------------------+
|   0    | id + id * id$ $ $ |            Shift to state 4             |
| 0 id 4 |  + id * id$ $ $   | Error: Stack underflow during Reduce R6 |
+--------+-------------------+-----------------------------------------+

The string is not parsable and is rejected.