import torch


def update_boltz(conf_fp, weight, boltz_nn):
    """
    Given a conformer fingerprint and Boltzmann weight,
    return a new updated fingerprint.
    Args:
        conf_fp (torch.Tensor): molecular finerprint of
            a conformer
        weight (float): Boltzmann weight
        boltz_nn (torch.nn.Module): network that converts
            the fingerprint and weight into a new
            fingerprint. If None, just multiply the Boltzmann
            factor with the fingerprint.
    Returns:
        boltzmann_fp (torch.Tensor): updated fingerprint
    """

    if boltz_nn is None:
        boltzmann_fp = conf_fp * weight
    # otherwise concatenate the weight with the fingerprint
    # and put it through the boltzmann nn
    else:
        weight_tens = torch.Tensor([weight]).to(conf_fp.device)
        new_fp = torch.cat((conf_fp, weight_tens))
        boltzmann_fp = boltz_nn(new_fp)
    return boltzmann_fp