"""
    Aho-Corasick string search algorithm.

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

import ahocorasick
import os

from ahocorasick import EMPTY, TRIE


def dump2dot(automaton, file):

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

    def nodename(nodeid):
        return 'node%x' % (nodeid & 0xffffffff)

    if automaton.kind == EMPTY:
        writeln("digraph empty {}")
        return

    if automaton.kind == TRIE:
        name = "trie"
    else:
        name = "ahocorasick"


    writeln("digraph %s {" % name)

    nodes, edges, fail = automaton.dump()

    # nodes
    for nodeid, end in nodes:
        if end:
            attr = '[shape=doublecircle, label=""]'
        else:
            attr = '[shape=circle, label=""]'

        writeln("\t%s %s" % (nodename(nodeid), attr))

    def format_label(label):
        label = str(label, 'ascii')
        label = label.replace('"', r'\"')
        return '"%s"' % label

    # trie edges
    for nodeid, label, destid in edges:
        writeln("\t%s -> %s [label=%s]" % (nodename(nodeid), nodename(destid), format_label(label)))

    # fail links
    for nodeid, failid in fail:
        writeln("\t%s -> %s [color=blue]" % (nodename(nodeid), nodename(failid)))

    writeln("}")


def show(automaton):
    path = '/dev/shm/%s.dot' % os.getpid()
    with open(path, 'wt') as f:
        dump2dot(automaton, f)

    os.system("xdot %s" % path)
    #os.system("dotty %s" % path)
    os.unlink(path)


if __name__ == '__main__':
    A = ahocorasick.Automaton(ahocorasick.STORE_LENGTH)

    A.add_word("he")
    A.add_word("her")
    A.add_word("hers")
    A.add_word("she")
    A.add_word("cat")
    A.add_word("shield")

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

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