from sets import Set
import itertools
import mml

def flatten(l):
	return list(itertools.chain(*l))

def rle_zeroes(l):
	new_l = []
	i = 0
	while i < len(l):
		if l[i] != 0:
			new_l.append(i)
			new_l.append(l[i])
		i += 1	
	if new_l == []:
		new_l = [0, 0]
	return new_l

def preprocess(text):
	words = text.split()
	words_set = Set()
	for word in words:
		words_set.add(word)
	
	word_count = len(words_set)

	word_to_index = {}
	for i, word in enumerate(words_set):
		word_to_index[word] = i

	freq = [[0]*word_count for i in range(word_count)]
	for i in range(len(words)-1):
		w, nw = words[i], words[i+1]
		i, ni = word_to_index[w], word_to_index[nw]
		freq[i][ni] += 1

	compr_freq = map(rle_zeroes, freq)	

	mml_output = mml.Node('markov_statistics', '_')
	mml_words = mml.Node('words', str(word_count))
	for word in word_to_index.keys():
		mml_words.children.append(mml.Node(word, str(word_to_index[word])))
	mml_freqs = mml.Node('frequencies', '_')
	for freq in compr_freq:
		mml_freqs.children.append(mml.Node('f', ' '.join(map(str, freq))))
	mml_output.children.append(mml_words)
	mml_output.children.append(mml_freqs)

	outfile = open('markov.mml', 'w')
	outfile.write(mml.serialize2(mml_output))
	outfile.close()

alice = open('alice.txt', 'r')
alice_str = alice.read()
alice.close()

preprocess(alice_str)

