import random


def select_subset(textdict, doccats, visids=[]):
    """
    select a random subset of the dataset if it contains more than 10000 examples

    Input and Returns:
        textdict: dict with {doc_id: text}
        doccats: dict with {doc_id: category}
        visids: a subset of docids for which the html visualization should be created
    """
    docids = sorted(textdict.keys())  # sort for consistency across OS
    random.seed(42)
    random.shuffle(docids)
    # visualize up to 1000 documents
    if not len(visids):
        visids = docids[:1000]
    elif len(visids) > 1000:
        print("WARNING: creating visualizations for %i, i.e. more than 1000 documents can be slow!" % len(visids))
        if len(visids) > 10000:
            print("You don't know what you're doing....Truncating visids to 5000 examples.")
            visids = visids[:5000]
    # select subsets of examples to speed up the computations
    if len(docids) > 10000:
        # always make sure you end up with exactly 10k random examples (incl visids) but also don't shuffle a lot more than 10k ids
        docids = list(set(docids[:10000+len(visids)]).difference(set(visids)))
        random.shuffle(docids)
        docids = docids[:10000-len(visids)] + visids
        textdict = {d: textdict[d] for d in docids}
        doccats = {d: doccats[d] for d in docids}
    return textdict, doccats, visids