import enum

from regex import lit, concat, bar, star
from nfa import NFA, copy_nfa, prettyprint

def remove_states(nfa):
	start, accept, transitions = nfa
	states = transitions.keys()

	states_to_remove = [i for i in states if i != start and i not in accept]

	while len(states_to_remove) > 0:
		# Select a state to remove this round
		removed_state = states_to_remove.pop()
		print('\nRemoving state:', removed_state)#debg

		# Remove loops from this state back into itself
		if removed_state in transitions[removed_state]:
			loop_condition = transitions[removed_state][removed_state]
			del transitions[removed_state][removed_state]

			# Prepend (condition)* to all transitions leading out
			# of this state
			for to_state in transitions[removed_state]:
				condition = transitions[removed_state][to_state]
				transitions[removed_state][to_state] = concat(star(loop_condition), condition)

		print(); prettyprint(nfa)#debg

		# Rewrite all transitions A→this→B as A→B transitions
		#
		# If the condition A→this is foo and this→B is bar, the
		# condition for A→B becomes simply foobar
		#
		# Since we've removed all loops back into this state, this
		# results in there being no transitions into this state
		for from_state in transitions:
			if removed_state in transitions[from_state]:
				# Create a list of new transitions to add to the
				# transition table for from_state
				new_transitions = {}
				condition_to_here = transitions[from_state][removed_state]
				for to_state in transitions[removed_state]:
					condition_from_here = transitions[removed_state][to_state]
					new_transitions[to_state] = concat(condition_to_here, condition_from_here)

				# Remove the transition to the state being deleted
				del transitions[from_state][removed_state]

				# Add the new transitions
				# Since they may lead to the same place as
				# already-existing transitions, we may need to
				# combine the conditions with pre-existing ones
				for to_state in new_transitions:
					if to_state in transitions[from_state]:
						# Already a transition leading
						# to the same state
						# If its condition is foo and
						# ours is bar, then the new
						# condition will be foo|bar
						other_condition = transitions[from_state][to_state]
						our_condition = new_transitions[to_state]
						transitions[from_state][to_state] = bar(other_condition, our_condition)

					else:
						# No pre-existing transition
						transitions[from_state][to_state] = new_transitions[to_state]

		# Finally, remove the state we no longer need
		del transitions[removed_state]

		print(); prettyprint(nfa)#debg

	return NFA(start, accept, transitions)

def to_regex(nfa):
	# Rewrite the NFA so that there are no transitions leading in to the
	# start state or any leading out of an accept state. The easy way to
	# do this is by creating a new start state that leads to the old one
	# with empty condition (i.e. it consumes no input), and creating a new
	# accept state that has similar empty condition transitions from all
	# the old ones. Since we have an NFA and not a DFA, that operation is
	# safe
	#
	# As a bonus, this rewrite gives us two useful properties:
	# a) There is exactly one start state and one accept state
	# b) After running remove_state() there will be only one transition,
	#    that of start to accept
	#
	# S
	class _(enum.Enum): start, end = range(2)

	start, accept, transitions = copy_nfa(nfa)

	# Add new start state
	transitions[_.start] = {start: lit('')}

	# Add new accept state and transitions to it
	transitions[_.end] = {}
	for state in accept:
		transitions[state][_.end] = lit('')

	# Package everything into a new NFA
	nfa = NFA(_.start, [_.end], transitions)

	print();prettyprint(nfa)#debg

	processed = remove_states(nfa)

	return processed.transitions[_.start][_.end]

def main():
	nfa = NFA('start', ['0'], {
		'start': {'1': lit('i'), '2': lit('d')},
		'0': {'1': lit('i'), '2': lit('d')},
		'1': {'0': lit('d'), '2': lit('i')},
		'2': {'0': lit('i'), '1': lit('d')}
	})

	prettyprint(nfa)

	regex = to_regex(nfa)

	print(repr(regex))
	print(regex)

if __name__ == '__main__':
	main()