"""
	Aho-Corasick string search algorithm.

	Author    : Wojciech Muła, wojciech_mula@poczta.onet.pl
	WWW       : http://0x80.pl
	License   : public domain
"""

import pyahocorasick


def exportdot(trie, file):

    def writeln(text=""):
        file.write(text + "\n")

    writeln("digraph ahocorasick {")

    def walk(node):
        queue = [node]
        while queue:
            node = queue.pop()
            yield node

            for child in node.children.itervalues():
                if child != node:
                    queue.append(child)

    nodes = list(walk(trie.root))

    # nodes
    for node in nodes:
        if node.output != pyahocorasick.nil:
            writeln("\tnode%d [shape=doublecircle, label=\"\"]" % id(node))
        else:
            writeln("\tnode%d [shape=circle, label=\"\"]" % id(node))

    # trie edges
    for node in nodes:
        for letter, child in node.children.iteritems():
            nodeid = id(node)
            destid = id(child)
            if destid == id(trie.root):
                # do not show self-links of root node created during make_automaton
                continue

            if letter.isalnum():
                label = letter
            else:
                label = '%02x' % ord(letter)

            writeln("\tnode%d -> node%d [label=\"%s\"]" % (nodeid, destid, label))

    # fail links
    for node in nodes:
        nodeid = id(node)
        failid = id(node.fail)

        if failid != pyahocorasick.nil:
            writeln("\tnode%d -> node%d [color=blue]" % (nodeid, failid))

    writeln("}")


if __name__ == '__main__':
    A = pyahocorasick.Trie()

    A.add_word("he", 0)
    A.add_word("her", 1)
    A.add_word("hers", 2)
    A.add_word("she", 3)
    A.add_word("cat", 4)
    A.add_word("shield", 5)

    with open('trie.dot', 'wt') as f:
        exportdot(A, f)

    A.make_automaton()

    with open('ahocorasick.dot', 'wt') as f:
        exportdot(A, f)
