from typing import Counter


def build_vocab(annotations, threshold=1):
    """ Create vocabulary for text captions.

    Args:
        annotations: Dictionary with data, including text captions.
        threshold: Threshold on word occurrences in all captions.

    Returns:
        word_to_idx: Vocabulary as python dictionary.

    """
    counter = Counter()
    max_len = 0
    for i, caption in enumerate(annotations['caption']):
        words = caption.split(' ') # caption contrains only lower-case words
        for w in words:
            counter[w] +=1
        
        if len(caption.split(" ")) > max_len:
            max_len = len(caption.split(" "))

    vocab = [word for word in counter if counter[word] >= threshold]
    print('Filtered %d words to %d words with word count threshold %d.' % (len(counter), len(vocab), threshold))

    word_to_idx = {u'<PAD>': 0, u'<START>': 1, u'<END>': 2}
    idx = 3

    for word in vocab:
        word_to_idx[word] = idx
        idx += 1
    print("Max length of caption: %d" %max_len)

    return word_to_idx