# # # # # # # # # # # # # # # # # # # # # # #
#  Framework provided by Daan de Graaf(UvA) #
#  For Automaten en Formele Talen           #
#  Under Guidance of Inge Bethke(UvA)       #
# # # # # # # # # # # # # # # # # # # # # # #


class FA(object):
    ''' Finite automaton object'''
    # key:name value:state object
    states = {}

    def __init__(self, Q, Sigma, delta, s, F):
        '''Creates an FA object with given data and checks the input'''
        if s not in Q:
            exit('NodeError: Node "' + s + '" not in states set Q')

        for state in F:
            if state not in Q:
                exit('NodeError: Final state "' + str(state) +
                     '" not in states set Q')

        for transition in delta:
            for symbol in delta[transition].keys():
                if symbol not in Sigma:
                    exit('TransitionError: "' + symbol +
                         '" not found in alphabet Sigma')

        self.alphabet = Sigma
        self.transition_table = delta
        self.transition_nodes = set(self.transition_table.keys())

        # Create states and assign values
        for name in Q:
            start_state = (name == s)
            if name in self.transition_nodes:
                table = self.transition_table[name]
            else:
                table = None
            state = self.State(name, table, start_state, (name in F))

            if start_state:
                self.current_state = state

            self.states[name] = state

        self.start_state = self.states[s]
        self.final_states = [self.states[f] for f in F]

    def move(self, input):
        ''' Follow a transition for given input, returns true if succeeded'''
        try:
            self.current_state = self.states[
                self.current_state.transition_table[input]]
            return True
        except:
            return False

    def reset(self):
        self.current_state = self.start_state

    def plot(self, current=False, filename='automaton'):
        ''' Gives a plot of the FA in 'img/filename' specified by user with an
        option to plot the current state in red'''
        import graphviz as gv
        G = gv.Digraph(format='svg')
        G.graph_attr['rankdir'] = 'LR'
        for node in self.states.values():
            color = 'red' if node == self.current_state and current else ''
            G.node(str(node.name), shape=(
                    'doublecircle' if node in self.final_states else 'circle'
                   ), color=color)
            if node.transition_table is not None:
                for edge in node.transition_table:
                    G.edge(str(node.name), str(node.transition_table[edge]),
                           label=edge)

        path = 'img/' + filename
        G.render(path)

    class State(object):
        transition_table = dict()

        def __init__(self, name, transitions, start_state, final_state):
            self.name = name
            self.transition_table = transitions
            self.start_state = start_state
            self.final_state = final_state