def onehot(idx, size):
    """One hot encoder
    """
    onehot = torch.zeros(size)
    # Fill the one-hot encoded sequence with 1 at the corresponding idx
    onehot[idx] = 1
    return np.array(onehot)

def format_node_edge_features(g):
    """Format the nodes and edges features computed with Graphein

    Args:
        g (object): graph

    Returns:
        object: updated graph
    """
    
    # one hot encoding
    residue_names = {'CYS': 0, 'HIS': 1, 'ASN': 2, 'GLN': 3, 'SER': 4, 'THR': 5, 'TYR': 6, 'TRP': 7,
                     'ALA': 8, 'PHE': 9, 'GLY': 10, 'ILE': 11, 'VAL': 12, 'MET': 13, 'PRO': 14, 'LEU': 15,
                     'GLU': 16, 'ASP': 17, 'LYS': 18, 'ARG': 19}
    
    edge_type_encoding = {
        'peptide_bond': 0, 'aromatic': 1, 'disulfide': 2, 'ionic': 3, 
        'aromatic_sulphur': 4, 'cation_pi' : 5, 'distance_threshold' : 6, 'hbond' : 7}
    
    # convert node information
    resname_onehot = []    
    for res in g.residue_name :
        # One hot encoding of the residue name
        resname_onehot.append(onehot(residue_names[res], len (residue_names)))

    g["residue"] = resname_onehot
    
    edge_onehot = []
    for res in g.kind :
        # One hot encoding of the edge type
        edge_onehot.append(onehot([edge_type_encoding[x] for x in res], len (edge_type_encoding)))

    g["edge_attr"] = edge_onehot

    return g