# IN1140 - gruppe 1 & 2 - annikaol
# 1: NLTK og spr?kmodeller

import nltk
from nltk.corpus import inaugural
from collections import Counter, defaultdict
import numpy as np
from nltk import bigrams, trigrams

#print(inaugural.fileids())

# Henter ut hele "2009-Obama.txt" som en lang streng
obama_raw = inaugural.raw("2009-Obama.txt")
# Henter ut en liste av alle ord fra "2009-Obama.txt"
obama_words = inaugural.words("2009-Obama.txt")
# Henter ut en liste av alle setninger fra "2009-Obama.txt"
obama_sents = inaugural.sents("2009-Obama.txt")

# Finne antall ord
total_words = len(obama_words)
print("Nr. of words", total_words)

# Finne antall unike ord (f.eks. "Eple" og "eple" telles bare 1 gang)
distinct_words = []
for w in obama_words:
    distinct_words.append(w.lower())
total_distinct_words = len(set(distinct_words))
print("Nr. of types", total_distinct_words)


# Oppretter en Counter som gir oss tellinger for hvert ord.
fd_obama_words = Counter(obama_words)

# Vi kan telle forekomster av ord med Counteren slik:
print("10 mest vanlige ord: ", fd_obama_words.most_common(10))
print("Forekomster av 'peace': ", fd_obama_words["peace"])
print("Forekomster av 'the': ", fd_obama_words["promise"])

# Lager en unigram-modell der vi ser p? frekvensen til hvert enkelt ord
# Legger til sannsynlighetene for hvert ord i en ordbok, slik at:
#      - N?kkelen: ordet
#      - Verdien som h?rer til en n?kkel: sannsynligheten til ordet
probabilities = {}
for word, count in fd_obama_words.items():
    probabilities[word] = count/total_words  # Beregn sannsynligheten til hvert ord.

#Summen av alle sannsynligheter er 1 - sjekk dette:
print("Summen av alle sannsynligheter skal bli ca. 1: ",sum(probabilities.values()))

# Genererer ny tekst basert p? unigram-modellen
text = np.random.choice(list(probabilities.keys()), 20,
                        p=list(probabilities.values()))
print(" ".join(text))

# Beregner sannsynligheten for den genererte teksten - versjon 1
word_probabilities = []
for w in text:
    word_probabilities.append(probabilities[w])
#                                                   multipliser alle sannsynlighetene sammen
print("Sannsynligheten for den genererte teksten, v1:", np.prod(word_probabilities))

# Beregner sannsynligheten for den genererte teksten - versjon 2
w_prob = 1
for w in text:
    w_prob *= probabilities[w]
print("Sannsynligheten for den genererte teksten, v2:", w_prob)

# Lage trigram-modell
# F? setningene fra teksten - vi trenger konteksten for ? beregne sekvenser
obama_sents = inaugural.sents("2009-Obama.txt")
first_sentence = obama_sents[0]
print(first_sentence)

# Bigrammer = sekvenser p? 2 ord; trigramer = sekvenser p? 3 ord
# Printer bigrammene til f?rste setning i "2009-Obama.txt"
print(bigrams(first_sentence))
print(list(bigrams(first_sentence)))
print(list(trigrams(first_sentence)))

# Husk ? ha med mark?rer for start og slutt p? setningene! Det gj?res med pad_left og pad_right, slik:
print(list(trigrams(first_sentence, pad_left=True, pad_right=True)))

# Initialiserer frekvens dictionary og sannsynlighet dictionary
# OBS: en defaultdict er en ordbok der vi gir en verdi (her 0), som er den verdien
# hver nye n?kkel har innledningsvis(det vil si at hvis vi ikke har gitt en n?kkel
# en verdi enn?, s? har den her verdien 0 eller 0.0).
trigram_counts = defaultdict(lambda: defaultdict(lambda: 0))  # Denne teller trigrammene
trigram_model = defaultdict(lambda: defaultdict(lambda: 0.0))  # Dette er den faktiske modellen v?r

# Itererer gjennom trigrammene ved ? se p? hvert element i trigrammet samtidig
for sentence in obama_sents:
    for w1, w2, w3 in trigrams(sentence, pad_right= True, pad_left = True):
        trigram_counts[(w1, w2)][w3] += 1

# Eksempler
print(trigram_counts[("My", "fellow")]["citizens"])
print(trigram_counts[("My", "fellow")]["nonexisting"])

# Bruker sannsynlighetsregler for ? beregne sannsynligheten p? hver trigramsekvens
# w1 er f?rste ord i trigrammet, w2 er andre, og w3 er tredje,
# w1_w2 er dermed sekvensen av w1 og w2, disse behandler vi her som en enhet.
for w1_w2 in trigram_counts:
    # Her finner vi totalt antall forekomster for bigrammet w1_w2. Denne bruker vi nedenfor
    # n?r vi skal inne sannsynligheten for trigrammet w1_w2_w3.
    total_trigramcounts = sum(trigram_counts[w1_w2].values())
    for w3 in trigram_counts[w1_w2]:
        # Her oppdaterer vi verdiene i modellen
        trigram_model[w1_w2][w3] = trigram_counts[w1_w2][w3]/total_trigramcounts

# Eksempler
print(trigram_model[("My", "fellow")]["citizens"])
print(trigram_model[("My", "fellow")]["nonexisting"])

# Generer tekst- initialiserer teksten med [None, None], som fungerer som startmark?rer.
# Vi bruker 2 None fordi vi jobber med trigrammer(kun 1 n?r vi jobber med bigrammer).
text = [None, None]
sentence_is_finished = False

while not sentence_is_finished:
    # Henter neste key: sekvens p? de f?rste 2 ordene i trigrammet
    key = tuple(text[-2:])
    # Lager en liste av mulige ord
    words = list(trigram_model[key].keys())
    # Henter sannsynlighetene til trigrammet
    probs = list(trigram_model[key].values())
    # Legger til ett og ett ord
    text.append(np.random.choice(words, p = probs))

    # Vi avslutter n?r de siste 2 ordene som ble generert er None (fungerer som sluttmark?rer)
    if text[-2:] == [None, None]:
        sentence_is_finished = True

# Slik ser den genererte teksten ut (liste) f?r vi har gjort den om til en streng.
print(text)

# Lager en streng fra listen og tar vekk None elementene, slik at det ser ut som en
# faktisk tekst. :)
print(' '.join([t for t in text if t]))