import torch


def balanced_spec_zhu(spec_dic,
                      zhu_p):
    """
    Get the Zhu weights assigned to each geom, such that
    the probability of getting a geom in species A
    is equal to the probability of getting a geom in
    species B [p(A) = p(B)], while the probabiltiies
    within a species are related by p(A, i) / p(A, j)
    = p_zhu(i) / p_zhu(j), where i and j are geometries in
    species A and p_zhu is the Zhu-Nakamura hopping probability.
    Args:
        spec_dic (dict): dictionary with indices of geoms in each
                    species
        zhu_p (torch.Tensor): Zhu-Nakamura hopping probabilities
                    for each geom.
    Returns:
                all_weights (torch.Tensor): sampling weights for each
                                geom in the dataset, normalized to 1.
    """

    num_geoms = sum([len(i) for i in spec_dic.values()])
    all_weights = torch.zeros(num_geoms)

    for lst_idx in spec_dic.values():

        idx = torch.LongTensor(lst_idx)
        this_zhu = zhu_p[idx]
        sum_zhu = this_zhu.sum()

        # make sure they're normalized first
        if sum_zhu != 0:
            this_zhu /= sum_zhu

        all_weights[idx] = this_zhu

    all_weights /= all_weights.sum()

    return all_weights