import collections


def scale_cluster_zero(filenames, clustering, size_multiplier=4):
    """
    Scales cluster zero to a multiple of the second largest class

    Parameters
    ----------
    filenames           :   list of str
                            list of strings containing file names of images
    clustering          :   list of int
                            list of integers giving the cluster assignment of each image in filenames
    size_multiplier     :   int, default=4
                            Determines how to scale cluster 0. Cluster 0 will have, at maximum, as many images as
                            size_multiplier times the size of the second largest cluster.

    Returns
    -------
    scaled_filenames    : list of str
    scaled_clustering   : list of str
    """
    scaled_filenames = []
    scaled_clustering = []

    cluster_sizes = collections.Counter(clustering)
    size_of_class_0_prior = cluster_sizes.most_common(2)[0][1]
    second_largest_cluster_size = cluster_sizes.most_common(2)[1][1]
    size_of_class_zero = size_multiplier * second_largest_cluster_size

    print("Reducing size of class 0 from {} to {} based on second largest cluster (size {})".format(
        size_of_class_0_prior, size_of_class_zero, second_largest_cluster_size
    ))

    admitted = 0
    for fn, cluster in zip(filenames, clustering):
        if cluster == 0:
            if admitted >= size_of_class_zero:
                continue
            else:
                admitted += 1
                scaled_filenames.append(fn)
                scaled_clustering.append(cluster)
        else:
            scaled_filenames.append(fn)
            scaled_clustering.append(cluster)

    return scaled_filenames, scaled_clustering