import re
import math
import sys
import codecs
import random
from collections import OrderedDict
import queue

PREFIX = '<s> '
SUFFIX = ' </s>'
UNKNOWN = '<unk>'

"""
n - size of n-grams
text - list of words/strings
returns: tuple of n-1 (word, context)
"""


def get_ngrams(n, text):  # - generator
    prefix = PREFIX
    suffix = SUFFIX
    for i in range(n - 2):
        prefix += PREFIX

    sentences = re.split(r'[?,.!:;]+', text)
    for sentence in sentences:
        sentence = re.sub(r'[^\w<>]', ' ', sentence)
        sentence = sentence.strip()
        sentence = prefix + sentence + suffix
        words = sentence.split()

        for i in range(n - 1, len(words)):
            context = ''
            for j in range(i - n + 1, i):
                context = context + ' ' + words[j].strip()
            yield words[i].strip(), context.strip()


"""
n - number of grams
corpus path - file path
returns: NGramLM trained model
"""


def create_ngramlm(n, corpus_path, delta=.0):
    ngramlm = NGramLM(n, delta)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        # corpus = mask_rare(corpus)
        ngramlm.update(corpus)

    return ngramlm


"""
n - number of grams
corpus path - file path
returns: NGramLM trained model
"""


def create_interpolator(n, lambdas, corpus_path):
    inter = NGramInterpolator(n, lambdas)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        corpus = mask_rare(corpus)
        inter.update(corpus)

    return inter


"""
returns log(e) prob of text (list of words/strings)
under NGramLM trained model
"""


def text_prob(model, text, delta=.0):
    log_prob = 0
    ngrams = get_ngrams(model.n, text)
    for word, context in ngrams:
        prob = model.word_prob(word, context, delta)
        log_prob += math.log(prob)

    return log_prob


"""
concatenate word with context
"""


def get_word_with_context(word, context):
    return context.strip() + ' ' + word.strip()


"""
corpus - content of file
returns: a copy of corpus with words 
that appear only once replaced by the `<unk>' token
"""


def mask_rare(corpus):
    all_words = {}
    rare_words = set()
    words = re.split(r'[^\w]', corpus)
    for word in words:
        count = all_words.get(word, 0)
        count += 1
        all_words[word] = count

    for word, count in all_words.items():
        if count == 1:
            rare_words.add(word)
    print("all: ", len(all_words), ", rare: ", len(rare_words))

    masked_corpus = []
    sentences = corpus.split()
    for sentence in sentences:
        words = re.split(r'[^\w]', sentence)
        for word in words:
            if word in rare_words:
                sentence = sentence.replace(word, UNKNOWN)
        masked_corpus.append(sentence)
        masked_corpus.append(' ')

    return ''.join(masked_corpus).strip()


"""
model - trained model
corpus_path - test data file path, no need to mask rare words
returns perplexity of a trained model
counts N, total tokens
"""


def perplexity(model, corpus_path):
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        num_tokens = 0

        sentences = re.split(r'[?,.!:;]+', corpus)
        for sentence in sentences:
            words = re.split(r'[^\w]', sentence)
            num_tokens += len(words)

        i = text_prob(model, corpus, model.delta) / num_tokens
        return 2 ** -i


def random_text(model, max_length, delta=0):
    context = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        context.append(PREFIX.strip())
    counter = 0
    sentence = ''

    while True:
        word = model.random_word(' '.join(context).strip(), delta)
        if word == suffix:
            return sentence.strip()

        sentence = sentence + ' ' + word
        counter += 1
        if counter > max_length:
            return sentence.strip()

        for i in range(len(context) - 1):
            context[i] = context[i + 1]
        context[len(context) - 1] = word


"""
almost identical to random_text
"""


def likeliest_text(model, max_length, delta=0):
    context = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        context.append(PREFIX.strip())
    counter = 0
    sentence = ''

    while True:
        word = model.likeliest_word(' '.join(context).strip(), delta)
        if word == suffix:
            return sentence.strip()

        sentence = sentence + ' ' + word
        counter += 1
        if counter > max_length:
            return sentence.strip()

        for i in range(len(context) - 1):
            context[i] = context[i + 1]
        context[len(context) - 1] = word


def beam_search_text(model, max_length, k, delta=.0):
    # here we first build the first candidate '<s> <s>'
    starter = []
    best_k = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        starter.append(PREFIX.strip())
    init_candidate = Candidate(' '.join(starter).strip(), 1)  # check the 'Candidate' class below

    # put the first candidate into queue and make it as the current best_k
    frontier = queue.Queue()
    best_k.append(init_candidate)
    frontier.put(init_candidate)

    # loop until we got the max_length of words
    for i in range(max_length):
        # use sub_k to store the next level best k words followed by the current best_k words
        # sub_k is reset at each level, i.e. for different i
        sub_k = []
        while not frontier.empty():
            # poll a text from the queue
            # if the last word of the text is suffix <\s>, skip it
            candidate = frontier.get()
            words = candidate.text.split()
            if words[len(words) - 1] == suffix:
                continue

            # use the last n words to build context
            # get the next k_words of the context
            # append each of the k_words to current text, which yields k new texts
            # remove the current text from best_k, since we already find its descendants
            context = words[len(words) - model.n + 1: len(words)]
            k_words = model.beam_search_word(' '.join(context).strip(), k, delta)
            sub_k = sub_k + get_new_candidates(model, candidate, k_words, delta)
            best_k.remove(candidate)

        # merge all the elements of best_k and sub_k to elect new best_k
        # put the new best_k texts to queue, to find their descendants
        best_k = merge(best_k, sub_k, k)
        for candidate in best_k:
            frontier.put(candidate)

    return best_k


def get_new_candidates(model, candidate, sub_k, delta=.0):
    # append each of the sub_k to current text,
    # returns candidate's k descendants
    new_candidates = []
    for item in sub_k:
        new_candidate = candidate.clone()
        new_candidate.append(item)
        new_candidate.prob = text_prob(model, new_candidate.text, delta)
        new_candidates.append(new_candidate)

    return new_candidates


def merge(best_k, sub_k, k):
    # elect new best_k
    temp = best_k + sub_k
    temp = sorted(temp, key=lambda x: x.prob, reverse=True)
    return temp[:k]


class NGramLM:
    def __init__(self, n, delta=.0):
        self.n = n  # size of n-grams
        self.ngram_counts = {}  # n-grams seen in the training data,
        self.context_counts = {}  # contexts seen in the training data,
        self.vocabulary = {}  # words seen in the traniing data
        self.delta = delta

    """
    updates NGramLM's internal counts
    and vocabulary for ngrams in text (list of words/strings)
    """

    def update(self, text):
        ngrams = get_ngrams(self.n, text)
        for word, context in ngrams:
            word_with_context = get_word_with_context(word, context)
            ngram_counter = self.ngram_counts.get(word_with_context, 0)
            context_counter = self.context_counts.get(context, 0)
            vocabulary_counter = self.vocabulary.get(word, 0)
            ngram_counter += 1
            context_counter += 1
            vocabulary_counter += 1
            self.ngram_counts[word_with_context] = ngram_counter
            self.context_counts[context] = context_counter
            self.vocabulary[word] = vocabulary_counter

        # print(self.vocabulary[UNKNOWN])

    """
    returns prob of n-gram (word, context) using internal counters.
    If context is previously unseen, returns 1/|V|, V is vocabulary model
    """

    def word_prob_delta_zero(self, word, context):
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        if context_counter == 0 or word not in self.vocabulary:
            return self.vocabulary[UNKNOWN] / (len(self.vocabulary) * 1.0)
        else:
            return ngram_counter / (context_counter * 1.0)


    """
    prob with smoothing
    returns: Laplace-smoothed probabilities 
    (need to modify to fit n-gram)
    """

    def word_prob(self, word, context, delta=.0):
        if delta == 0:
            return self.word_prob_delta_zero(word, context)

        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        return (ngram_counter + delta) / ((context_counter + delta * len(self.vocabulary)) * 1.0)


    """
    returns a word sampled from the model's 
    probability distribution for context
    """

    def random_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        vocabulary = sorted(vocabulary)

        distrib = []
        total_counts = 0
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            total_counts += count
            distrib.append(total_counts)

        r = random.random()
        for i in range(len(distrib)):
            distrib[i] = distrib[i] / (total_counts * 1.0)
            if r < distrib[i]:
                return vocabulary[i]

        return vocabulary[len(vocabulary) - 1]


    """
        returns n-gram with the highest prob for context
    """

    def likeliest_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())

        distrib = []
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            distrib.append(count)

        max_value = 0
        max_Index = 0
        for i in range(len(distrib)):
            if max_value < distrib[i]:
                max_value = distrib[i]
                max_Index = i

        return vocabulary[max_Index]


    """
        return the best k words of the given context
    """

    def beam_search_word(self, context, k, delta=.0):
        distrib = {}
        best_k = []

        # this loop loads every possible word following the given context with its counts into distrib{}
        for word, _ in self.vocabulary.items():
            distrib[word] = self.word_prob(word, context, delta)

        # rank the probabilities in the order of high to low
        ordered_distrib = OrderedDict(sorted(distrib.items(), key=lambda kv: kv[1], reverse=True))

        # this loop loads the top ranked k words into best_k
        i = 0
        for word, prob in ordered_distrib.items():
            if i >= k:
                break
            best_k.append(Candidate(word, prob))
            i += 1

        return best_k


"""
	linear interpolation.
"""

class NGramInterpolator:


    """
        n - size of the largest n-gram
        lambdas - a list of length n containing the interpolation factors
            (floats) in descending order of n-gram size.

        save n and lambdas
        initialize n internal NGramLMs, one for each n-gram size
    """

    def __init__(self, n, lambdas):
        self.n = n
        self.lambdas = lambdas
        self.models = []
        for i in range(n):
            self.models.append(NGramLM(n - i))


    """
        update all of the internal NGramLMs
    """

    def update(self, text):
        for model in self.models:
            model.update(text)


    """
        returns the linearly interpolated prob using lambdas
        and prob given by NGramLMs
    """

    def word_prob(self, word, context, delta=0):
        prob = 0
        for i in range(self.n):
            prob += self.lambdas[i] * self.models[i].word_prob(word, context, delta)

        return prob


class Candidate:

    """
        This class is just for convenience,
        avoid to keep return tuple
    """

    def __init__(self, text, prob):
        self.text = text
        self.prob = prob

    def append(self, word, prob):
        # append a word to text
        # the prob is overwritten later by text_prob,
        # since I original design it to be the real probability.
        # this has the potential of underflow
        self.text = self.text + ' ' + word.strip()
        self.prob *= prob

    def append(self, other):
        # append another candidate to this candidate
        # text is being appended
        # the prob is overwritten later by text_prob,
        # since I original design it to be the real probability.
        # this has the potential of underflow
        self.text = self.text + ' ' + other.text.strip()
        self.prob *= other.prob

    def clone(self):
        # create a copy of current candidate
        return Candidate(self.text, self.prob)


def main(argv):
    trimodel = create_ngramlm(3, "shakespeare.txt")
    for sentence in beam_search_text(trimodel, 10, 5, 1):
        print(sentence.text)

    pass


if __name__ == '__main__':
    main(sys.argv)

149d8f8b1ddbb11763d9e163192be2e25cf77323
<mask token>


def get_ngrams(n, text):
    prefix = PREFIX
    suffix = SUFFIX
    for i in range(n - 2):
        prefix += PREFIX
    sentences = re.split('[?,.!:;]+', text)
    for sentence in sentences:
        sentence = re.sub('[^\\w<>]', ' ', sentence)
        sentence = sentence.strip()
        sentence = prefix + sentence + suffix
        words = sentence.split()
        for i in range(n - 1, len(words)):
            context = ''
            for j in range(i - n + 1, i):
                context = context + ' ' + words[j].strip()
            yield words[i].strip(), context.strip()


<mask token>


def get_new_candidates(model, candidate, sub_k, delta=0.0):
    new_candidates = []
    for item in sub_k:
        new_candidate = candidate.clone()
        new_candidate.append(item)
        new_candidate.prob = text_prob(model, new_candidate.text, delta)
        new_candidates.append(new_candidate)
    return new_candidates


<mask token>


class NGramLM:

    def __init__(self, n, delta=0.0):
        self.n = n
        self.ngram_counts = {}
        self.context_counts = {}
        self.vocabulary = {}
        self.delta = delta
    """
    updates NGramLM's internal counts
    and vocabulary for ngrams in text (list of words/strings)
    """

    def update(self, text):
        ngrams = get_ngrams(self.n, text)
        for word, context in ngrams:
            word_with_context = get_word_with_context(word, context)
            ngram_counter = self.ngram_counts.get(word_with_context, 0)
            context_counter = self.context_counts.get(context, 0)
            vocabulary_counter = self.vocabulary.get(word, 0)
            ngram_counter += 1
            context_counter += 1
            vocabulary_counter += 1
            self.ngram_counts[word_with_context] = ngram_counter
            self.context_counts[context] = context_counter
            self.vocabulary[word] = vocabulary_counter
    """
    returns prob of n-gram (word, context) using internal counters.
    If context is previously unseen, returns 1/|V|, V is vocabulary model
    """

    def word_prob_delta_zero(self, word, context):
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        if context_counter == 0 or word not in self.vocabulary:
            return self.vocabulary[UNKNOWN] / (len(self.vocabulary) * 1.0)
        else:
            return ngram_counter / (context_counter * 1.0)
    """
    prob with smoothing
    returns: Laplace-smoothed probabilities 
    (need to modify to fit n-gram)
    """

    def word_prob(self, word, context, delta=0.0):
        if delta == 0:
            return self.word_prob_delta_zero(word, context)
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        return (ngram_counter + delta) / ((context_counter + delta * len(
            self.vocabulary)) * 1.0)
    """
    returns a word sampled from the model's 
    probability distribution for context
    """

    def random_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        vocabulary = sorted(vocabulary)
        distrib = []
        total_counts = 0
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            total_counts += count
            distrib.append(total_counts)
        r = random.random()
        for i in range(len(distrib)):
            distrib[i] = distrib[i] / (total_counts * 1.0)
            if r < distrib[i]:
                return vocabulary[i]
        return vocabulary[len(vocabulary) - 1]
    """
        returns n-gram with the highest prob for context
    """

    def likeliest_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        distrib = []
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            distrib.append(count)
        max_value = 0
        max_Index = 0
        for i in range(len(distrib)):
            if max_value < distrib[i]:
                max_value = distrib[i]
                max_Index = i
        return vocabulary[max_Index]
    """
        return the best k words of the given context
    """

    def beam_search_word(self, context, k, delta=0.0):
        distrib = {}
        best_k = []
        for word, _ in self.vocabulary.items():
            distrib[word] = self.word_prob(word, context, delta)
        ordered_distrib = OrderedDict(sorted(distrib.items(), key=lambda kv:
            kv[1], reverse=True))
        i = 0
        for word, prob in ordered_distrib.items():
            if i >= k:
                break
            best_k.append(Candidate(word, prob))
            i += 1
        return best_k


<mask token>


class NGramInterpolator:
    """
        n - size of the largest n-gram
        lambdas - a list of length n containing the interpolation factors
            (floats) in descending order of n-gram size.

        save n and lambdas
        initialize n internal NGramLMs, one for each n-gram size
    """

    def __init__(self, n, lambdas):
        self.n = n
        self.lambdas = lambdas
        self.models = []
        for i in range(n):
            self.models.append(NGramLM(n - i))
    """
        update all of the internal NGramLMs
    """

    def update(self, text):
        for model in self.models:
            model.update(text)
    """
        returns the linearly interpolated prob using lambdas
        and prob given by NGramLMs
    """

    def word_prob(self, word, context, delta=0):
        prob = 0
        for i in range(self.n):
            prob += self.lambdas[i] * self.models[i].word_prob(word,
                context, delta)
        return prob


class Candidate:
    """
        This class is just for convenience,
        avoid to keep return tuple
    """

    def __init__(self, text, prob):
        self.text = text
        self.prob = prob

    def append(self, word, prob):
        self.text = self.text + ' ' + word.strip()
        self.prob *= prob

    def append(self, other):
        self.text = self.text + ' ' + other.text.strip()
        self.prob *= other.prob

    def clone(self):
        return Candidate(self.text, self.prob)


<mask token>

<mask token>


def get_ngrams(n, text):
    prefix = PREFIX
    suffix = SUFFIX
    for i in range(n - 2):
        prefix += PREFIX
    sentences = re.split('[?,.!:;]+', text)
    for sentence in sentences:
        sentence = re.sub('[^\\w<>]', ' ', sentence)
        sentence = sentence.strip()
        sentence = prefix + sentence + suffix
        words = sentence.split()
        for i in range(n - 1, len(words)):
            context = ''
            for j in range(i - n + 1, i):
                context = context + ' ' + words[j].strip()
            yield words[i].strip(), context.strip()


<mask token>


def create_ngramlm(n, corpus_path, delta=0.0):
    ngramlm = NGramLM(n, delta)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        ngramlm.update(corpus)
    return ngramlm


<mask token>


def text_prob(model, text, delta=0.0):
    log_prob = 0
    ngrams = get_ngrams(model.n, text)
    for word, context in ngrams:
        prob = model.word_prob(word, context, delta)
        log_prob += math.log(prob)
    return log_prob


<mask token>


def get_new_candidates(model, candidate, sub_k, delta=0.0):
    new_candidates = []
    for item in sub_k:
        new_candidate = candidate.clone()
        new_candidate.append(item)
        new_candidate.prob = text_prob(model, new_candidate.text, delta)
        new_candidates.append(new_candidate)
    return new_candidates


def merge(best_k, sub_k, k):
    temp = best_k + sub_k
    temp = sorted(temp, key=lambda x: x.prob, reverse=True)
    return temp[:k]


class NGramLM:

    def __init__(self, n, delta=0.0):
        self.n = n
        self.ngram_counts = {}
        self.context_counts = {}
        self.vocabulary = {}
        self.delta = delta
    """
    updates NGramLM's internal counts
    and vocabulary for ngrams in text (list of words/strings)
    """

    def update(self, text):
        ngrams = get_ngrams(self.n, text)
        for word, context in ngrams:
            word_with_context = get_word_with_context(word, context)
            ngram_counter = self.ngram_counts.get(word_with_context, 0)
            context_counter = self.context_counts.get(context, 0)
            vocabulary_counter = self.vocabulary.get(word, 0)
            ngram_counter += 1
            context_counter += 1
            vocabulary_counter += 1
            self.ngram_counts[word_with_context] = ngram_counter
            self.context_counts[context] = context_counter
            self.vocabulary[word] = vocabulary_counter
    """
    returns prob of n-gram (word, context) using internal counters.
    If context is previously unseen, returns 1/|V|, V is vocabulary model
    """

    def word_prob_delta_zero(self, word, context):
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        if context_counter == 0 or word not in self.vocabulary:
            return self.vocabulary[UNKNOWN] / (len(self.vocabulary) * 1.0)
        else:
            return ngram_counter / (context_counter * 1.0)
    """
    prob with smoothing
    returns: Laplace-smoothed probabilities 
    (need to modify to fit n-gram)
    """

    def word_prob(self, word, context, delta=0.0):
        if delta == 0:
            return self.word_prob_delta_zero(word, context)
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        return (ngram_counter + delta) / ((context_counter + delta * len(
            self.vocabulary)) * 1.0)
    """
    returns a word sampled from the model's 
    probability distribution for context
    """

    def random_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        vocabulary = sorted(vocabulary)
        distrib = []
        total_counts = 0
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            total_counts += count
            distrib.append(total_counts)
        r = random.random()
        for i in range(len(distrib)):
            distrib[i] = distrib[i] / (total_counts * 1.0)
            if r < distrib[i]:
                return vocabulary[i]
        return vocabulary[len(vocabulary) - 1]
    """
        returns n-gram with the highest prob for context
    """

    def likeliest_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        distrib = []
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            distrib.append(count)
        max_value = 0
        max_Index = 0
        for i in range(len(distrib)):
            if max_value < distrib[i]:
                max_value = distrib[i]
                max_Index = i
        return vocabulary[max_Index]
    """
        return the best k words of the given context
    """

    def beam_search_word(self, context, k, delta=0.0):
        distrib = {}
        best_k = []
        for word, _ in self.vocabulary.items():
            distrib[word] = self.word_prob(word, context, delta)
        ordered_distrib = OrderedDict(sorted(distrib.items(), key=lambda kv:
            kv[1], reverse=True))
        i = 0
        for word, prob in ordered_distrib.items():
            if i >= k:
                break
            best_k.append(Candidate(word, prob))
            i += 1
        return best_k


<mask token>


class NGramInterpolator:
    """
        n - size of the largest n-gram
        lambdas - a list of length n containing the interpolation factors
            (floats) in descending order of n-gram size.

        save n and lambdas
        initialize n internal NGramLMs, one for each n-gram size
    """

    def __init__(self, n, lambdas):
        self.n = n
        self.lambdas = lambdas
        self.models = []
        for i in range(n):
            self.models.append(NGramLM(n - i))
    """
        update all of the internal NGramLMs
    """

    def update(self, text):
        for model in self.models:
            model.update(text)
    """
        returns the linearly interpolated prob using lambdas
        and prob given by NGramLMs
    """

    def word_prob(self, word, context, delta=0):
        prob = 0
        for i in range(self.n):
            prob += self.lambdas[i] * self.models[i].word_prob(word,
                context, delta)
        return prob


class Candidate:
    """
        This class is just for convenience,
        avoid to keep return tuple
    """

    def __init__(self, text, prob):
        self.text = text
        self.prob = prob

    def append(self, word, prob):
        self.text = self.text + ' ' + word.strip()
        self.prob *= prob

    def append(self, other):
        self.text = self.text + ' ' + other.text.strip()
        self.prob *= other.prob

    def clone(self):
        return Candidate(self.text, self.prob)


<mask token>

<mask token>


def get_ngrams(n, text):
    prefix = PREFIX
    suffix = SUFFIX
    for i in range(n - 2):
        prefix += PREFIX
    sentences = re.split('[?,.!:;]+', text)
    for sentence in sentences:
        sentence = re.sub('[^\\w<>]', ' ', sentence)
        sentence = sentence.strip()
        sentence = prefix + sentence + suffix
        words = sentence.split()
        for i in range(n - 1, len(words)):
            context = ''
            for j in range(i - n + 1, i):
                context = context + ' ' + words[j].strip()
            yield words[i].strip(), context.strip()


<mask token>


def create_ngramlm(n, corpus_path, delta=0.0):
    ngramlm = NGramLM(n, delta)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        ngramlm.update(corpus)
    return ngramlm


<mask token>


def create_interpolator(n, lambdas, corpus_path):
    inter = NGramInterpolator(n, lambdas)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        corpus = mask_rare(corpus)
        inter.update(corpus)
    return inter


<mask token>


def text_prob(model, text, delta=0.0):
    log_prob = 0
    ngrams = get_ngrams(model.n, text)
    for word, context in ngrams:
        prob = model.word_prob(word, context, delta)
        log_prob += math.log(prob)
    return log_prob


<mask token>


def perplexity(model, corpus_path):
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        num_tokens = 0
        sentences = re.split('[?,.!:;]+', corpus)
        for sentence in sentences:
            words = re.split('[^\\w]', sentence)
            num_tokens += len(words)
        i = text_prob(model, corpus, model.delta) / num_tokens
        return 2 ** -i


<mask token>


def get_new_candidates(model, candidate, sub_k, delta=0.0):
    new_candidates = []
    for item in sub_k:
        new_candidate = candidate.clone()
        new_candidate.append(item)
        new_candidate.prob = text_prob(model, new_candidate.text, delta)
        new_candidates.append(new_candidate)
    return new_candidates


def merge(best_k, sub_k, k):
    temp = best_k + sub_k
    temp = sorted(temp, key=lambda x: x.prob, reverse=True)
    return temp[:k]


class NGramLM:

    def __init__(self, n, delta=0.0):
        self.n = n
        self.ngram_counts = {}
        self.context_counts = {}
        self.vocabulary = {}
        self.delta = delta
    """
    updates NGramLM's internal counts
    and vocabulary for ngrams in text (list of words/strings)
    """

    def update(self, text):
        ngrams = get_ngrams(self.n, text)
        for word, context in ngrams:
            word_with_context = get_word_with_context(word, context)
            ngram_counter = self.ngram_counts.get(word_with_context, 0)
            context_counter = self.context_counts.get(context, 0)
            vocabulary_counter = self.vocabulary.get(word, 0)
            ngram_counter += 1
            context_counter += 1
            vocabulary_counter += 1
            self.ngram_counts[word_with_context] = ngram_counter
            self.context_counts[context] = context_counter
            self.vocabulary[word] = vocabulary_counter
    """
    returns prob of n-gram (word, context) using internal counters.
    If context is previously unseen, returns 1/|V|, V is vocabulary model
    """

    def word_prob_delta_zero(self, word, context):
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        if context_counter == 0 or word not in self.vocabulary:
            return self.vocabulary[UNKNOWN] / (len(self.vocabulary) * 1.0)
        else:
            return ngram_counter / (context_counter * 1.0)
    """
    prob with smoothing
    returns: Laplace-smoothed probabilities 
    (need to modify to fit n-gram)
    """

    def word_prob(self, word, context, delta=0.0):
        if delta == 0:
            return self.word_prob_delta_zero(word, context)
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        return (ngram_counter + delta) / ((context_counter + delta * len(
            self.vocabulary)) * 1.0)
    """
    returns a word sampled from the model's 
    probability distribution for context
    """

    def random_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        vocabulary = sorted(vocabulary)
        distrib = []
        total_counts = 0
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            total_counts += count
            distrib.append(total_counts)
        r = random.random()
        for i in range(len(distrib)):
            distrib[i] = distrib[i] / (total_counts * 1.0)
            if r < distrib[i]:
                return vocabulary[i]
        return vocabulary[len(vocabulary) - 1]
    """
        returns n-gram with the highest prob for context
    """

    def likeliest_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        distrib = []
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            distrib.append(count)
        max_value = 0
        max_Index = 0
        for i in range(len(distrib)):
            if max_value < distrib[i]:
                max_value = distrib[i]
                max_Index = i
        return vocabulary[max_Index]
    """
        return the best k words of the given context
    """

    def beam_search_word(self, context, k, delta=0.0):
        distrib = {}
        best_k = []
        for word, _ in self.vocabulary.items():
            distrib[word] = self.word_prob(word, context, delta)
        ordered_distrib = OrderedDict(sorted(distrib.items(), key=lambda kv:
            kv[1], reverse=True))
        i = 0
        for word, prob in ordered_distrib.items():
            if i >= k:
                break
            best_k.append(Candidate(word, prob))
            i += 1
        return best_k


<mask token>


class NGramInterpolator:
    """
        n - size of the largest n-gram
        lambdas - a list of length n containing the interpolation factors
            (floats) in descending order of n-gram size.

        save n and lambdas
        initialize n internal NGramLMs, one for each n-gram size
    """

    def __init__(self, n, lambdas):
        self.n = n
        self.lambdas = lambdas
        self.models = []
        for i in range(n):
            self.models.append(NGramLM(n - i))
    """
        update all of the internal NGramLMs
    """

    def update(self, text):
        for model in self.models:
            model.update(text)
    """
        returns the linearly interpolated prob using lambdas
        and prob given by NGramLMs
    """

    def word_prob(self, word, context, delta=0):
        prob = 0
        for i in range(self.n):
            prob += self.lambdas[i] * self.models[i].word_prob(word,
                context, delta)
        return prob


class Candidate:
    """
        This class is just for convenience,
        avoid to keep return tuple
    """

    def __init__(self, text, prob):
        self.text = text
        self.prob = prob

    def append(self, word, prob):
        self.text = self.text + ' ' + word.strip()
        self.prob *= prob

    def append(self, other):
        self.text = self.text + ' ' + other.text.strip()
        self.prob *= other.prob

    def clone(self):
        return Candidate(self.text, self.prob)


<mask token>

import re
import math
import sys
import codecs
import random
from collections import OrderedDict
import queue
PREFIX = '<s> '
SUFFIX = ' </s>'
UNKNOWN = '<unk>'
<mask token>


def get_ngrams(n, text):
    prefix = PREFIX
    suffix = SUFFIX
    for i in range(n - 2):
        prefix += PREFIX
    sentences = re.split('[?,.!:;]+', text)
    for sentence in sentences:
        sentence = re.sub('[^\\w<>]', ' ', sentence)
        sentence = sentence.strip()
        sentence = prefix + sentence + suffix
        words = sentence.split()
        for i in range(n - 1, len(words)):
            context = ''
            for j in range(i - n + 1, i):
                context = context + ' ' + words[j].strip()
            yield words[i].strip(), context.strip()


<mask token>


def create_ngramlm(n, corpus_path, delta=0.0):
    ngramlm = NGramLM(n, delta)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        ngramlm.update(corpus)
    return ngramlm


<mask token>


def create_interpolator(n, lambdas, corpus_path):
    inter = NGramInterpolator(n, lambdas)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        corpus = mask_rare(corpus)
        inter.update(corpus)
    return inter


<mask token>


def text_prob(model, text, delta=0.0):
    log_prob = 0
    ngrams = get_ngrams(model.n, text)
    for word, context in ngrams:
        prob = model.word_prob(word, context, delta)
        log_prob += math.log(prob)
    return log_prob


<mask token>


def get_word_with_context(word, context):
    return context.strip() + ' ' + word.strip()


<mask token>


def mask_rare(corpus):
    all_words = {}
    rare_words = set()
    words = re.split('[^\\w]', corpus)
    for word in words:
        count = all_words.get(word, 0)
        count += 1
        all_words[word] = count
    for word, count in all_words.items():
        if count == 1:
            rare_words.add(word)
    print('all: ', len(all_words), ', rare: ', len(rare_words))
    masked_corpus = []
    sentences = corpus.split()
    for sentence in sentences:
        words = re.split('[^\\w]', sentence)
        for word in words:
            if word in rare_words:
                sentence = sentence.replace(word, UNKNOWN)
        masked_corpus.append(sentence)
        masked_corpus.append(' ')
    return ''.join(masked_corpus).strip()


<mask token>


def perplexity(model, corpus_path):
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        num_tokens = 0
        sentences = re.split('[?,.!:;]+', corpus)
        for sentence in sentences:
            words = re.split('[^\\w]', sentence)
            num_tokens += len(words)
        i = text_prob(model, corpus, model.delta) / num_tokens
        return 2 ** -i


def random_text(model, max_length, delta=0):
    context = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        context.append(PREFIX.strip())
    counter = 0
    sentence = ''
    while True:
        word = model.random_word(' '.join(context).strip(), delta)
        if word == suffix:
            return sentence.strip()
        sentence = sentence + ' ' + word
        counter += 1
        if counter > max_length:
            return sentence.strip()
        for i in range(len(context) - 1):
            context[i] = context[i + 1]
        context[len(context) - 1] = word


<mask token>


def likeliest_text(model, max_length, delta=0):
    context = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        context.append(PREFIX.strip())
    counter = 0
    sentence = ''
    while True:
        word = model.likeliest_word(' '.join(context).strip(), delta)
        if word == suffix:
            return sentence.strip()
        sentence = sentence + ' ' + word
        counter += 1
        if counter > max_length:
            return sentence.strip()
        for i in range(len(context) - 1):
            context[i] = context[i + 1]
        context[len(context) - 1] = word


def beam_search_text(model, max_length, k, delta=0.0):
    starter = []
    best_k = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        starter.append(PREFIX.strip())
    init_candidate = Candidate(' '.join(starter).strip(), 1)
    frontier = queue.Queue()
    best_k.append(init_candidate)
    frontier.put(init_candidate)
    for i in range(max_length):
        sub_k = []
        while not frontier.empty():
            candidate = frontier.get()
            words = candidate.text.split()
            if words[len(words) - 1] == suffix:
                continue
            context = words[len(words) - model.n + 1:len(words)]
            k_words = model.beam_search_word(' '.join(context).strip(), k,
                delta)
            sub_k = sub_k + get_new_candidates(model, candidate, k_words, delta
                )
            best_k.remove(candidate)
        best_k = merge(best_k, sub_k, k)
        for candidate in best_k:
            frontier.put(candidate)
    return best_k


def get_new_candidates(model, candidate, sub_k, delta=0.0):
    new_candidates = []
    for item in sub_k:
        new_candidate = candidate.clone()
        new_candidate.append(item)
        new_candidate.prob = text_prob(model, new_candidate.text, delta)
        new_candidates.append(new_candidate)
    return new_candidates


def merge(best_k, sub_k, k):
    temp = best_k + sub_k
    temp = sorted(temp, key=lambda x: x.prob, reverse=True)
    return temp[:k]


class NGramLM:

    def __init__(self, n, delta=0.0):
        self.n = n
        self.ngram_counts = {}
        self.context_counts = {}
        self.vocabulary = {}
        self.delta = delta
    """
    updates NGramLM's internal counts
    and vocabulary for ngrams in text (list of words/strings)
    """

    def update(self, text):
        ngrams = get_ngrams(self.n, text)
        for word, context in ngrams:
            word_with_context = get_word_with_context(word, context)
            ngram_counter = self.ngram_counts.get(word_with_context, 0)
            context_counter = self.context_counts.get(context, 0)
            vocabulary_counter = self.vocabulary.get(word, 0)
            ngram_counter += 1
            context_counter += 1
            vocabulary_counter += 1
            self.ngram_counts[word_with_context] = ngram_counter
            self.context_counts[context] = context_counter
            self.vocabulary[word] = vocabulary_counter
    """
    returns prob of n-gram (word, context) using internal counters.
    If context is previously unseen, returns 1/|V|, V is vocabulary model
    """

    def word_prob_delta_zero(self, word, context):
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        if context_counter == 0 or word not in self.vocabulary:
            return self.vocabulary[UNKNOWN] / (len(self.vocabulary) * 1.0)
        else:
            return ngram_counter / (context_counter * 1.0)
    """
    prob with smoothing
    returns: Laplace-smoothed probabilities 
    (need to modify to fit n-gram)
    """

    def word_prob(self, word, context, delta=0.0):
        if delta == 0:
            return self.word_prob_delta_zero(word, context)
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        return (ngram_counter + delta) / ((context_counter + delta * len(
            self.vocabulary)) * 1.0)
    """
    returns a word sampled from the model's 
    probability distribution for context
    """

    def random_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        vocabulary = sorted(vocabulary)
        distrib = []
        total_counts = 0
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            total_counts += count
            distrib.append(total_counts)
        r = random.random()
        for i in range(len(distrib)):
            distrib[i] = distrib[i] / (total_counts * 1.0)
            if r < distrib[i]:
                return vocabulary[i]
        return vocabulary[len(vocabulary) - 1]
    """
        returns n-gram with the highest prob for context
    """

    def likeliest_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        distrib = []
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            distrib.append(count)
        max_value = 0
        max_Index = 0
        for i in range(len(distrib)):
            if max_value < distrib[i]:
                max_value = distrib[i]
                max_Index = i
        return vocabulary[max_Index]
    """
        return the best k words of the given context
    """

    def beam_search_word(self, context, k, delta=0.0):
        distrib = {}
        best_k = []
        for word, _ in self.vocabulary.items():
            distrib[word] = self.word_prob(word, context, delta)
        ordered_distrib = OrderedDict(sorted(distrib.items(), key=lambda kv:
            kv[1], reverse=True))
        i = 0
        for word, prob in ordered_distrib.items():
            if i >= k:
                break
            best_k.append(Candidate(word, prob))
            i += 1
        return best_k


<mask token>


class NGramInterpolator:
    """
        n - size of the largest n-gram
        lambdas - a list of length n containing the interpolation factors
            (floats) in descending order of n-gram size.

        save n and lambdas
        initialize n internal NGramLMs, one for each n-gram size
    """

    def __init__(self, n, lambdas):
        self.n = n
        self.lambdas = lambdas
        self.models = []
        for i in range(n):
            self.models.append(NGramLM(n - i))
    """
        update all of the internal NGramLMs
    """

    def update(self, text):
        for model in self.models:
            model.update(text)
    """
        returns the linearly interpolated prob using lambdas
        and prob given by NGramLMs
    """

    def word_prob(self, word, context, delta=0):
        prob = 0
        for i in range(self.n):
            prob += self.lambdas[i] * self.models[i].word_prob(word,
                context, delta)
        return prob


class Candidate:
    """
        This class is just for convenience,
        avoid to keep return tuple
    """

    def __init__(self, text, prob):
        self.text = text
        self.prob = prob

    def append(self, word, prob):
        self.text = self.text + ' ' + word.strip()
        self.prob *= prob

    def append(self, other):
        self.text = self.text + ' ' + other.text.strip()
        self.prob *= other.prob

    def clone(self):
        return Candidate(self.text, self.prob)


def main(argv):
    trimodel = create_ngramlm(3, 'shakespeare.txt')
    for sentence in beam_search_text(trimodel, 10, 5, 1):
        print(sentence.text)
    pass


if __name__ == '__main__':
    main(sys.argv)

import re
import math
import sys
import codecs
import random
from collections import OrderedDict
import queue

PREFIX = '<s> '
SUFFIX = ' </s>'
UNKNOWN = '<unk>'

"""
n - size of n-grams
text - list of words/strings
returns: tuple of n-1 (word, context)
"""


def get_ngrams(n, text):  # - generator
    prefix = PREFIX
    suffix = SUFFIX
    for i in range(n - 2):
        prefix += PREFIX

    sentences = re.split(r'[?,.!:;]+', text)
    for sentence in sentences:
        sentence = re.sub(r'[^\w<>]', ' ', sentence)
        sentence = sentence.strip()
        sentence = prefix + sentence + suffix
        words = sentence.split()

        for i in range(n - 1, len(words)):
            context = ''
            for j in range(i - n + 1, i):
                context = context + ' ' + words[j].strip()
            yield words[i].strip(), context.strip()


"""
n - number of grams
corpus path - file path
returns: NGramLM trained model
"""


def create_ngramlm(n, corpus_path, delta=.0):
    ngramlm = NGramLM(n, delta)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        # corpus = mask_rare(corpus)
        ngramlm.update(corpus)

    return ngramlm


"""
n - number of grams
corpus path - file path
returns: NGramLM trained model
"""


def create_interpolator(n, lambdas, corpus_path):
    inter = NGramInterpolator(n, lambdas)
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        corpus = mask_rare(corpus)
        inter.update(corpus)

    return inter


"""
returns log(e) prob of text (list of words/strings)
under NGramLM trained model
"""


def text_prob(model, text, delta=.0):
    log_prob = 0
    ngrams = get_ngrams(model.n, text)
    for word, context in ngrams:
        prob = model.word_prob(word, context, delta)
        log_prob += math.log(prob)

    return log_prob


"""
concatenate word with context
"""


def get_word_with_context(word, context):
    return context.strip() + ' ' + word.strip()


"""
corpus - content of file
returns: a copy of corpus with words 
that appear only once replaced by the `<unk>' token
"""


def mask_rare(corpus):
    all_words = {}
    rare_words = set()
    words = re.split(r'[^\w]', corpus)
    for word in words:
        count = all_words.get(word, 0)
        count += 1
        all_words[word] = count

    for word, count in all_words.items():
        if count == 1:
            rare_words.add(word)
    print("all: ", len(all_words), ", rare: ", len(rare_words))

    masked_corpus = []
    sentences = corpus.split()
    for sentence in sentences:
        words = re.split(r'[^\w]', sentence)
        for word in words:
            if word in rare_words:
                sentence = sentence.replace(word, UNKNOWN)
        masked_corpus.append(sentence)
        masked_corpus.append(' ')

    return ''.join(masked_corpus).strip()


"""
model - trained model
corpus_path - test data file path, no need to mask rare words
returns perplexity of a trained model
counts N, total tokens
"""


def perplexity(model, corpus_path):
    with codecs.open(corpus_path, 'r', 'utf-8') as file:
        corpus = file.read()
        num_tokens = 0

        sentences = re.split(r'[?,.!:;]+', corpus)
        for sentence in sentences:
            words = re.split(r'[^\w]', sentence)
            num_tokens += len(words)

        i = text_prob(model, corpus, model.delta) / num_tokens
        return 2 ** -i


def random_text(model, max_length, delta=0):
    context = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        context.append(PREFIX.strip())
    counter = 0
    sentence = ''

    while True:
        word = model.random_word(' '.join(context).strip(), delta)
        if word == suffix:
            return sentence.strip()

        sentence = sentence + ' ' + word
        counter += 1
        if counter > max_length:
            return sentence.strip()

        for i in range(len(context) - 1):
            context[i] = context[i + 1]
        context[len(context) - 1] = word


"""
almost identical to random_text
"""


def likeliest_text(model, max_length, delta=0):
    context = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        context.append(PREFIX.strip())
    counter = 0
    sentence = ''

    while True:
        word = model.likeliest_word(' '.join(context).strip(), delta)
        if word == suffix:
            return sentence.strip()

        sentence = sentence + ' ' + word
        counter += 1
        if counter > max_length:
            return sentence.strip()

        for i in range(len(context) - 1):
            context[i] = context[i + 1]
        context[len(context) - 1] = word


def beam_search_text(model, max_length, k, delta=.0):
    # here we first build the first candidate '<s> <s>'
    starter = []
    best_k = []
    suffix = SUFFIX.strip()
    for i in range(model.n - 1):
        starter.append(PREFIX.strip())
    init_candidate = Candidate(' '.join(starter).strip(), 1)  # check the 'Candidate' class below

    # put the first candidate into queue and make it as the current best_k
    frontier = queue.Queue()
    best_k.append(init_candidate)
    frontier.put(init_candidate)

    # loop until we got the max_length of words
    for i in range(max_length):
        # use sub_k to store the next level best k words followed by the current best_k words
        # sub_k is reset at each level, i.e. for different i
        sub_k = []
        while not frontier.empty():
            # poll a text from the queue
            # if the last word of the text is suffix <\s>, skip it
            candidate = frontier.get()
            words = candidate.text.split()
            if words[len(words) - 1] == suffix:
                continue

            # use the last n words to build context
            # get the next k_words of the context
            # append each of the k_words to current text, which yields k new texts
            # remove the current text from best_k, since we already find its descendants
            context = words[len(words) - model.n + 1: len(words)]
            k_words = model.beam_search_word(' '.join(context).strip(), k, delta)
            sub_k = sub_k + get_new_candidates(model, candidate, k_words, delta)
            best_k.remove(candidate)

        # merge all the elements of best_k and sub_k to elect new best_k
        # put the new best_k texts to queue, to find their descendants
        best_k = merge(best_k, sub_k, k)
        for candidate in best_k:
            frontier.put(candidate)

    return best_k


def get_new_candidates(model, candidate, sub_k, delta=.0):
    # append each of the sub_k to current text,
    # returns candidate's k descendants
    new_candidates = []
    for item in sub_k:
        new_candidate = candidate.clone()
        new_candidate.append(item)
        new_candidate.prob = text_prob(model, new_candidate.text, delta)
        new_candidates.append(new_candidate)

    return new_candidates


def merge(best_k, sub_k, k):
    # elect new best_k
    temp = best_k + sub_k
    temp = sorted(temp, key=lambda x: x.prob, reverse=True)
    return temp[:k]


class NGramLM:
    def __init__(self, n, delta=.0):
        self.n = n  # size of n-grams
        self.ngram_counts = {}  # n-grams seen in the training data,
        self.context_counts = {}  # contexts seen in the training data,
        self.vocabulary = {}  # words seen in the traniing data
        self.delta = delta

    """
    updates NGramLM's internal counts
    and vocabulary for ngrams in text (list of words/strings)
    """

    def update(self, text):
        ngrams = get_ngrams(self.n, text)
        for word, context in ngrams:
            word_with_context = get_word_with_context(word, context)
            ngram_counter = self.ngram_counts.get(word_with_context, 0)
            context_counter = self.context_counts.get(context, 0)
            vocabulary_counter = self.vocabulary.get(word, 0)
            ngram_counter += 1
            context_counter += 1
            vocabulary_counter += 1
            self.ngram_counts[word_with_context] = ngram_counter
            self.context_counts[context] = context_counter
            self.vocabulary[word] = vocabulary_counter

        # print(self.vocabulary[UNKNOWN])

    """
    returns prob of n-gram (word, context) using internal counters.
    If context is previously unseen, returns 1/|V|, V is vocabulary model
    """

    def word_prob_delta_zero(self, word, context):
        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        if context_counter == 0 or word not in self.vocabulary:
            return self.vocabulary[UNKNOWN] / (len(self.vocabulary) * 1.0)
        else:
            return ngram_counter / (context_counter * 1.0)


    """
    prob with smoothing
    returns: Laplace-smoothed probabilities 
    (need to modify to fit n-gram)
    """

    def word_prob(self, word, context, delta=.0):
        if delta == 0:
            return self.word_prob_delta_zero(word, context)

        word_with_context = get_word_with_context(word, context)
        ngram_counter = self.ngram_counts.get(word_with_context, 0)
        context_counter = self.context_counts.get(context, 0)
        return (ngram_counter + delta) / ((context_counter + delta * len(self.vocabulary)) * 1.0)


    """
    returns a word sampled from the model's 
    probability distribution for context
    """

    def random_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())
        vocabulary = sorted(vocabulary)

        distrib = []
        total_counts = 0
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            total_counts += count
            distrib.append(total_counts)

        r = random.random()
        for i in range(len(distrib)):
            distrib[i] = distrib[i] / (total_counts * 1.0)
            if r < distrib[i]:
                return vocabulary[i]

        return vocabulary[len(vocabulary) - 1]


    """
        returns n-gram with the highest prob for context
    """

    def likeliest_word(self, context, delta=0):
        vocabulary = list(self.vocabulary.keys())

        distrib = []
        for word in vocabulary:
            word_with_context = get_word_with_context(word, context)
            count = self.ngram_counts.get(word_with_context, 0)
            distrib.append(count)

        max_value = 0
        max_Index = 0
        for i in range(len(distrib)):
            if max_value < distrib[i]:
                max_value = distrib[i]
                max_Index = i

        return vocabulary[max_Index]


    """
        return the best k words of the given context
    """

    def beam_search_word(self, context, k, delta=.0):
        distrib = {}
        best_k = []

        # this loop loads every possible word following the given context with its counts into distrib{}
        for word, _ in self.vocabulary.items():
            distrib[word] = self.word_prob(word, context, delta)

        # rank the probabilities in the order of high to low
        ordered_distrib = OrderedDict(sorted(distrib.items(), key=lambda kv: kv[1], reverse=True))

        # this loop loads the top ranked k words into best_k
        i = 0
        for word, prob in ordered_distrib.items():
            if i >= k:
                break
            best_k.append(Candidate(word, prob))
            i += 1

        return best_k


"""
	linear interpolation.
"""

class NGramInterpolator:


    """
        n - size of the largest n-gram
        lambdas - a list of length n containing the interpolation factors
            (floats) in descending order of n-gram size.

        save n and lambdas
        initialize n internal NGramLMs, one for each n-gram size
    """

    def __init__(self, n, lambdas):
        self.n = n
        self.lambdas = lambdas
        self.models = []
        for i in range(n):
            self.models.append(NGramLM(n - i))


    """
        update all of the internal NGramLMs
    """

    def update(self, text):
        for model in self.models:
            model.update(text)


    """
        returns the linearly interpolated prob using lambdas
        and prob given by NGramLMs
    """

    def word_prob(self, word, context, delta=0):
        prob = 0
        for i in range(self.n):
            prob += self.lambdas[i] * self.models[i].word_prob(word, context, delta)

        return prob


class Candidate:

    """
        This class is just for convenience,
        avoid to keep return tuple
    """

    def __init__(self, text, prob):
        self.text = text
        self.prob = prob

    def append(self, word, prob):
        # append a word to text
        # the prob is overwritten later by text_prob,
        # since I original design it to be the real probability.
        # this has the potential of underflow
        self.text = self.text + ' ' + word.strip()
        self.prob *= prob

    def append(self, other):
        # append another candidate to this candidate
        # text is being appended
        # the prob is overwritten later by text_prob,
        # since I original design it to be the real probability.
        # this has the potential of underflow
        self.text = self.text + ' ' + other.text.strip()
        self.prob *= other.prob

    def clone(self):
        # create a copy of current candidate
        return Candidate(self.text, self.prob)


def main(argv):
    trimodel = create_ngramlm(3, "shakespeare.txt")
    for sentence in beam_search_text(trimodel, 10, 5, 1):
        print(sentence.text)

    pass


if __name__ == '__main__':
    main(sys.argv)
