# IN1140 - gruppe 1 & 2 - annikaol
# Eksempel: bigrams, trigrams

import nltk
nltk.download('inaugural')
from nltk.corpus import inaugural
from collections import Counter, defaultdict
import numpy as np
from nltk import bigrams, trigrams


#print(inaugural.fileids())
print()

# Henter ut hele "2017-Trump.txt" som en lang streng
trump_raw = inaugural.raw("2017-Trump.txt")
print(trump_raw)
# Henter ut en liste av alle ord fra "2017-Trump.txt"
trump_words = inaugural.words("2017-Trump.txt")
# Henter ut en liste av alle setninger fra "2017-Trump.txt"
trump_sents = inaugural.sents("2017-Trump.txt")

# Finne antall ord
total_words = len(trump_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 trump_words:
    distinct_words.append(w.lower())
total_distinct_words = len(set(distinct_words))
print("Nr. of types", total_distinct_words)
print()

# Oppretter en Counter som gir oss tellinger for hvert ord.
fd_trump_words = Counter(trump_words)

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

# 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_trump_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()))
print()

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

# 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)
print()

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

# Bigrammer = sekvenser p? 2 ord; trigramer = sekvenser p? 3 ord
# Printer bigrammene til f?rste setning
print("Bigram gir oss et generator-objekt: ", bigrams(first_sentence))
print()
print("Bigram for f?rste setning: ", list(bigrams(first_sentence)))
print("Trigram for f?rste setning: ", list(trigrams(first_sentence)))
print()

# Husk ? ha med mark?rer for start og slutt p? setningene! Det gj?res med pad_left og pad_right, slik:
print("Bigram med <s> og </s>: ", list(bigrams(first_sentence, pad_left=True, pad_right=True)))
print("Trigram med <s> og </s>: ", 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, med alle sannsynlighetene

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

# Eksempler
print()
print("trigrams_counts[('God', 'bless')]['America'] =", trigram_counts[("God", "bless")]["America"])
print("trigrams_counts[('My', 'fellow')]['nonexisting'] =",trigram_counts[("My", "fellow")]["nonexisting"])
print()

# 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]:
        if total_trigramcounts:
            # Her oppdaterer vi verdiene i modellen
            trigram_model[w1_w2][w3] = trigram_counts[w1_w2][w3]/total_trigramcounts

# Eksempler
print("trigram_model[('God', 'bless')]['America'] =", trigram_model[("God", "bless")]["America"])
print('trigram_model[("My", "fellow")]["nonexisting"] =', 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

print()
# Slik ser den genererte teksten ut (liste) f?r vi har gjort den om til en streng.
print("Text as list: \n", text)
print()

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