import torch

import torch.nn as nn

import torch.nn.functional as F

from torch import Tensor

from torch_sparse import SparseTensor

from torch_cluster import random_walk

from torch_geometric.data.sampler import EdgeIndex, Adj

from torch.nn.utils.rnn import pad_sequence

from torch.utils.data import Dataset

from torch_geometric.utils import add_self_loops, add_remaining_self_loops

from torch_geometric.data import Data, DataLoader, NeighborSampler

from typing import List, Optional, Tuple, NamedTuple, Union, Callable, Dict

from collections import defaultdict

import time

import random

import pickle

from collections import Counter

from operator import itemgetter

import copy

import numpy as np

from utils.pretrain_utils import get_indices_into_edge_index, HeterogeneousEdgeIndex 

from sklearn.preprocessing import label_binarize

import project_config

class NeighborSampler(torch.utils.data.DataLoader):

    r"""The neighbor sampler from the `"Inductive Representation Learning on

    Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper, which allows

    for mini-batch training of GNNs on large-scale graphs where full-batch

    training is not feasible.

    Given a GNN with :math:`L` layers and a specific mini-batch of nodes

    :obj:`node_idx` for which we want to compute embeddings, this module

    iteratively samples neighbors and constructs bipartite graphs that simulate

    the actual computation flow of GNNs.

    More specifically, :obj:`sizes` denotes how much neighbors we want to

    sample for each node in each layer.

    This module then takes in these :obj:`sizes` and iteratively samples

    :obj:`sizes[l]` for each node involved in layer :obj:`l`.

    In the next layer, sampling is repeated for the union of nodes that were

    already encountered.

    The actual computation graphs are then returned in reverse-mode, meaning

    that we pass messages from a larger set of nodes to a smaller one, until we

    reach the nodes for which we originally wanted to compute embeddings.

    Hence, an item returned by :class:`NeighborSampler` holds the current

    :obj:`batch_size`, the IDs :obj:`n_id` of all nodes involved in the

    computation, and a list of bipartite graph objects via the tuple

    :obj:`(edge_index, e_id, size)`, where :obj:`edge_index` represents the

    bipartite edges between source and target nodes, :obj:`e_id` denotes the

    IDs of original edges in the full graph, and :obj:`size` holds the shape

    of the bipartite graph.

    For each bipartite graph, target nodes are also included at the beginning

    of the list of source nodes so that one can easily apply skip-connections

    or add self-loops.

    .. note::

        For an example of using :obj:`NeighborSampler`, see

        `examples/reddit.py

        <https://github.com/rusty1s/pytorch_geometric/blob/master/examples/

        reddit.py>`_ or

        `examples/ogbn_products_sage.py

        <https://github.com/rusty1s/pytorch_geometric/blob/master/examples/

        ogbn_products_sage.py>`_.

    Args:

        edge_index (Tensor or SparseTensor): A :obj:`torch.LongTensor` or a

            :obj:`torch_sparse.SparseTensor` that defines the underlying graph

            connectivity/message passing flow.

            :obj:`edge_index` holds the indices of a (sparse) symmetric

            adjacency matrix.

            If :obj:`edge_index` is of type :obj:`torch.LongTensor`, its shape

            must be defined as :obj:`[2, num_edges]`, where messages from nodes

            :obj:`edge_index[0]` are sent to nodes in :obj:`edge_index[1]`

            (in case :obj:`flow="source_to_target"`).

            If :obj:`edge_index` is of type :obj:`torch_sparse.SparseTensor`,

            its sparse indices :obj:`(row, col)` should relate to

            :obj:`row = edge_index[1]` and :obj:`col = edge_index[0]`.

            The major difference between both formats is that we need to input

            the *transposed* sparse adjacency matrix.

        sizes ([int]): The number of neighbors to sample for each node in each

            layer. If set to :obj:`sizes[l] = -1`, all neighbors are included

            in layer :obj:`l`.

        node_idx (LongTensor, optional): The nodes that should be considered

            for creating mini-batches. If set to :obj:`None`, all nodes will be

            considered.

        num_nodes (int, optional): The number of nodes in the graph.

            (default: :obj:`None`)

        return_e_id (bool, optional): If set to :obj:`False`, will not return

            original edge indices of sampled edges. This is only useful in case

            when operating on graphs without edge features to save memory.

            (default: :obj:`True`)

        transform (callable, optional): A function/transform that takes in

            an a sampled mini-batch and returns a transformed version.

            (default: :obj:`None`)

        **kwargs (optional): Additional arguments of

            :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,

            :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.

    """

    def __init__(self, dataset_type: str, edge_index: Union[Tensor, SparseTensor], 

                sample_edge_index: Union[Tensor, SparseTensor],

                 sizes: List[int],

                 node_idx: Optional[Tensor] = None,

                 num_nodes: Optional[int] = None, return_e_id: bool = True,

                 transform: Callable = None,

                 do_filter_edges: bool = True, 

                 **kwargs):

        edge_index = edge_index.to('cpu')

        sample_edge_index = sample_edge_index.to('cpu')

        # add self loops

        sample_edge_index, _ = add_self_loops(sample_edge_index)

        if 'collate_fn' in kwargs:

            del kwargs['collate_fn']

        # Save for Pytorch Lightning...

        self.dataset_type = dataset_type

        self.edge_index = edge_index #always train edge index

        self.sample_edge_index = sample_edge_index # depends on train/val/test

        self.node_idx = node_idx

        self.num_nodes = num_nodes

        self.sizes = sizes

        self.return_e_id = return_e_id

        self.transform = transform

        self.is_sparse_tensor = isinstance(edge_index, SparseTensor)

        self.__val__ = None

        self.do_filter_edges = do_filter_edges

        # Obtain a *transposed* `SparseTensor` instance.

        if not self.is_sparse_tensor:

            if (num_nodes is None and node_idx is not None

                    and node_idx.dtype == torch.bool):

                num_nodes = node_idx.size(0)

                sample_num_nodes = num_nodes

            if (num_nodes is None and node_idx is not None

                    and node_idx.dtype == torch.long):

                num_nodes = max(int(edge_index.max()), int(node_idx.max())) + 1

                sample_num_nodes = num_nodes

            if num_nodes is None:

                num_nodes = int(edge_index.max()) + 1

                sample_num_nodes = int(sample_edge_index.max()) + 1

            value = torch.arange(edge_index.size(1)) if return_e_id else None

            sample_value = torch.arange(sample_edge_index.size(1)) if return_e_id else None

            self.adj_t = SparseTensor(row=edge_index[0], col=edge_index[1],

                                      value=value,

                                      sparse_sizes=(num_nodes, num_nodes)).t()

            self.adj_t_sample = SparseTensor(row=sample_edge_index[0], col=sample_edge_index[1],

                                      value=sample_value,

                                      sparse_sizes=(sample_num_nodes, sample_num_nodes)).t()

        else:

            adj_t = edge_index

            adj_t_sample = sample_edge_index

            if return_e_id:

                self.__val__ = adj_t.storage.value()

                value = torch.arange(adj_t.nnz())

                adj_t = adj_t.set_value(value, layout='coo')

                adj_t_sample = adj_t_sample.set_value(torch.arange(adj_t_sample.nnz()), layout='coo')

            self.adj_t = adj_t

            self.adj_t_sample = adj_t_sample

        self.adj_t.storage.rowptr()

        self.adj_t_sample.storage.rowptr()

        if node_idx is None:

            node_idx = torch.arange(self.adj_t_sample.sparse_size(0)) 

        elif node_idx.dtype == torch.bool:

            node_idx = node_idx.nonzero(as_tuple=False).view(-1)

        super(NeighborSampler, self).__init__(

            node_idx.view(-1).tolist(), collate_fn=self.sample, **kwargs)

    def filter_edges(self, edge_index, e_id, source_nodes, target_nodes):

        '''

        Filter out the edges we're trying to predict in the current batch from the edge index

        NOTE: edge_index here is re-indexed

        '''

        reindex_source_nodes = torch.arange(source_nodes.size(0))

        reindex_target_nodes = torch.arange(start = source_nodes.size(0), end = source_nodes.size(0) + target_nodes.size(0))

        # get reverse edges to filter as well

        all_source_nodes = torch.cat([reindex_source_nodes, reindex_target_nodes])

        all_target_nodes = torch.cat([reindex_target_nodes, reindex_source_nodes])

        ind_to_edge_index, ind_to_nodes = get_indices_into_edge_index(edge_index, all_source_nodes, all_target_nodes) #get index into the original edge index (this returns e_ids)

        mask = torch.ones(edge_index.size(1), dtype=torch.bool)

        mask[ind_to_edge_index] = False

        return edge_index[:, mask], e_id[mask]

    def sample(self, source_batch):

        #convert to tensor

        if not isinstance(source_batch, Tensor):

            source_batch = torch.tensor(source_batch)

        # sample nodes to form positive edges. we will try to predict these edges

        row, col, e_id = self.adj_t_sample.coo()    

        target_batch = random_walk(row, col, source_batch, walk_length=1, coalesced=False)[:, 1] #NOTE: only does self loops when no edges in the current partition of the dataset

        batch = torch.cat([source_batch, target_batch], dim=0) 

        batch_size: int = len(batch)

        adjs = []

        n_id = batch

        for size in self.sizes:

            adj_t, n_id = self.adj_t.sample_adj(n_id, size, replace=False) 

            e_id = adj_t.storage.value()

            size = adj_t.sparse_sizes()[::-1]

            if self.__val__ is not None:

                adj_t.set_value_(self.__val__[e_id], layout='coo')

            if self.is_sparse_tensor: #TODO: implement filter_edges if sparse tensor

                adjs.append(Adj(adj_t, e_id, size))

            else:

                row, col, _ = adj_t.coo()

                edge_index = torch.stack([col, row], dim=0)

                if self.do_filter_edges and self.dataset_type == 'train':

                    edge_index, e_id = self.filter_edges(edge_index, e_id, source_batch, target_batch)

                adjs.append(EdgeIndex(edge_index, e_id, size))

        adjs = adjs[0] if len(adjs) == 1 else adjs[::-1]

        out = (batch_size, n_id, adjs)

        out = self.transform(*out) if self.transform is not None else out

        return out

    def __repr__(self):

        return '{}(sizes={})'.format(self.__class__.__name__, self.sizes)

class PatientNeighborSampler(torch.utils.data.DataLoader):

    def __init__(self, dataset_type: str, edge_index: Union[Tensor, SparseTensor], 

                 sample_edge_index: Union[Tensor, SparseTensor],

                 sizes: List[int], 

                 patient_dataset,

                 all_edge_attributes,

                 n_nodes: int,

                 relevant_node_idx = None,

                 do_filter_edges: Optional[bool] = False,

                 num_nodes: Optional[int] = None, 

                 return_e_id: bool = True,

                 sparse_sample: Optional[int] = 0,

                 train_phenotype_counter: Dict = None,

                 train_gene_counter: Dict = None,

                 sample_edges_from_train_patients=False,

                 upsample_cand: Optional[int] = 0,

                 n_cand_diseases=-1,

                 use_diseases=False,

                 nid_to_spl_dict = None,

                 gp_spl = None,

                 spl_indexing_dict=None,

                 gene_similarity_dict=None,

                 gene_deg_dict = None,

                 hparams=None,

                 transform: Callable = None, 

                 **kwargs):

        edge_index = edge_index.to('cpu')

        sample_edge_index = sample_edge_index.to('cpu')

        # add self loops

        sample_edge_index = torch.cat((sample_edge_index, torch.stack([edge_index.unique(), edge_index.unique()])),1 )

        sample_edge_index, _ = add_remaining_self_loops(sample_edge_index)

        if 'collate_fn' in kwargs:

            del kwargs['collate_fn']

        # Save for Pytorch Lightning...

        self.do_filter_edges = do_filter_edges

        self.relevant_node_idx = relevant_node_idx

        self.n_nodes = n_nodes

        self.all_edge_attr = all_edge_attributes

        self.dataset_type = dataset_type

        self.sparse_sample = sparse_sample

        self.edge_index = edge_index #always train edge index

        self.sample_edge_index = sample_edge_index # depends on train/val/test

        self.patient_dataset = patient_dataset

        self.num_nodes = num_nodes

        self.train_phenotype_counter = train_phenotype_counter

        self.train_gene_counter = train_gene_counter

        self.sample_edges_from_train_patients = sample_edges_from_train_patients

        self.sizes = sizes

        self.return_e_id = return_e_id

        self.transform = transform

        self.is_sparse_tensor = isinstance(edge_index, SparseTensor)

        self.__val__ = None

        # For SPL

        self.nid_to_spl_dict = nid_to_spl_dict 

        if hparams["alpha"] < 1: self.gp_spl = gp_spl

        else: self.gp_spl = None

        self.spl_indexing_dict = spl_indexing_dict

        # Up-sample candidate genes

        self.upsample_cand = upsample_cand

        self.cand_gene_freq = Counter([])

        with open(str(project_config.KG_DIR  / f'ensembl_to_idx_dict_{project_config.CURR_KG}.pkl'), 'rb') as handle:

            ensembl_to_idx_dict = pickle.load(handle) # create ensembl to node_idx map

        idx_to_ensembl_dict = {v: k for k, v in ensembl_to_idx_dict.items()}

        self.cand_gene_freq = Counter([k for k in nid_to_spl_dict if k in idx_to_ensembl_dict]) # Upsample from all gene nodes in the KG

        self.n_cand_diseases = n_cand_diseases

        self.use_diseases = use_diseases

        self.hparams = hparams

        self.gene_similarity_dict = gene_similarity_dict

        self.gene_deg_dict = gene_deg_dict

        # Obtain a *transposed* `SparseTensor` instance.

        if not self.is_sparse_tensor:

            if num_nodes is None:

                num_nodes = int(edge_index.max()) + 1

                sample_num_nodes = int(sample_edge_index.max()) + 1

            value = torch.arange(edge_index.size(1)) if return_e_id else None

            sample_value = torch.arange(sample_edge_index.size(1)) if return_e_id else None

            self.adj_t = SparseTensor(row=edge_index[0], col=edge_index[1],

                                      value=value,

                                      sparse_sizes=(num_nodes, num_nodes)).t()

            self.adj_t_sample = SparseTensor(row=sample_edge_index[0], col=sample_edge_index[1],

                                      value=sample_value,

                                      sparse_sizes=(sample_num_nodes, sample_num_nodes)).t()

        else:

            adj_t = edge_index

            adj_t_sample = sample_edge_index

            if return_e_id:

                self.__val__ = adj_t.storage.value()

                value = torch.arange(adj_t.nnz())

                adj_t = adj_t.set_value(value, layout='coo')

                adj_t_sample = adj_t_sample.set_value(torch.arange(adj_t_sample.nnz()), layout='coo')

            self.adj_t = adj_t

            self.adj_t_sample = adj_t_sample

        self.adj_t.storage.rowptr()

        self.adj_t_sample.storage.rowptr()

        super(PatientNeighborSampler, self).__init__(

            self.patient_dataset, collate_fn=self.collate, **kwargs)

    def filter_edges(self, edge_index, e_id, source_nodes, target_nodes):

        '''

        Filter out the edges we're trying to predict in the current batch from the edge index

        NOTE: edge_index here is re-indexed

        '''

        reindex_source_nodes = torch.arange(source_nodes.size(0))

        reindex_target_nodes = torch.arange(start = source_nodes.size(0), end = source_nodes.size(0) + target_nodes.size(0))

        # get reverse edges to filter as well

        all_source_nodes = torch.cat([reindex_source_nodes, reindex_target_nodes])

        all_target_nodes = torch.cat([reindex_target_nodes, reindex_source_nodes])

        ind_to_edge_index, ind_to_nodes = get_indices_into_edge_index(edge_index, all_source_nodes, all_target_nodes) #get index into the original edge index (this returns e_ids)

        mask = torch.ones(edge_index.size(1), dtype=torch.bool)

        mask[ind_to_edge_index] = False

        return edge_index[:, mask], e_id[mask]

    def get_source_nodes(self, phenotype_node_idx, candidate_gene_node_idx, correct_genes_node_idx, disease_node_idx, candidate_disease_node_idx, sim_gene_node_idx): 

        # Get batch node indices based on patient phenotypes and genes

        if sim_gene_node_idx is not None:

            source_batch = torch.cat(phenotype_node_idx +  candidate_gene_node_idx +  correct_genes_node_idx + disease_node_idx + candidate_disease_node_idx + sim_gene_node_idx)

        else:

            source_batch = torch.cat(phenotype_node_idx +  candidate_gene_node_idx +  correct_genes_node_idx + disease_node_idx + candidate_disease_node_idx)

         # Randomly sample nodes in KG 

        if self.sparse_sample > 0:

            if self.relevant_node_idx == None:

                rand_idx = torch.randint(high=self.n_nodes, size=(self.sparse_sample,)) # NOTE that this can sample duplicates, but has the benefit of randomly sampling new nodes each epoch

            else:

                rand_idx = self.relevant_node_idx[torch.randint(high=self.relevant_node_idx.size(0), size=(self.sparse_sample,))]

            source_batch = torch.cat([source_batch, rand_idx])

            source_batch = torch.unique(source_batch)

            sparse_idx = torch.unique(rand_idx)

        else:

            source_batch = torch.unique(source_batch)

            sparse_idx = torch.Tensor([])

        return source_batch, sparse_idx

    def sample_target_nodes(self, source_batch):

        row, col, e_id = self.adj_t_sample.coo() 

        if self.sample_edges_from_train_patients:

            train_patient_nodes = torch.tensor(list(self.train_phenotype_counter.keys()) + list(self.train_gene_counter.keys())) 

            ind_with_train_patient_nodes = (col == train_patient_nodes.unsqueeze(-1)).nonzero(as_tuple=True)[1]

            subset_row = row[ind_with_train_patient_nodes]

            subset_col = col[ind_with_train_patient_nodes]

            try:

                # first try to find an edge that connects back to the training set patient data

                targets = random_walk(subset_row, subset_col, source_batch, walk_length=1, coalesced=False)[:, 1] #NOTE: only does self loops when no edges in the current partition of the dataset

                source_batch_1 = source_batch[~torch.eq(source_batch, targets)]

                targets_1 = targets[~torch.eq(source_batch, targets)]

                # if no edges are found, use all available edges in this split of the data

                source_batch_2 = source_batch[torch.eq(source_batch, targets)]

                targets_2 = random_walk(row, col, source_batch_2, walk_length=1, coalesced=False)[:, 1] #NOTE: only does self loops when no edges in the current partition of the dataset

                #concat the two together

                source_batch = torch.cat([source_batch_1, source_batch_2])

                targets = torch.cat([targets_1, targets_2])

            except:

                targets = random_walk(row, col, source_batch, walk_length=1, coalesced=False)[:, 1] #NOTE: only does self loops when no edges in the current partition of the dataset

        else:

            targets = random_walk(row, col, source_batch, walk_length=1, coalesced=False)[:, 1] #NOTE: only does self loops when no edges in the current partition of the dataset

        return source_batch, targets

    def add_patient_information(self, patient_ids, phenotype_node_idx, candidate_gene_node_idx, correct_genes_node_idx, sim_gene_node_idx, gene_sims, gene_degs, disease_node_idx, candidate_disease_node_idx, labels, disease_labels, patient_labels, additional_labels, adjs, batch_size, n_id, sparse_idx, target_batch): #candidate_disease_node_idx

        # Create Data Object & Add patient level information

        adjs = [HeterogeneousEdgeIndex(adj.edge_index, adj.e_id, self.all_edge_attr[adj.e_id], adj.size) for adj in adjs] 

        max_n_candidates = max([len(l) for l in candidate_gene_node_idx])

        data = Data(adjs = adjs, 

                batch_size = batch_size,

                patient_ids = patient_ids,

                n_id = n_id

                )

        if self.hparams['loss'] != 'patient_disease_NCA' and self.hparams['loss'] != 'patient_patient_NCA':

            if None in list(labels): data['one_hot_labels'] = None

            else: data['one_hot_labels'] = torch.LongTensor(label_binarize(labels, classes = list(range(max_n_candidates))))

        if self.use_diseases:

            data['disease_one_hot_labels'] = disease_labels 

        if self.hparams['loss'] == 'patient_patient_NCA':

            if patient_labels is None: data['patient_labels'] = None

            else: data['patient_labels'] = torch.stack(patient_labels)

        # Get candidate genes to phenotypes SPL

        if not self.gp_spl is None:

            if not self.spl_indexing_dict is None:

                patient_ids = np.vectorize(self.spl_indexing_dict.get)(patient_ids).astype(int)

            gene_to_phenotypes_spl = -torch.Tensor(self.gp_spl[patient_ids,:])

            # get gene idx to spl information

            cand_gene_idx_to_spl = [torch.LongTensor(np.vectorize(self.nid_to_spl_dict.get)(cand_genes)) for cand_genes in list(candidate_gene_node_idx)]

            # get SPLs for each patient's candidate genes

            batch_cand_gene_to_phenotypes_spl = [gene_spls[cand_genes] for cand_genes, gene_spls in zip(cand_gene_idx_to_spl, gene_to_phenotypes_spl)]

            # pad to same # of candidate genes

            data['batch_cand_gene_to_phenotypes_spl'] = pad_sequence(batch_cand_gene_to_phenotypes_spl, batch_first=True, padding_value=0)

            # get unique gene idx across all patients in the batch

            cand_gene_idx_flattened_unique = torch.unique(torch.cat(cand_gene_idx_to_spl)).flatten()

            # get SPLs for unique genes in the batch

            data['batch_concat_cand_gene_to_phenotypes_spl'] = gene_to_phenotypes_spl[:, cand_gene_idx_flattened_unique]

        else:

            data['batch_cand_gene_to_phenotypes_spl'] = None

            data['batch_concat_cand_gene_to_phenotypes_spl'] = None

        # Create mapping from KG node IDs to batch indices

        node2batch = {n+1: int(i+1) for i, n in enumerate(data.n_id.tolist())}

        node2batch[0] = 0

        # add phenotype / gene / disease names

        data['phenotype_names'] = [[(self.patient_dataset.node_idx_to_name(p.item()), self.patient_dataset.node_idx_to_degree(p.item())) for p in p_list] for p_list in phenotype_node_idx ]

        data['cand_gene_names'] = [[self.patient_dataset.node_idx_to_name(g.item()) for g in g_list] for g_list in candidate_gene_node_idx ]

        data['corr_gene_names'] = [[self.patient_dataset.node_idx_to_name(g.item()) for g in g_list] for g_list in correct_genes_node_idx  ]

        data['disease_names'] = [[self.patient_dataset.node_idx_to_name(d.item()) for d in d_list] for d_list in disease_node_idx ]

        if self.use_diseases:

            data['cand_disease_names'] = [[self.patient_dataset.node_idx_to_name(d.item()) for d in d_list] for d_list in candidate_disease_node_idx ]

        #reindex nodes to make room for padding

        phenotype_node_idx = [p + 1 for p in phenotype_node_idx]

        candidate_gene_node_idx = [g + 1 for g in candidate_gene_node_idx]

        correct_genes_node_idx = [g + 1 for g in correct_genes_node_idx]

        if self.use_diseases:

            disease_node_idx = [d + 1 for d in disease_node_idx]

            candidate_disease_node_idx = [d + 1 for d in candidate_disease_node_idx]

        if 'augment_genes' in self.hparams and self.hparams['augment_genes']:

            sim_gene_node_idx = [g + 1 for g in sim_gene_node_idx]

        # if there aren't any disease idx in the batch, we add filler

        if self.use_diseases:

            if all(len(t) == 0 for t in disease_node_idx):

                disease_node_idx = [torch.LongTensor([0]) for i in range(len(disease_node_idx))]

            if all(len(t) == 0 for t in candidate_disease_node_idx):

                candidate_disease_node_idx = [torch.LongTensor([0]) for i in range(len(candidate_disease_node_idx))]

        # add padding to patient phenotype and gene node idx

        data['batch_pheno_nid'] = pad_sequence(phenotype_node_idx, batch_first=True, padding_value=0) 

        if len(candidate_gene_node_idx[0]) > 0:

            data['batch_cand_gene_nid'] = pad_sequence(candidate_gene_node_idx, batch_first=True, padding_value=0) 

        data['batch_corr_gene_nid'] = pad_sequence(correct_genes_node_idx, batch_first=True, padding_value=0) 

        if self.use_diseases:

            data['batch_disease_nid'] = pad_sequence(disease_node_idx, batch_first=True, padding_value=0) 

            data['batch_cand_disease_nid'] = pad_sequence(candidate_disease_node_idx, batch_first=True, padding_value=0) 

        if 'augment_genes' in self.hparams and self.hparams['augment_genes']:

            data['batch_cand_gene_degs'] = pad_sequence(gene_degs, batch_first=True, padding_value=0) 

            data['batch_sim_gene_nid'] = pad_sequence(sim_gene_node_idx, batch_first=True, padding_value=0) 

            data['batch_sim_gene_sims'] = pad_sequence(gene_sims, batch_first=True, padding_value=0)

            # Normalize

            data['batch_sim_gene_sims'] = data['batch_sim_gene_sims'] / torch.sum(data['batch_sim_gene_sims'], dim=1, keepdim=True)

        else:

            if len(candidate_gene_node_idx[0]) > 0:

                data['batch_cand_gene_nid'] = pad_sequence(candidate_gene_node_idx, batch_first=True, padding_value=0) 

        # Convert KG node IDs to batch IDs

        # When performing inference (i.e., predict.py), use the original node IDs because the full KG is used in forward pass of node model

        if self.dataset_type != "predict":

            data['batch_pheno_nid']  = torch.LongTensor(np.vectorize(node2batch.get)(data['batch_pheno_nid']))

            if len(candidate_gene_node_idx[0]) > 0:

                data['batch_cand_gene_nid'] = torch.LongTensor(np.vectorize(node2batch.get)(data['batch_cand_gene_nid']))

            if len(correct_genes_node_idx[0]) > 0:

                data['batch_corr_gene_nid'] = torch.LongTensor(np.vectorize(node2batch.get)(data['batch_corr_gene_nid']))

            if self.use_diseases:

                data['batch_disease_nid'] = torch.LongTensor(np.vectorize(node2batch.get)(data['batch_disease_nid']))

                data['batch_cand_disease_nid'] = torch.LongTensor(np.vectorize(node2batch.get)(data['batch_cand_disease_nid']))

            if 'augment_genes' in self.hparams and self.hparams['augment_genes']:

                data['batch_sim_gene_nid'] = torch.LongTensor(np.vectorize(node2batch.get)(data['batch_sim_gene_nid']))

        return data

    def get_candidate_diseases(self, disease_node_idx, candidate_gene_node_idx):

        cand_diseases = self.patient_dataset.get_candidate_diseases(cand_type=self.hparams['candidate_disease_type'])

        if self.n_cand_diseases != -1: cand_diseases = cand_diseases[torch.randperm(len(cand_diseases))][0:self.n_cand_diseases] 

        if self.hparams['only_hard_distractors']: #add candidates to every patient

            candidate_disease_node_idx = tuple(torch.unique(torch.cat([corr_dis, cand_diseases ]), sorted=False) for corr_dis in disease_node_idx)

            candidate_disease_node_idx = tuple(torch.unique(dis[torch.randperm(len(dis))], sorted=False, return_inverse=False, return_counts=False) for dis in candidate_disease_node_idx)

        else: # split candidates across all patients in the batch

            all_correct_diseases = torch.cat(disease_node_idx)

            all_diseases = torch.unique(torch.cat([all_correct_diseases, cand_diseases]))

            all_diseases = all_diseases[torch.randperm(len(all_diseases))]

            candidate_disease_node_idx = np.array_split(all_diseases, len(candidate_gene_node_idx))

            candidate_disease_node_idx = tuple(candidate_disease_node_idx)

        max_n_dis_candidates = max([len(l) for l in candidate_disease_node_idx])

        if max_n_dis_candidates == 0: 

            max_n_dis_candidates = 1

            print('WARNING: there are no disease candidates')

        disease_ind = [(dis.unsqueeze(1) == corr_dis.unsqueeze(0)).nonzero(as_tuple=True)[0] if len(corr_dis) > 0 else torch.tensor(-1) for dis, corr_dis in zip(candidate_disease_node_idx, disease_node_idx)]

        disease_labels = torch.zeros((len(candidate_disease_node_idx), max_n_dis_candidates))

        for i, ind in enumerate(disease_ind): disease_labels[i,ind[ind != -1]] = 1

        return candidate_disease_node_idx, disease_labels

    def get_candidate_patients(self, patient_ids):

        # get patients with the same disease/gene

        similar_pat_ids = [self.patient_dataset.get_similar_patients(p_id, similarity_type=self.hparams['patient_similarity_type']) for p_id in patient_ids]

        # shuffle patients & subset to n_sim_pats so we have X similar patients per patient in the batch

        similar_pat_ids = [p[:self.hparams['n_similar_patients']] for p in similar_pat_ids] #[torch.randperm(len(p))]

        # Retrieve the patients for each of the sampled patient ids if they aren't already in the batch

        patient_ids = list(patient_ids) 

        similar_pats = [self.patient_dataset[self.patient_dataset.patient_id_to_index[p_id.item()]] for p_ids in similar_pat_ids for p_id in p_ids if p_id.item() not in patient_ids]

        return similar_pats

    def sample(self, batch, source_batch, target_batch):

        batch_size: int = len(batch)

        adjs = []

        n_id = batch

        for size in self.sizes:

            adj_t, n_id = self.adj_t.sample_adj(n_id, size, replace=False)

            e_id = adj_t.storage.value()

            size = adj_t.sparse_sizes()[::-1]

            if self.__val__ is not None:

                adj_t.set_value_(self.__val__[e_id], layout='coo')

            if self.is_sparse_tensor: #TODO: implement filter_edges if sparse tensor

                adjs.append(Adj(adj_t, e_id, size))

            else:

                row, col, _ = adj_t.coo()

                edge_index = torch.stack([col, row], dim=0)

                if self.do_filter_edges and self.dataset_type == 'train':

                    edge_index, e_id = self.filter_edges(edge_index, e_id, source_batch, target_batch)

                adjs.append(EdgeIndex(edge_index, e_id, size))

        adjs = [adjs[0]] if len(adjs) == 1 else adjs[::-1]

        return adjs, batch_size, n_id

    def get_similar_genes(self, patient_ids, candidate_gene_node_idx):

        k = self.hparams['n_sim_genes']

        gene_ids = []

        sims = []

        degs = []

        assert len(patient_ids) == len(candidate_gene_node_idx)

        for p, p_cand_genes in zip(patient_ids, candidate_gene_node_idx):

            p_genes = []

            p_sims = []

            p_degs = []

            for g in p_cand_genes:

                p_genes.append(torch.LongTensor([idx for idx, sim in list(self.gene_similarity_dict[int(g)])[:k]]))

                p_sims.append(torch.LongTensor([sim for idx, sim in list(self.gene_similarity_dict[int(g)])[:k]]))

                p_degs.append(self.gene_deg_dict[int(g)])

            gene_ids.append(torch.stack(p_genes))

            sims.append(torch.stack(p_sims))

            degs.append(torch.LongTensor(p_degs))

        assert len(gene_ids) == len(patient_ids)

        assert len(sims) == len(patient_ids)

        unique_genes = torch.unique(torch.cat(gene_ids).flatten()).unsqueeze(-1)

        return tuple(gene_ids), tuple(sims), tuple(degs), tuple(unique_genes)

    def collate(self, batch):

        t00 = time.time()

        phenotype_node_idx, candidate_gene_node_idx, correct_genes_node_idx, disease_node_idx, labels, additional_labels, patient_ids = zip(*batch)

        # Up-sample under-represented candidate genes

        t0 = time.time()

        if self.upsample_cand > 0:

            curr_cand_gene_freq = Counter(torch.cat(candidate_gene_node_idx).flatten().tolist())

            self.cand_gene_freq += curr_cand_gene_freq

            num_patients = len(candidate_gene_node_idx) * self.upsample_cand

            lowest_k_cand = self.cand_gene_freq.most_common()[:-num_patients-1:-1]

            lowest_k_cand = np.array_split([g[0] for g in lowest_k_cand], len(candidate_gene_node_idx))

            upsampled_candidate_gene_node_idx = []

            added_cand_gene = []

            for patient, cand_gene, corr_gene_idx in zip(candidate_gene_node_idx, lowest_k_cand, labels):

                # Remove correct genes from list of upsampled candidate genes

                corr_gene_nid = patient[corr_gene_idx]

                cand_gene = cand_gene[~np.isin(cand_gene, corr_gene_nid)].flatten()

                # Remove duplicates

                unique_cand_genes, new_cand_genes_freq = torch.unique(torch.tensor(patient.tolist() + list(cand_gene)), return_counts = True)

                unique_cand_genes = unique_cand_genes[new_cand_genes_freq == 1]

                cand_gene = cand_gene[np.isin(cand_gene, unique_cand_genes)]                

                # Add upsampled candidate genes

                added_cand_gene.extend(list(cand_gene))

                new_cand_list = torch.tensor(patient.tolist() + list(cand_gene))

                upsampled_candidate_gene_node_idx.append(new_cand_list)

            candidate_gene_node_idx = tuple(upsampled_candidate_gene_node_idx)

            self.cand_gene_freq += Counter(added_cand_gene)

        # Add similar patients to batch (for "patients like me" head)

        if self.hparams['add_similar_patients']:

            similar_pats = self.get_candidate_patients(patient_ids)

            # merge original batch with sampled patients

            phenotype_node_idx_sim, candidate_gene_node_idx_sim, correct_genes_node_idx_sim, disease_node_idx_sim, labels_sim, additional_labels_sim, patient_ids_sim = zip(*similar_pats)

            phenotype_node_idx = phenotype_node_idx + phenotype_node_idx_sim

            candidate_gene_node_idx = candidate_gene_node_idx + candidate_gene_node_idx_sim

            correct_genes_node_idx = correct_genes_node_idx + correct_genes_node_idx_sim

            disease_node_idx = disease_node_idx + disease_node_idx_sim

            labels = labels + labels_sim

            additional_labels = additional_labels + additional_labels_sim

            patient_ids = patient_ids + patient_ids_sim

        # get patient labels

        patient_labels = correct_genes_node_idx

        # Add candidate diseases to batch

        if self.hparams['add_cand_diseases']:

            candidate_disease_node_idx, disease_labels = self.get_candidate_diseases(disease_node_idx, candidate_gene_node_idx)

        else: 

            candidate_disease_node_idx = disease_node_idx

            disease_labels = torch.tensor([1] * len(candidate_disease_node_idx))

        if self.hparams['augment_genes']:

            sim_gene_node_idx, gene_sims, gene_degs, unique_sim_genes = self.get_similar_genes(patient_ids, candidate_gene_node_idx)

        else:

            unique_sim_genes = gene_degs = gene_sims = sim_gene_node_idx = None

        t1 = time.time()

        # get nodes from patients + randomly sampled nodes

        source_batch, sparse_idx = self.get_source_nodes(phenotype_node_idx, candidate_gene_node_idx, correct_genes_node_idx, disease_node_idx, candidate_disease_node_idx, unique_sim_genes)

        # sample nodes to form positive edges

        source_batch, target_batch = self.sample_target_nodes(source_batch) 

        batch = torch.cat([source_batch, target_batch], dim=0) 

        t2 = time.time()

        # get k hop adj graph

        adjs, batch_size, n_id = self.sample(batch, source_batch, target_batch)

        t3 = time.time()

        # add patient information to data object

        data = self.add_patient_information(patient_ids, phenotype_node_idx, candidate_gene_node_idx, correct_genes_node_idx, sim_gene_node_idx, gene_sims, gene_degs, disease_node_idx, candidate_disease_node_idx, labels, disease_labels, patient_labels, additional_labels, adjs, batch_size, n_id, sparse_idx, target_batch) #candidate_disease_node_idx

        t4 = time.time()

        if self.hparams['time']:

            print(f'It takes {t0-t00:0.4f}s to unzip batch, {t1-t0:0.4f}s to upsample candidate gene nodes, {t2-t1:0.4f}s to sample positive nodes, {t3-t2:0.4f}s to get k-hop adjs, and {t4-t3:0.4f}s to add patient information')

        return data        

    def __repr__(self):

        return '{}(sizes={})'.format(self.__class__.__name__, self.sizes)