class LanguageModel:

    def __init__(self, unigram, bigram_next, trigram_topk):
        self.unigram = unigram
        self.bigram_next = self._to_score_map(bigram_next)
        self.trigram_topk = self._to_score_map(trigram_topk)
        self.total = sum(unigram.values())

    @staticmethod
    def _to_score_map(ngram_source):
        score_map = {}
        for ctx, next_words in ngram_source.items():
            if isinstance(next_words, dict):
                score_map[ctx] = next_words
            else:
                score_map[ctx] = {word: count for word, count in next_words}
        return score_map

    def score(self, w1, w2, w3):
        tri = 0.0
        bi = 0.0
        uni = self.unigram.get(w3, 0)

        key = f"{w1}\t{w2}"
        tri = self.trigram_topk.get(key, {}).get(w3, 0)
        bi = self.bigram_next.get(w2, {}).get(w3, 0)
        return 0.6 * tri + 0.3 * bi + 0.1 * uni
