import torch


def construct_edge_list(sim_score_mat, thres):                
    """ Constructs edge lists for a PyG graph (COO representation) 
    based on the pairwise similarity matrix of the 
    molecule and the molecular features
    """
    print('Constructing COO edge list based on similarity matrix.')
    srcs, dsts = [], []               
    n = len(sim_score_mat)
    for i in range(n):
        for j in range(n):
            if sim_score_mat[i][j] > thres:
                srcs.append(i)
                dsts.append(j)
    edge_index = torch.tensor([srcs,dsts], dtype=torch.long)
    print('Done.')
    return edge_index