from __future__ import print_function

#
# Compute prefix-suffix maps for a file
#
import string
non_word = string.punctuation + string.whitespace

def process_file(filename):
    one_gram = {}
    two_gram = {}
    prefix = ( None, None )
    fp = open(filename)
    for line in fp:
        prefix = process_line(line, prefix, one_gram, two_gram)
    return one_gram, two_gram

def process_line(line, prefix, m1, m2):
    from markov3_io import shift
    line = line.replace('-', ' ')
    for word in line.split():
        word = word.strip(non_word).lower()
        if prefix[1] is not None:
            m1.setdefault(prefix[1], []).append(word)
        if prefix[0] is not None and prefix[1] is not None:
            m2.setdefault(prefix, []).append(word)
        prefix = shift(prefix, word)
    return prefix

#
# Compute uni- and bigrams for all .txt files
#
def walk(dir):
    import os
    import os.path
    from markov3_io import write_grams
    for name in os.listdir(dir):
        path = os.path.join(dir, name)
        if os.path.isdir(path):
            walk(path)
        else:
            #print(path)
            root, ext = os.path.splitext(name)
            if ext == ".txt":
                m1, m2 = process_file(path)
                write_grams(path, m1, m2)
                print(path, "prepped")

walk(".")