import torch


def to_xyz_idx(batch_dic):
    """
    Get the indices of the nxyz corresponding to atoms in conformers
    we want to keep.
    Args:
        batch_dic (dict): Dictionary with extra conformer 
            information about the batch
    Returns:
        xyz_conf_all_idx (torch.LongTensor): nxyz indices of atoms
            in the conformers we want to keep.
    """

    confs_in_batch = batch_dic["confs_in_batch"]
    mol_size = batch_dic["mol_size"]
    # Indices of the conformers we're keeping (usually of the form [0, 1, ...
    # `max_confs`], unless specified otherwise)
    conf_idx = batch_dic["conf_idx"]

    # the xyz indices of where each conformer starts
    xyz_conf_start_idx = [i * mol_size for i in range(confs_in_batch + 1)]
    # a list of the full set of indices for each conformer
    xyz_conf_all_idx = []

    # go through each conformer index, get the corresponding xyz indices,
    # and append them to xyz_conf_all_idx

    for conf_num in conf_idx:

        start_idx = xyz_conf_start_idx[conf_num]
        end_idx = xyz_conf_start_idx[conf_num + 1]
        full_idx = torch.arange(start_idx, end_idx)

        xyz_conf_all_idx.append(full_idx)

    # concatenate into tensor
    xyz_conf_all_idx = torch.cat(xyz_conf_all_idx)

    return xyz_conf_all_idx