import torch


def get_rand_weights(spec_dic):
    """
    Generate weights for random sampling of geoms - i.e., equal weights
    for every geometry.
    Args:
        spec_dic (dict): dictionary with indices of geoms in each species.
    Returns:
        balanced_spec_weights (torch.Tensor): weights generated so that
                P(A) = P(B) for species A and B, and p(A, i) = p(A, j)
                for any geoms within A.
        imbalanced_spec_weights (torch.Tensor): weights generated so that
                P(A) != P(B) in general for species A and B, and p(i) = p(j)
                for any geoms.
    """

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

    imbalanced_spec_weights = torch.ones(num_geoms)
    imbalanced_spec_weights /= imbalanced_spec_weights.sum()

    balanced_spec_weights = torch.zeros(num_geoms)
    for idx in spec_dic.values():
        if len(idx) == 0:
            continue
        balanced_spec_weights[idx] = 1 / len(idx)

    total = balanced_spec_weights.sum()
    if total != 0:
        balanced_spec_weights /= total

    return balanced_spec_weights, imbalanced_spec_weights