__author__ = 'Justin'

import nltk
import numpy
import random
import string
import re
import platform
from collections import defaultdict

if platform.system() == 'Darwin':
    win = False
else:
    win = True

order = 2

def weighted_pick(freqs):
    """
    Given a list of words and frequencies (weights), return a word randomly based on weight.
    Credit to user97370 at stackoverflow: http://stackoverflow.com/a/2570802

    :param freqs: a dict with structure {str word: int frequency, ...} generated by count_freqs()
    :return:
    """
    total = sum(freqs.itervalues())
    pick = random.randint(0, total - 1)
    tmp = 0
    for key, weight in freqs.iteritems():
        tmp += weight
        if pick < tmp:
            return key


def count_freqs(tokens):
    """
    Iterates through a list of strings with a phrase window of size (order),
    Appending to outcomes such that outcomes[phrase] = {str next_word: int frequency}

    :param tokens: A list of strings tokenized with NLTK
    :return: outcomes[phrase] = {str next_word: int frequency}
    """
    outcomes = dict()
    for i in xrange(0, len(tokens) - order - 1):
        phrase = tuple(tokens[i:i + order])
        if phrase not in outcomes:
            outcomes[phrase] = defaultdict(int)
        try:
            outcomes[phrase][tokens[i + order]] += 1
        except IndexError:
            print "out of bounds"

    return outcomes


def create_text(outcomes, text=None):
    """
    Given a starting text of length (order), or picking randomly from all (order)-length phrases,
    find the next word based on the relative probabilities assigned in count_freqs().  Advance the
    pos window by 1, and repeat the process with the new (order)-length phrase until a '.' or '?'.

    :param outcomes: dict of dicts from count_freqs()
    :param text: list of strings forming the beginning of the text
    :return: text: list of strings forming the '?' or '.'-terminated text
    """

    #start routine

    if text == None:
        text = outcomes.keys()[random.randint(0,len(outcomes))]

    pos = 0

    while True:
        seed = text[pos:pos+order]

        text = text + (weighted_pick(outcomes[seed]),)
        pos += 1
        #print str(pos)+' '+text[pos]+' '+str(text[pos] == '?')
        if text[pos] == '?' or text[pos] == '.':
            return text[0:pos+1]


txt = open(os.path.join('.', 'inputs', 'chs', '17 - ithaca.txt', 'rU'))

txt = txt.read().translate(None, "`()'")

q = re.findall("(?<=[\n]).*\?", txt)
a = re.findall("(?<=[\?]\n\n)[^\r\n]*", txt)

q_no_punct = [re.sub("\?", "", x) for x in q]
a_no_punct = [re.sub("\.", "", x) for x in a]

q_first = defaultdict(int)
a_first = defaultdict(int)

for x in [w.partition(' ')[0] for w in q_no_punct]: q_first[x] += 1
for x in [w.partition(' ')[0] for w in a_no_punct]: a_first[x] += 1

q_tokens = nltk.word_tokenize(' '.join(q))
a_tokens = nltk.word_tokenize(' '.join(a))

q_freq = count_freqs(q_tokens)
a_freq = count_freqs(a_tokens)

questions = []
answers = []

for i in xrange(1,3):
    print ' '.join(create_text(q_freq))
    print ' '.join(create_text(a_freq))
    print ''

