from collections import Counter, defaultdict

class LanguageModelTrainer:

    def __init__(self, min_count=5, top_k=5):
        self.unigram = Counter()
        self.bigram = defaultdict(Counter)
        self.trigram = defaultdict(Counter)
        self.min_count = min_count
        self.top_k = top_k

    def update(self, tokens):

        self.unigram.update(tokens)

        for i in range(len(tokens) - 1):
            self.bigram[tokens[i]][tokens[i+1]] += 1

        for i in range(len(tokens) - 2):
            key = (tokens[i], tokens[i+1])
            self.trigram[key][tokens[i+2]] += 1

    def finalize(self):

        # prune unigram
        self.unigram = Counter({
            w:c for w,c in self.unigram.items()
            if c >= self.min_count
        })

        vocab = set(self.unigram.keys())

        # prune bigram
        bigram_topk = {}
        for w2, ctr in self.bigram.items():
            if w2 not in vocab:
                continue
            filtered = Counter({
                w3:c for w3,c in ctr.items()
                if w3 in vocab
            })
            if filtered:
                bigram_topk[w2] = filtered.most_common(self.top_k)

        # prune trigram
        trigram_topk = {}
        for (w1,w2), ctr in self.trigram.items():
            if w1 not in vocab or w2 not in vocab:
                continue
            filtered = Counter({
                w3:c for w3,c in ctr.items()
                if w3 in vocab
            })
            if filtered:
                trigram_topk[f"{w1}\t{w2}"] = filtered.most_common(self.top_k)

        return self.unigram, bigram_topk, trigram_topk