from typing import Counter
import pickle


def build_dict(examples, filepath, max_words=50000):
    """
        Build a dictionary for the words in `sentences` and a dictionary for the entities.
        Only the max_words ones are kept and the remaining will be mapped to <UNK>.
    """
    documents, questions, answers = examples
    word_count = Counter()
    for sent in documents + questions:
        for w in sent.split(' '):
            word_count[w] += 1

    ls = word_count.most_common(max_words)
    print('Words: %d -> %d' % (len(word_count), len(ls)))

    # leave 0 to PAD
    # leave 1 to UNK
    word_dict = {w[0]: idx + 2 for (idx, w) in enumerate(ls)}
    word_dict['PAD'] = 0
    word_dict['UNK'] = 1
    entity_markers = ['<unk_entity>'] + list(set([w for w in word_dict.keys() if w.startswith('@entity')] + answers))
    entity_dict = {w: idx for (idx, w) in enumerate(entity_markers)}

    # save the dicts
    with open('{}-word-dict.pkl'.format(filepath), 'wb') as f:
        pickle.dump(word_dict, f)
    with open('{}-entity-dict.pkl'.format(filepath), 'wb') as f:
        pickle.dump(entity_dict, f)

    return word_dict, entity_dict