from typing import Dict
from typing import List
from typing import Any
import random


def make_balanced_labeled_dict(data_set: Dict[int, List[Any]], max_count: int, randomizer: random.Random) -> dict:
    """
    Takes a dict of data with labels (int:list()) and returns a dict where each
    class is balanced size-wise where the number of elements in each class does not
    exceed max_count. If max_count is -1, then the smallest class element count found
    is used.  The number of elements per class is returned.
    :param data_set: Input data which is list of (str,int)
    :param max_count: The maximum elements per class or -1 to use existing size.
    :param randomizer: Randomizer used when sampling the categories.
    :return: Dict[int, List[Any]]
    """

    # Find smallest count
    smallest_count = min([len(v) for k, v in data_set.items()])
    if max_count == -1:
        max_count = smallest_count

    if smallest_count < max_count:
        max_count = smallest_count

    # Make a new dict with that subset
    new_dict = {}
    for k, v in data_set.items():
        new_dict[k] = randomizer.sample(v, max_count)

    return new_dict