# -*- coding: utf-8 -*-
"""
Code to accompany the chapter "Natural Language Corpus Data"
from the book "Beautiful Data" (Segaran and Hammerbacher, 2009)
http://oreilly.com/catalog/9780596157111/

Code copyright (c) 2008-2009 by Peter Norvig

You are free to use this code under the MIT licencse: 
http://www.opensource.org/licenses/mit-license.php
"""
import codecs
import re, string, random, glob, operator, heapq, sys
from collections import defaultdict
from math import log10

sys.setrecursionlimit(100000000)

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

@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)

# increased max length from 20 to 25
def splits(text, L=20):
    "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 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 count,key 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=' '):
    "Read key,value pairs from file."
    for line in codecs.open(name,'r','utf-8').read().split('\n')[:-1]:
    #for line in file(name):
        space_pos = line.find(' ')
        yield [line[:space_pos],line[space_pos+1:]]

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

#### segment2: second version, with bigram counts, (p. 226-227)

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

@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, (Prem, rem)): 
    "Combine first and rem results into one (probability, words) pair." 
    return Pfirst+Prem, [first]+rem 

text = codecs.open('hi-wordseg.txt','r','utf-8').read()
# append end of sentence marker to end of each
#text = ''.join([line+'\n' for line in text.split('\n')])
text = ''.join([line+u' </s>\n' for line in text.split('\n')])
N = 15900757 ## Number of tokens

fd = codecs.open('output','w','utf-8')
#fd = open('output','w')
Pw  = Pdist(datafile('common_unigrams.txt'), N, avoid_long_words)
#for word in segment(text):
#    if '/' in word:
#        #fd.write(u'\n')
#        continue
#    if 's' in word: continue
#    fd.write(word + u' ')
#

P2w = Pdist(datafile('common_bigrams.txt'), N)
for word in segment2(text)[1]:
    word = re.sub('<s>','',word)
    word = re.sub('</s>','',word)
    word = re.sub('</s >','',word)
 #   if '/s' in word:
 #       fd.write(u'\n')
 #       continue
 #   if 's' in word: continue
    fd.write(word + u' ')
fd.close()