from typing import Counter


def create_vocabulary(train_set, vocab_size):
    """Creates vocabulary out of the training set tokens.

    Arguments:
        train_set: CoNLL2003 train_set from HuggingFace
        vocab_size (int): Maximum number of tokens in the vocab
    Returns:
        vocab (dict): Vocabulary of all tokens in the training set
            key: token
            value: ordinal number of token in the vocabulary
    """
    all_tokens = []
    for token_subseq in train_set["tokens"]:
        all_tokens += token_subseq

    # Perform some pre-processing of the tokens
    all_tokens_lower = list(map(str.lower, all_tokens))
    all_tokens_strip = list(map(str.strip, all_tokens_lower))

    # Count the occurence of every word
    counter = Counter(all_tokens_strip)
    # Extract VOCAB_SIZE - 2 since we will define tokens for padding elements
    # and words which aren't present in the training set
    most_frequent = counter.most_common(vocab_size - 2)

    # Initialize the vocabulary
    vocab = {
        "UNK": 0,
        "PADD": 1
    }
    ind = len(vocab)
    # Populate the vocab
    for token, _ in most_frequent:
        vocab[token] = ind
        ind += 1

    print("\nCreated vocabulary of {} tokens.".format(ind))
    return vocab