#!/usr/bin/env python

import pygraphviz as pyg

a = "q0 b a a a b"
func = {
    ("q0", "a"):[("q0", "_", "R"), ("q0", "b", "R")],
    ("q0", "b"):[("q1", "_", "R")],
    ("q1", "a"):[("q0", "_", "R"), ("q0", "b", "L"), ("q1", "a", "R")],
    ("q1", "b"):[("q1", "_", "R")],
    ("q1", "_"):[("qacc", "_", "R")],
}

def next_conf(config, function, accept="qacc", reject="qrej"):
    """Returns a list of configurations as result of applying the function
    to the current configuration.
    
    Arguments:

    config - string with tokens as the tape. States are represented by q*
    function - dict like (state, symbol):[(state, symbol, R|L)]
    """

    tokens = config.split()

    ret = []

    state = [x for x in tokens if x.startswith("q")][0]

    try:
        symbol = tokens[tokens.index(state)+1]
    except IndexError:
        symbol = "_"

    if (state, symbol) not in function:
        return []

    for new_state, new_symbol, direction in function[(state, symbol)]:

        state_index = tokens.index(state)
        new_tokens = tokens[:]

        try:
            new_tokens.pop(state_index+1)
        except IndexError:
            pass

        new_tokens.insert(state_index+1, new_symbol)
        new_tokens.remove(state)

        if direction == "L": # <--
            new_tokens.insert(max(0, state_index-1), new_state)
        else: # -->
            new_tokens.insert(state_index+1, new_state)

        if new_state == accept:
            status = "accept"
        elif new_state == reject:
            status = "reject"
        else:
            status = "running"
            
        ret.append((" ".join(new_tokens), status))

    return ret


class Node:
    def __init__(self, data, parent, visited=False):
        self.data = data
        self.parent = parent
        self.childs = []
        self.visited = visited

    def add_child(self, data):
        node = Node(data, self)
        self.childs.append(node)

        return node

def run_breadth_first(init_config, function):
    """Runs a turing machine, building a tree with the history of computation.
    """
    root = Node(init_config, None)
    nodes = [root]
    i = 0

    while True:
        try:
            node = nodes[i]
        except IndexError:
            break

        if node.visited:
            i += 1
            continue

        configs = next_conf(node.data, function)

        for conf in configs:
            new_node = node.add_child(conf[0])

            nodes.append(new_node)

            if conf[1] in ["accept", "reject"]:
                new_node.visited = True

        node.visited = True
        i += 1

    return nodes

def tree2dot(nodes):
    """Writes a tree of configurations to a png file."""
    g = pyg.AGraph(directed=True)
    for node in nodes:
        print node.data
        for child in node.childs:
            g.add_edge(node.data, child.data)

    g.layout(prog="dot")
    g.draw("file.png")
    return g

if __name__ == "__main__":
    nodes = run_breadth_first(a, func) 
    tree2dot(nodes)