from collections import Counter


def gen_idx_byclass(labels):
    """
    Neatly organize indices of labeled samples by their classes.

    Parameters
    ----------
    labels : list
        Note that labels should be a simple Python list instead of a tensor.

    Returns
    -------
    idx_byclass : dictionary {[class_label (int) : indices (list)]}
    """
    # print("in gen_idx_byclass...")
    classes = Counter(labels).keys()  # obtain a list of classes
    idx_byclass = {}

    for class_label in classes:
        # Find samples of this class:
        class_idx = []  # indices for samples that belong to this class
        for idx in range(len(labels)):
            if labels[idx] == class_label:
                class_idx.append(idx)
        idx_byclass[class_label] = class_idx

    return idx_byclass