import os

# Increment counts for ngram dicitonaries
def addToDics((first, second), dicBi, dicMono):
	def addToDic(entry, dic):
		if entry in dic:
			dic[entry] = dic[entry] + 1
		else:
			dic[entry] = 1

	addToDic(first, dicMono)
	addToDic((first, second), dicBi)


# Adds words in story to ngram dictionary
def fromStory(dicMono, dicBi, story):
	for token in ["...", "?", "!", "\n", "\r"]:
		story.replace(token, ".")

	for token in [",", "(", ")", "'", '"']:
		story.replace(token, "")

	lines = story.split(".")

	for line in lines:
		bi = ('<s>', '<s>')
		for word in filter(lambda x : x.isalpha(), line.lower().split(" ")):
			bi = (bi[1], word)
			addToDics(bi, dicBi, dicMono)

# Generete ngrams from a collection of storeis.
def getDictionaries():
	files = os.listdir('input/texts/')

	dicBi = dict()
	dicMono = dict()

	for fileName in files:
		f = open('input/texts/' + fileName, 'r')
		story = "".join( [ line for line in f.readlines() ] )
		f.close()
		fromStory(dicMono, dicBi, story)

	return (dicMono, dicBi)

(monoGramDic, biGramDic) = getDictionaries()

# This function ranks a set of words based on their swipe path distances similarity and the language model.
def rank(previous, sett, wlm=True):
	# Has a version with and without language model.
	if wlm:
		sett = map(lambda x : (x[0], x[1]**2.1), sett)

		p = previous[len(previous)-1]
		bigramChances = []
		monoFrequencies = []
		if p in monoGramDic:
			monoFreq = float(monoGramDic[p]) + 1.0
		else:
			monoFreq = 1.0

		for i in xrange(len(sett)):
			if (p, sett[i][0]) in biGramDic:
				biFreq = biGramDic[(p, sett[i][0])] + 1.0
				# print (p, sett[i][0]), "occured:", biFreq
			else:
				biFreq = 1.0
				# print (p, sett[i][0]), "didn't occur"

			bigramChances.append( (biFreq / monoFreq) )

		bigramSum = sum(bigramChances)
		bigramChances = [ f / bigramSum for f in bigramChances ]

		swipeSum = float(sum([ pair[1] for pair in sett]))
		swipeChances = [ 1 - (f/swipeSum) for f in [ e[1] for e in sett] ]

		# We now have the (normalized) probabilities based on swipe similarity and language model.
		# Combine them to get a final ranking on the words.
		finalChances = [ e[0]*e[1] for e in zip(swipeChances, bigramChances) ]

	else:
		swipeSum = float(sum([ pair[1] for pair in sett]))
		swipeChances = [ 1 - (f/swipeSum) for f in [ e[1] for e in sett] ]
		
		finalChances = swipeChances

	newSett = zip([e[0] for e in sett], finalChances)
	newSett.sort( key = lambda tup : tup[1], reverse=True )
	
	return newSett

