#!/usr/local/bin/python3

# Matteo Del Vecchio

from utils import preprocessing
from unigram import unigram
from bigram import bigram
from trigram import trigram

from pprint import pprint

def main():
	# Reading corpus
	wordsList = []
	corpus = open("corpus.txt", "r")
	i = 0
	for line in corpus.readlines():
		# Every line in the corpus will be preprocessed
		wordsList += preprocessing(line)
		if i >= 3: # Added to avoid overloading when loading the dataset
			break
		i+=1

	##### Exercise 1 (b) #####

	# Calculating occurrences of each word in the corpus
	occurrenceCounts = list(map(lambda word: wordsList.count(word), wordsList))

	# Calculating UNIGRAM frequencies and probabilities
	singleWordOcc = dict()	# For each word, it takes track of number of occurrences
	# Initialization of unigram data structure containing probabilities
	# Key: string (the word), value: probability value
	unigramData = dict()
	i = 0
	for word in wordsList:
		# Obtaining word frequency
		singleWordOcc[word] = occurrenceCounts[i]
		# Computing P(wi)
		unigramData[word] = singleWordOcc[word] / len(occurrenceCounts)
		i+=1

	# Calculating BIGRAM frequencies and probabilities
	doubleWordsOcc = dict()	# For pair of contiguous words, it takes track of number of occurrences
	# Initialization of bigram data structure containing probabilities
	# Key: tuple (wi-1, wi), value: probability value
	bigramData = dict()
	for i in range(1, len(wordsList)):
		# Counting frequencies for each bigram
		key = (wordsList[i-1], wordsList[i])
		doubleWordsOcc[key] = doubleWordsOcc.get(key, 0) + 1

	# Calculating probabilities for each bigram
	for key in doubleWordsOcc:
		precWord = key[0]
		bigramData[key] = doubleWordsOcc[key] / singleWordOcc[precWord]		# f(wi-1, wi) / f(wi-1)

	# Calculating TRIGRAM frequencies and probabilities
	triWordsOcc = dict()	# For triplets of contiguous words, it takes track of number of occurrences
	# Initialization of trigram data structure containing probabilities
	# Key: tuple (wi-2, wi-1, wi), value: probability value
	trigramData = dict()
	for i in range(2, len(wordsList)):
		# Counting frequencies for each trigram
		key = (wordsList[i-2], wordsList[i-1], wordsList[i])
		triWordsOcc[key] = triWordsOcc.get(key, 0) + 1

	# Calculating probabilities for each trigram
	for key in triWordsOcc:
		firstPrecWord = key[0]		# word i-2
		secondPrecWord = key[1]		# word i-1
		bigramKey = (firstPrecWord, secondPrecWord)
		trigramData[key] = triWordsOcc[key] / doubleWordsOcc[bigramKey]		# f(wi-2, wi-1, wi) / f(wi-2, wi-1)
	
	##### END Exercise 1 (b) #####

	pprint(unigramData)
	print()
	pprint(bigramData)
	print()
	pprint(trigramData)

	print("Generated by UNIGRAM:\n" + " ".join(unigram(unigramData)) + "\n")
	print("Generated by BIGRAM:\n" + " ".join(bigram(bigramData)) + "\n")
	print("Generated by TRIGRAM:\n" + " ".join(trigram(trigramData, bigramData)) + "\n")

if __name__ == "__main__":
	main()
