
import string, random
import text
import gnoetics

class Markov:

    def __init__(self):
        self.__next  = {}
        self.__prev  = {}
        self.__start = gnoetics.token_lookup("<start>")
        self.__stop  = gnoetics.token_lookup("<stop>")
        self.__starters = {}
        self.__stoppers = {}

    def add_text(self, txt):
        prev = txt.get_token(0)
        for i in xrange(1, txt.length()):
            curr = txt.get_token(i)
            self.__next.setdefault(prev, []).append(curr)
            self.__prev.setdefault(curr, []).append(prev)
            if prev == self.__start:
                self.__starters[curr] = 1
            if curr == self.__stop:
                self.__stoppers[prev] = 1
            prev = curr

    def get_start(self):
        return self.__start
                
    def get_stop(self):
        return self.__stop

    def __pick(self, choices,
              min_syllables=0,
              max_syllables=0,
              force_starter=0,
              force_stopper=0):

        screened_choices = []
        for c in choices:
            if (min_syllables <= c.syllables()) \
               and (max_syllables == 0 or c.syllables() <= max_syllables) \
               and (force_starter == 0 or c.syllables() < max_syllables or  self.__starters.has_key(c)) \
               and (force_stopper == 0 or c.syllables() < max_syllables or self.__stoppers.has_key(c)):
                screened_choices.append(c)

        if not screened_choices:
            return None
        c = screened_choices[random.randint(0, len(screened_choices)-1)]
        return c
                

        

    def choose_next(self, w, **filters):
        L = self.__next[w]
        return self.__pick(L, **filters)

    def choose_prev(self, w, **filters):
        L = self.__prev[w]
        return self.__pick(L, **filters)

    def generate(self, basis=None,
                 min_syllables=-1,
                 max_syllables=-1):

        walk = []

        done = 0
        while not done:
            if not walk:
                walk = [basis or self.get_start()]
                syllables = 0

            done = 0
            if walk[0] != self.__start:
                w = self.choose_prev(walk[0])
                i = 0
            elif walk[-1] != self.__stop:
                w = self.choose_next(walk[-1])
                i = len(walk)
            else:
                done = 1

            walk.insert(i, w)
            s = w.syllables()
            syllables += s

            # Invalidation
            if (done and min_syllables > 0 and syllables < min_syllables) \
               or (max_syllables > 0 and syllables > max_syllables):
                walk = []

            if not walk:
                done = 0

        walk_str = ""
        for w in walk:
            walk_str = w.to_string(prefix=walk_str)
        return walk_str

        
        
            
