import torch


def featurize_nodes_and_compute_combo_scores(
        node_featurizer, reactant_mol, valid_candidate_combos):
    """Featurize atoms in reactants and compute scores for combos of bond changes

    Parameters
    ----------
    node_featurizer : callable, rdkit.Chem.rdchem.Mol -> dict
        Featurization for nodes like atoms in a molecule, which can be used to update
        ndata for a DGLGraph.
    reactant_mol : rdkit.Chem.rdchem.Mol
        RDKit molecule instance for reactants in a reaction
    valid_candidate_combos : list
        valid_candidate_combos[i] gives a list of tuples, which is the i-th valid combo
        of candidate bond changes for the reaction.

    Returns
    -------
    node_feats : float32 tensor of shape (N, M)
        Node features for reactants, N for the number of nodes and M for the feature size
    combo_bias : float32 tensor of shape (B, 1)
        Scores for combos of bond changes, B equals len(valid_candidate_combos)
    """
    node_feats = node_featurizer(reactant_mol)['hv']
    combo_bias = torch.zeros(len(valid_candidate_combos), 1).float()
    for combo_id, combo in enumerate(valid_candidate_combos):
        combo_bias[combo_id] = sum([
            score for (atom1, atom2, change_type, score) in combo])

    return node_feats, combo_bias