import copy
import random
import logging


def get_random_subset(dataset_object, num_examples_per_class):
    """
    Return a copy of the dataset object that only contains a random subset of the original data items
    Parameters
    ----------
    dataset_object: tv_net.dataset.Dataset
        instance of the dataset class to take subset from
    num_examples_per_class: int
        number of examples to include in subset

    Returns
    -------
    subset: tv_net.dataset.Dataset
        new dataset object that contains only a random subset
    """
    subset = copy.copy(dataset_object)
    subset.items = list()  # re-initialise to empty list
    for class_name in dataset_object.class_names:
        # Get a list of the examples belonging to this class
        class_examples = [example for example in dataset_object.items if example.class_name == class_name]
        random.shuffle(class_examples)
        # take subset
        subset.items += class_examples[:num_examples_per_class]

        if len(class_examples) < num_examples_per_class:
            logging.warning('Class: {} only contains {} training examples'.format(class_name, len(class_examples)))
            logging.warning('This is less than the {} examples per class specified'.format(
                num_examples_per_class))
            
    subset.num_examples = len(subset.items)
    return subset