import math


def get_confs_per_split(batch,
                        num_confs,
                        sub_batch_size):
    """
    Get the number of conformers per sub-batch.
    Args:
        batch (dict): batched sample of species
        num_confs (int): number of conformers in the species
        sub_batch_size (int): maximum number of conformers per
            sub-batch.
    Returns:
        confs_per_split (list[int]): number of conformers in each
            sub-batch.
    """

    val_len = len(batch["nxyz"])
    inherent_val_len = val_len // num_confs
    split_list = [sub_batch_size * inherent_val_len] * math.floor(
        num_confs / sub_batch_size)

    # if there's a remainder

    if sum(split_list) != val_len:
        split_list.append(val_len - sum(split_list))

    confs_per_split = [i // inherent_val_len for i in split_list]

    return confs_per_split