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, h1, h2):
    line = line.replace('-', ' ')
    for word in line.split():
        word = word.strip(non_word).lower()
        if prefix[1] is not None:
            h1.setdefault(prefix[1], []).append(word)
        if prefix[0] is not None and prefix[1] is not None:
            h2.setdefault(prefix, []).append(word)
        prefix = shift(prefix, word)
    return prefix

def shift(prefix, word):
    return prefix[1:] + (word,)

#
# Compute uni- and bigrams for all .txt files
#
def walk(dir):
    import os
    import os.path
    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":
                prep_file(path)

def prep_file(path):
    h1, h2 = process_file(path)
    import os.path
    root, ext = os.path.splitext(path)
    tsv_name = root + ".tsv"
    with open(tsv_name, "w") as f:
        for prefix, suffix in h1.items():
            print("%s\t%s" % (prefix, '\t'.join(suffix)), file=f)
        print(file=f)
        for prefix, suffix in h2.items():
            print("%s\t%s\t%s" % (prefix[0], prefix[1], '\t'.join(suffix)),
                  file=f)
    print(path, "prepped")

walk(".")