import re, string, random, glob, operator, heapq
from collections import defaultdict
from math import log10
import functools

def memo(f):
    "Memoize function f."
    table = {}
    def fmemo(*args):
        if args not in table:
            table[args] = f(*args)
        return table[args]
    fmemo.memo = table
    return fmemo

def test(verbose=None):
    """Run some tests, taken from the chapter.
    Since the hillclimbing algorithm is randomized, some tests may fail."""
    import doctest
    #print 'Running tests...'
    doctest.testfile('ngrams-test.txt', verbose=verbose)

@memo
def segment(text):
    "Return a list of words that is the best segmentation of text."
    if not text: return []
    candidates = ([first]+segment(rem) for first,rem in splits(text))
    return max(candidates, key=Pwords)

def splits(text, L=11):
    "Return a list of all possible (first, rem) pairs, len(first)<=L."
    return [(text[:i+1], text[i+1:]) 
            for i in range(min(len(text), L))]

def Pwords(words): 
    "The Naive Bayes probability of a sequence of words."
    return product(Pw(w) for w in words)

def product(nums):
    "Return the product of a sequence of numbers."
    return functools.reduce(operator.mul, nums, 1)

class Pdist(dict):
    "A probability distribution estimated from counts in datafile."
    def __init__(self, data=[], N=None, missingfn=None):
        for key,count in data:
            self[key] = self.get(key, 0) + int(count)
        self.N = float(N or sum(self.itervalues()))
        self.missingfn = missingfn or (lambda k, N: 1./N)
    def __call__(self, key): 
        if key in self: return self[key]/self.N  
        else: return self.missingfn(key, self.N)

def datafile(name, sep='\t'):
    "Read key,value pairs from file."
    f = open(name, encoding='utf-8')
    for line in f:
        yield line.split(sep)

def avoid_long_words(key, N):
    "Estimate the probability of an unknown word."
    return 10./(N * 10**len(key))

#N = 1024908267229 ## Number of tokens
#N = 8905000 ## Number of tokens
N = 13412600 ## Number of tokens
# N = 13517000
#Pw  = Pdist(datafile('count_1w.txt'), N, avoid_long_words)
Pw  = Pdist(datafile('../unigrams/vocab.txt'), N, avoid_long_words)

def cPw(word, prev):
    "Conditional probability of word, given previous word."
    try:
        return P2w[prev + ' ' + word]/float(Pw[prev])
    except KeyError:
        return Pw(word)

# P2w = Pdist(datafile('count_2w.txt'), N)
P2w = Pdist(datafile('vocab_big.txt'), N)

@memo 
def segment2(text, prev='<S>'): 
    "Return (log P(words), words), where words is the best segmentation." 
    if not text: return 0.0, [] 
    candidates = [combine(log10(cPw(first, prev)), first, segment2(rem, first)) 
                  for first,rem in splits(text)] 
    return max(candidates) 

def combine(Pfirst, first, l): 
    Prem = l[0]
    rem = l[1]
    "Combine first and rem results into one (probability, words) pair." 
    return Pfirst+Prem, [first]+rem 
