import math
import warnings
from functools import reduce
from collections import defaultdict

import networkx as nx

from matplotlib import pyplot as plt
import torch
from torch_scatter import scatter_add, scatter_min

from torchdrug import core, utils
from torchdrug.data import Dictionary
from torchdrug.utils import pretty

plt.switch_backend("agg")


class Graph(core._MetaContainer):
    r"""
    Basic container for sparse graphs.

    To batch graphs with variadic sizes, use :meth:`data.Graph.pack <torchdrug.data.Graph.pack>`.
    This will return a PackedGraph object with the following block diagonal adjacency matrix.

    .. math::

        \begin{bmatrix}
            A_1    & \cdots & 0      \\
            \vdots & \ddots & \vdots \\
            0      & \cdots & A_n
        \end{bmatrix}

    where :math:`A_i` is the adjacency of :math:`i`-th graph.

    You may register dynamic attributes for each graph.
    The registered attributes will be automatically processed during packing.

    .. warning::

        This class doesn't enforce any order on the edges.

    Example::

        >>> graph = data.Graph(torch.randint(10, (30, 2)))
        >>> with graph.node():
        >>>     graph.my_node_attr = torch.rand(10, 5, 5)

    Parameters:
        edge_list (array_like, optional): list of edges of shape :math:`(|E|, 2)` or :math:`(|E|, 3)`.
            Each tuple is (node_in, node_out) or (node_in, node_out, relation).
        edge_weight (array_like, optional): edge weights of shape :math:`(|E|,)`
        num_node (int, optional): number of nodes.
            By default, it will be inferred from the largest id in `edge_list`
        num_relation (int, optional): number of relations
        node_feature (array_like, optional): node features of shape :math:`(|V|, ...)`
        edge_feature (array_like, optional): edge features of shape :math:`(|E|, ...)`
        graph_feature (array_like, optional): graph feature of any shape
    """

    _meta_types = {"node", "edge", "graph", "node reference", "edge reference", "graph reference"}

    def __init__(self, edge_list=None, edge_weight=None, num_node=None, num_relation=None,
                 node_feature=None, edge_feature=None, graph_feature=None, **kwargs):
        super(Graph, self).__init__(**kwargs)
        # edge_list: N * [h, t] or N * [h, t, r]
        edge_list, num_edge = self._standarize_edge_list(edge_list, num_relation)
        edge_weight = self._standarize_edge_weight(edge_weight, edge_list)

        num_node = self._standarize_num_node(num_node, edge_list)
        num_relation = self._standarize_num_relation(num_relation, edge_list)

        self._edge_list = edge_list
        self._edge_weight = edge_weight
        self.num_node = num_node
        self.num_edge = num_edge
        self.num_relation = num_relation

        if node_feature is not None:
            with self.node():
                self.node_feature = torch.as_tensor(node_feature, device=self.device)
        if edge_feature is not None:
            with self.edge():
                self.edge_feature = torch.as_tensor(edge_feature, device=self.device)
        if graph_feature is not None:
            with self.graph():
                self.graph_feature = torch.as_tensor(graph_feature, device=self.device)

    def node(self):
        """
        Context manager for node attributes.
        """
        return self.context("node")

    def edge(self):
        """
        Context manager for edge attributes.
        """
        return self.context("edge")

    def graph(self):
        """
        Context manager for graph attributes.
        """
        return self.context("graph")

    def node_reference(self):
        """
        Context manager for node references.
        """
        return self.context("node reference")

    def edge_reference(self):
        """
        Context manager for edge references.
        """
        return self.context("edge reference")

    def graph_reference(self):
        """
        Context manager for graph references.
        """
        return self.context("graph reference")

    def _check_attribute(self, key, value):
        for type in self._meta_contexts:
            if "reference" in type:
                if value.dtype != torch.long:
                    raise TypeError("Tensors used as reference must be long tensors")
            if type == "node":
                if len(value) != self.num_node:
                    raise ValueError("Expect node attribute `%s` to have shape (%d, *), but found %s" %
                                     (key, self.num_node, value.shape))
            elif type == "edge":
                if len(value) != self.num_edge:
                    raise ValueError("Expect edge attribute `%s` to have shape (%d, *), but found %s" %
                                     (key, self.num_edge, value.shape))
            elif type == "node reference":
                is_valid = (value >= -1) & (value < self.num_node)
                if not is_valid.all():
                    error_value = value[~is_valid]
                    raise ValueError("Expect node reference in [-1, %d), but found %d" %
                                     (self.num_node, error_value[0]))
            elif type == "edge reference":
                is_valid = (value >= -1) & (value < self.num_edge)
                if not is_valid.all():
                    error_value = value[~is_valid]
                    raise ValueError("Expect edge reference in [-1, %d), but found %d" %
                                     (self.num_edge, error_value[0]))
            elif type == "graph reference":
                is_valid = (value >= -1) & (value < self.batch_size)
                if not is_valid.all():
                    error_value = value[~is_valid]
                    raise ValueError("Expect graph reference in [-1, %d), but found %d" %
                                     (self.batch_size, error_value[0]))

    def __setattr__(self, key, value):
        if hasattr(self, "meta_dict"):
            self._check_attribute(key, value)
        super(Graph, self).__setattr__(key, value)

    def _standarize_edge_list(self, edge_list, num_relation):
        if edge_list is not None and len(edge_list):
            if isinstance(edge_list, torch.Tensor) and edge_list.dtype != torch.long:
                try:
                    edge_list = torch.LongTensor(edge_list)
                except TypeError:
                    raise TypeError("Can't convert `edge_list` to torch.long")
            else:
                edge_list = torch.as_tensor(edge_list, dtype=torch.long)
        else:
            num_element = 2 if num_relation is None else 3
            if isinstance(edge_list, torch.Tensor):
                device = edge_list.device
            else:
                device = "cpu"
            edge_list = torch.zeros(0, num_element, dtype=torch.long, device=device)
        if (edge_list < 0).any():
            raise ValueError("`edge_list` should only contain non-negative indexes")
        num_edge = torch.tensor(len(edge_list), device=edge_list.device)
        return edge_list, num_edge

    def _standarize_edge_weight(self, edge_weight, edge_list):
        if edge_weight is not None:
            edge_weight = torch.as_tensor(edge_weight, dtype=torch.float, device=edge_list.device)
            if len(edge_list) != len(edge_weight):
                raise ValueError("`edge_list` and `edge_weight` should be the same size, but found %d and %d"
                                 % (len(edge_list), len(edge_weight)))
        else:
            edge_weight = torch.ones(len(edge_list), device=edge_list.device)
        return edge_weight

    def _standarize_num_node(self, num_node, edge_list):
        if num_node is None:
            num_node = self._maybe_num_node(edge_list)
        num_node = torch.as_tensor(num_node, device=edge_list.device)
        if (edge_list[:, :2] >= num_node).any():
            raise ValueError("`num_node` is %d, but found node %d in `edge_list`" % (num_node, edge_list[:, :2].max()))
        return num_node

    def _standarize_num_relation(self, num_relation, edge_list):
        if num_relation is None and edge_list.shape[1] > 2:
            num_relation = self._maybe_num_relation(edge_list)
        if num_relation is not None:
            num_relation = torch.as_tensor(num_relation, device=edge_list.device)
            if edge_list.shape[1] <= 2:
                raise ValueError("`num_relation` is provided, but the number of dims of `edge_list` is less than 3.")
            elif (edge_list[:, 2] >= num_relation).any():
                raise ValueError("`num_relation` is %d, but found relation %d in `edge_list`" % (num_relation, edge_list[:, 2].max()))
        return num_relation

    def _maybe_num_node(self, edge_list):
        warnings.warn("_maybe_num_node() is used to determine the number of nodes. "
                      "This may underestimate the count if there are isolated nodes.")
        if len(edge_list):
            return edge_list[:, :2].max().item() + 1
        else:
            return 0

    def _maybe_num_relation(self, edge_list):
        warnings.warn("_maybe_num_relation() is used to determine the number of relations. "
                      "This may underestimate the count if there are unseen relations.")
        return edge_list[:, 2].max().item() + 1

    def _standarize_index(self, index, count):
        if isinstance(index, slice):
            start = index.start or 0
            if start < 0:
                start += count
            stop = index.stop or count
            if stop < 0:
                stop += count
            step = index.step or 1
            index = torch.arange(start, stop, step, device=self.device)
        else:
            index = torch.as_tensor(index, device=self.device)
            if index.ndim == 0:
                index = index.unsqueeze(0)
            if index.dtype == torch.bool:
                if index.shape != (count,):
                    raise IndexError("Invalid mask. Expect mask to have shape %s, but found %s" %
                                     ((int(count),), tuple(index.shape)))
                index = index.nonzero().squeeze(-1)
            else:
                index = index.long()
                max_index = -1 if len(index) == 0 else index.max().item()
                if max_index >= count:
                    raise IndexError("Invalid index. Expect index smaller than %d, but found %d" % (count, max_index))
        return index

    def _get_mapping(self, index, count):
        index = self._standarize_index(index, count)
        if (index.bincount() > 1).any():
            raise ValueError("Can't create mapping for duplicate index")
        mapping = -torch.ones(count + 1, dtype=torch.long, device=self.device)
        mapping[index] = torch.arange(len(index), device=self.device)
        return mapping

    def _get_repeat_pack_offsets(self, num_xs, repeats):
        new_num_xs = num_xs.repeat_interleave(repeats)
        cum_repeats_shifted = repeats.cumsum(0) - repeats
        new_num_xs[cum_repeats_shifted] -= num_xs
        offsets = new_num_xs.cumsum(0)
        return offsets

    @classmethod
    def from_dense(cls, adjacency, node_feature=None, edge_feature=None):
        """
        Create a sparse graph from a dense adjacency matrix.
        For zero entries in the adjacency matrix, their edge features will be ignored.

        Parameters:
            adjacency (array_like): adjacency matrix of shape :math:`(|V|, |V|)` or :math:`(|V|, |V|, |R|)`
            node_feature (array_like): node features of shape :math:`(|V|, ...)`
            edge_feature (array_like): edge features of shape :math:`(|V|, |V|, ...)` or :math:`(|V|, |V|, |R|, ...)`
        """
        adjacency = torch.as_tensor(adjacency)
        if adjacency.shape[0] != adjacency.shape[1]:
            raise ValueError("`adjacency` should be a square matrix, but found %d and %d" % adjacency.shape[:2])

        edge_list = adjacency.nonzero()
        edge_weight = adjacency[tuple(edge_list.t())]
        num_node = adjacency.shape[0]
        num_relation = adjacency.shape[2] if adjacency.ndim > 2 else None
        if edge_feature is not None:
            edge_feature = torch.as_tensor(edge_feature)
            edge_feature = edge_feature[tuple(edge_list.t())]

        return cls(edge_list, edge_weight, num_node, num_relation, node_feature, edge_feature)

    def connected_components(self):
        """
        Split this graph into connected components.

        Returns:
            (PackedGraph, LongTensor): connected components, number of connected components per graph
        """
        node_in, node_out = self.edge_list.t()[:2]
        range = torch.arange(self.num_node, device=self.device)
        node_in, node_out = torch.cat([node_in, node_out, range]), torch.cat([node_out, node_in, range])

        # find connected component
        # O(|E|d), d is the diameter of the graph
        min_neighbor = torch.arange(self.num_node, device=self.device)
        last = torch.zeros_like(min_neighbor)
        while not torch.equal(min_neighbor, last):
            last = min_neighbor
            min_neighbor = scatter_min(min_neighbor[node_out], node_in, dim_size=self.num_node)[0]
        anchor = torch.unique(min_neighbor)
        num_cc = self.node2graph[anchor].bincount(minlength=self.batch_size)
        return self.split(min_neighbor), num_cc

    def split(self, node2graph):
        """
        Split a graph into multiple disconnected graphs.

        Parameters:
            node2graph (array_like): ID of the graph each node belongs to

        Returns:
            PackedGraph
        """
        node2graph = torch.as_tensor(node2graph, dtype=torch.long, device=self.device)
        # coalesce arbitrary graph IDs to [0, n)
        _, node2graph = torch.unique(node2graph, return_inverse=True)
        num_graph = node2graph.max() + 1
        index = node2graph.argsort()
        mapping = torch.zeros_like(index)
        mapping[index] = torch.arange(len(index), device=self.device)

        node_in, node_out = self.edge_list.t()[:2]
        edge_mask = node2graph[node_in] == node2graph[node_out]
        edge2graph = node2graph[node_in]
        edge_index = edge2graph.argsort()
        edge_index = edge_index[edge_mask[edge_index]]

        prepend = torch.tensor([-1], device=self.device)
        is_first_node = torch.diff(node2graph[index], prepend=prepend) > 0
        graph_index = self.node2graph[index[is_first_node]]

        edge_list = self.edge_list.clone()
        edge_list[:, :2] = mapping[edge_list[:, :2]]

        num_nodes = node2graph.bincount(minlength=num_graph)
        num_edges = edge2graph[edge_index].bincount(minlength=num_graph)

        num_cum_nodes = num_nodes.cumsum(0)
        offsets = (num_cum_nodes - num_nodes)[edge2graph[edge_index]]

        data_dict, meta_dict = self.data_mask(index, edge_index, graph_index=graph_index, exclude="graph reference")

        return self.packed_type(edge_list[edge_index], edge_weight=self.edge_weight[edge_index], num_nodes=num_nodes,
                                num_edges=num_edges, num_relation=self.num_relation, offsets=offsets,
                                meta_dict=meta_dict, **data_dict)

    @classmethod
    def pack(cls, graphs):
        """
        Pack a list of graphs into a PackedGraph object.

        Parameters:
            graphs (list of Graph): list of graphs

        Returns:
            PackedGraph
        """
        edge_list = []
        edge_weight = []
        num_nodes = []
        num_edges = []
        num_relation = -1
        num_cum_node = 0
        num_cum_edge = 0
        num_graph = 0
        data_dict = defaultdict(list)
        meta_dict = graphs[0].meta_dict
        for graph in graphs:
            edge_list.append(graph.edge_list)
            edge_weight.append(graph.edge_weight)
            num_nodes.append(graph.num_node)
            num_edges.append(graph.num_edge)
            for k, v in graph.data_dict.items():
                for type in meta_dict[k]:
                    if type == "graph":
                        v = v.unsqueeze(0)
                    elif type == "node reference":
                        v = v + num_cum_node
                    elif type == "edge reference":
                        v = v + num_cum_edge
                    elif type == "graph reference":
                        v = v + num_graph
                data_dict[k].append(v)
            if num_relation == -1:
                num_relation = graph.num_relation
            elif num_relation != graph.num_relation:
                raise ValueError("Inconsistent `num_relation` in graphs. Expect %d but got %d."
                                 % (num_relation, graph.num_relation))
            num_cum_node += graph.num_node
            num_cum_edge += graph.num_edge
            num_graph += 1

        edge_list = torch.cat(edge_list)
        edge_weight = torch.cat(edge_weight)
        data_dict = {k: torch.cat(v) for k, v in data_dict.items()}

        return cls.packed_type(edge_list, edge_weight=edge_weight, num_nodes=num_nodes, num_edges=num_edges,
                               num_relation=num_relation, meta_dict=meta_dict, **data_dict)

    def repeat(self, count):
        """
        Repeat this graph.

        Parameters:
            count (int): number of repetitions

        Returns:
            PackedGraph
        """
        edge_list = self.edge_list.repeat(count, 1)
        edge_weight = self.edge_weight.repeat(count)
        num_nodes = [self.num_node] * count
        num_edges = [self.num_edge] * count
        num_relation = self.num_relation

        data_dict = {}
        for k, v in self.data_dict.items():
            if "graph" in self.meta_dict[k]:
                v = v.unsqueeze(0)
            shape = [1] * v.ndim
            shape[0] = count
            length = len(v)
            v = v.repeat(shape)
            for type in self.meta_dict[k]:
                if type == "node reference":
                    offsets = torch.arange(count, device=self.device) * self.num_node
                    v = v + offsets.repeat_interleave(length)
                elif type == "edge reference":
                    offsets = torch.arange(count, device=self.device) * self.num_edge
                    v = v + offsets.repeat_interleave(length)
                elif type == "graph reference":
                    offsets = torch.arange(count, device=self.device)
                    v = v + offsets.repeat_interleave(length)
            data_dict[k] = v

        return self.packed_type(edge_list, edge_weight=edge_weight, num_nodes=num_nodes, num_edges=num_edges,
                                num_relation=num_relation, meta_dict=self.meta_dict, **data_dict)

    def get_edge(self, edge):
        """
        Get the weight of of an edge.

        Parameters:
            edge (array_like): index of shape :math:`(2,)` or :math:`(3,)`

        Returns:
            Tensor: weight of the edge
        """
        if len(edge) != self.edge_list.shape[1]:
            raise ValueError("Incorrect edge index. Expect %d axes but got %d axes"
                             % (self.edge_list.shape[1], len(edge)))

        edge_index, num_match = self.match(edge)
        return self.edge_weight[edge_index].sum()

    def match(self, pattern):
        """
        Return all matched indexes for each pattern. Support patterns with ``-1`` as the wildcard.

        Parameters:
            pattern (array_like): index of shape :math:`(N, 2)` or :math:`(N, 3)`

        Returns:
            (LongTensor, LongTensor): matched indexes, number of matches per edge

        Examples::

            >>> graph = data.Graph([[0, 1], [1, 0], [1, 2], [2, 1], [2, 0], [0, 2]])
            >>> index, num_match = graph.match([[0, -1], [1, 2]])
            >>> assert (index == torch.tensor([0, 5, 2])).all()
            >>> assert (num_match == torch.tensor([2, 1])).all()

        """
        if len(pattern) == 0:
            index = num_match = torch.zeros(0, dtype=torch.long, device=self.device)
            return index, num_match

        if not hasattr(self, "edge_inverted_index"):
            self.edge_inverted_index = {}
        pattern = torch.as_tensor(pattern, dtype=torch.long, device=self.device)
        if pattern.ndim == 1:
            pattern = pattern.unsqueeze(0)
        mask = pattern != -1
        scale = 2 ** torch.arange(pattern.shape[-1], device=self.device)
        query_type = (mask * scale).sum(dim=-1)
        query_index = query_type.argsort()
        num_query = query_type.unique(return_counts=True)[1]
        query_ends = num_query.cumsum(0)
        query_starts = query_ends - num_query
        mask_set = mask[query_index[query_starts]].tolist()

        type_ranges = []
        type_orders = []
        # get matched range for each query type
        for i, mask in enumerate(mask_set):
            query_type = tuple(mask)
            type_index = query_index[query_starts[i]: query_ends[i]]
            type_edge = pattern[type_index][:, mask]
            if query_type not in self.edge_inverted_index:
                self.edge_inverted_index[query_type] = self._build_edge_inverted_index(mask)
            inverted_range, order = self.edge_inverted_index[query_type]
            ranges = inverted_range.get(type_edge, default=0)
            type_ranges.append(ranges)
            type_orders.append(order)
        ranges = torch.cat(type_ranges)
        orders = torch.stack(type_orders)
        types = torch.arange(len(mask_set), device=self.device)
        types = types.repeat_interleave(num_query)

        # reorder matched ranges according to the query order
        ranges = scatter_add(ranges, query_index, dim=0, dim_size=len(pattern))
        types = scatter_add(types, query_index, dim_size=len(pattern))
        # convert range to indexes
        starts, ends = ranges.t()
        num_match = ends - starts
        offsets = num_match.cumsum(0) - num_match
        types = types.repeat_interleave(num_match)
        ranges = torch.arange(num_match.sum(), device=self.device)
        ranges = ranges + (starts - offsets).repeat_interleave(num_match)
        index = orders[types, ranges]

        return index, num_match

    def _build_edge_inverted_index(self, mask):
        keys = self.edge_list[:, mask]
        base = torch.tensor(self.shape, device=self.device)
        base = base[mask]
        max = reduce(int.__mul__, base.tolist())
        if max > torch.iinfo(torch.int64).max:
            raise ValueError("Fail to build an inverted index table based on sorting. "
                             "The graph is too large.")
        scale = base.cumprod(0)
        scale = torch.div(scale[-1], scale, rounding_mode="floor")
        key = (keys * scale).sum(dim=-1)
        order = key.argsort()
        num_keys = key.unique(return_counts=True)[1]
        ends = num_keys.cumsum(0)
        starts = ends - num_keys
        ranges = torch.stack([starts, ends], dim=-1)
        keys_set = keys[order[starts]]
        inverted_range = Dictionary(keys_set, ranges)
        return inverted_range, order

    def __getitem__(self, index):
        # why do we check tuple?
        # case 1: x[0, 1] is parsed as (0, 1)
        # case 2: x[[0, 1]] is parsed as [0, 1]
        if not isinstance(index, tuple):
            index = (index,)
        index = list(index)

        while len(index) < 2:
            index.append(slice(None))
        if len(index) > 2:
            raise ValueError("Graph has only 2 axis, but %d axis is indexed" % len(index))

        if all([isinstance(axis_index, int) for axis_index in index]):
            return self.get_edge(index)

        edge_list = self.edge_list.clone()
        for i, axis_index in enumerate(index):
            axis_index = self._standarize_index(axis_index, self.num_node)
            mapping = -torch.ones(self.num_node, dtype=torch.long, device=self.device)
            mapping[axis_index] = axis_index
            edge_list[:, i] = mapping[edge_list[:, i]]
        edge_index = (edge_list >= 0).all(dim=-1)

        return self.edge_mask(edge_index)

    def __len__(self):
        return 1

    @property
    def batch_size(self):
        """Batch size."""
        return 1

    def subgraph(self, index):
        """
        Return a subgraph based on the specified nodes.
        Equivalent to :meth:`node_mask(index, compact=True) <node_mask>`.

        Parameters:
            index (array_like): node index

        Returns:
            Graph

        See also:
            :meth:`Graph.node_mask`
        """
        return self.node_mask(index, compact=True)

    def data_mask(self, node_index=None, edge_index=None, graph_index=None, include=None, exclude=None):
        data_dict, meta_dict = self.data_by_meta(include, exclude)
        node_mapping = None
        edge_mapping = None
        graph_mapping = None
        for k, v in data_dict.items():
            for type in meta_dict[k]:
                if type == "node" and node_index is not None:
                    v = v[node_index]
                elif type == "edge" and edge_index is not None:
                    v = v[edge_index]
                elif type == "graph" and graph_index is not None:
                    v = v.unsqueeze(0)[graph_index]
                elif type == "node reference" and node_index is not None:
                    if node_mapping is None:
                        node_mapping = self._get_mapping(node_index, self.num_node)
                    v = node_mapping[v]
                elif type == "edge reference" and edge_index is not None:
                    if edge_mapping is None:
                        edge_mapping = self._get_mapping(edge_index, self.num_edge)
                    v = edge_mapping[v]
                elif type == "graph reference" and graph_index is not None:
                    if graph_mapping is None:
                        graph_mapping = self._get_mapping(graph_index, self.batch_size)
                    v = graph_mapping[v]
            data_dict[k] = v

        return data_dict, meta_dict

    def node_mask(self, index, compact=False):
        """
        Return a masked graph based on the specified nodes.

        This function can also be used to re-order the nodes.

        Parameters:
            index (array_like): node index
            compact (bool, optional): compact node ids or not

        Returns:
            Graph

        Examples::

            >>> graph = data.Graph.from_dense(torch.eye(3))
            >>> assert graph.node_mask([1, 2]).adjacency.shape == (3, 3)
            >>> assert graph.node_mask([1, 2], compact=True).adjacency.shape == (2, 2)

        """
        index = self._standarize_index(index, self.num_node)
        mapping = -torch.ones(self.num_node, dtype=torch.long, device=self.device)
        if compact:
            mapping[index] = torch.arange(len(index), device=self.device)
            num_node = len(index)
        else:
            mapping[index] = index
            num_node = self.num_node

        edge_list = self.edge_list.clone()
        edge_list[:, :2] = mapping[edge_list[:, :2]]
        edge_index = (edge_list[:, :2] >= 0).all(dim=-1)

        if compact:
            data_dict, meta_dict = self.data_mask(index, edge_index)
        else:
            data_dict, meta_dict = self.data_mask(edge_index=edge_index)

        return type(self)(edge_list[edge_index], edge_weight=self.edge_weight[edge_index], num_node=num_node,
                          num_relation=self.num_relation, meta_dict=meta_dict, **data_dict)

    def compact(self):
        """
        Remove isolated nodes and compact node ids.

        Returns:
            Graph
        """
        index = self.degree_out + self.degree_in > 0
        return self.subgraph(index)

    def edge_mask(self, index):
        """
        Return a masked graph based on the specified edges.

        This function can also be used to re-order the edges.

        Parameters:
            index (array_like): edge index

        Returns:
            Graph
        """
        index = self._standarize_index(index, self.num_edge)
        data_dict, meta_dict = self.data_mask(edge_index=index)

        return type(self)(self.edge_list[index], edge_weight=self.edge_weight[index], num_node=self.num_node,
                          num_relation=self.num_relation, meta_dict=meta_dict, **data_dict)

    def line_graph(self):
        """
        Construct a line graph of this graph.
        The node feature of the line graph is inherited from the edge feature of the original graph.

        In the line graph, each node corresponds to an edge in the original graph.
        For a pair of edges (a, b) and (b, c) that share the same intermediate node in the original graph,
        there is a directed edge (a, b) -> (b, c) in the line graph.

        Returns:
            Graph
        """
        node_in, node_out = self.edge_list.t()[:2]
        edge_index = torch.arange(self.num_edge, device=self.device)
        edge_in = edge_index[node_out.argsort()]
        edge_out = edge_index[node_in.argsort()]

        degree_in = node_in.bincount(minlength=self.num_node)
        degree_out = node_out.bincount(minlength=self.num_node)
        size = degree_out * degree_in
        starts = (size.cumsum(0) - size).repeat_interleave(size)
        range = torch.arange(size.sum(), device=self.device)
        # each node u has degree_out[u] * degree_in[u] local edges
        local_index = range - starts
        local_inner_size = degree_in.repeat_interleave(size)
        edge_in_offset = (degree_out.cumsum(0) - degree_out).repeat_interleave(size)
        edge_out_offset = (degree_in.cumsum(0) - degree_in).repeat_interleave(size)
        edge_in_index = torch.div(local_index, local_inner_size, rounding_mode="floor") + edge_in_offset
        edge_out_index = local_index % local_inner_size + edge_out_offset

        edge_in = edge_in[edge_in_index]
        edge_out = edge_out[edge_out_index]
        edge_list = torch.stack([edge_in, edge_out], dim=-1)
        node_feature = getattr(self, "edge_feature", None)
        num_node = self.num_edge
        num_edge = size.sum()

        return Graph(edge_list, num_node=num_node, num_edge=num_edge, node_feature=node_feature)

    def full(self):
        """
        Return a fully connected graph over the nodes.

        Returns:
            Graph
        """
        index = torch.arange(self.num_node, device=self.device)
        if self.num_relation:
            edge_list = torch.meshgrid(index, index, torch.arange(self.num_relation, device=self.device))
        else:
            edge_list = torch.meshgrid(index, index)
        edge_list = torch.stack(edge_list).flatten(1)
        edge_weight = torch.ones(len(edge_list))

        data_dict, meta_dict = self.data_by_meta(exclude="edge")

        return type(self)(edge_list, edge_weight=edge_weight, num_node=self.num_node, num_relation=self.num_relation,
                          meta_dict=meta_dict, **data_dict)

    def directed(self, order=None):
        """
        Mask the edges to create a directed graph.
        Edges that go from a node index to a larger or equal node index will be kept.

        Parameters:
            order (Tensor, optional): topological order of the nodes
        """
        node_in, node_out = self.edge_list.t()[:2]
        if order is not None:
            edge_index = order[node_in] <= order[node_out]
        else:
            edge_index = node_in <= node_out

        return self.edge_mask(edge_index)

    def undirected(self, add_inverse=False):
        """
        Flip all the edges to create an undirected graph.

        For knowledge graphs, the flipped edges can either have the original relation or an inverse relation.
        The inverse relation for relation :math:`r` is defined as :math:`|R| + r`.

        Parameters:
            add_inverse (bool, optional): whether to use inverse relations for flipped edges
        """
        edge_list = self.edge_list.clone()
        edge_list[:, :2] = edge_list[:, :2].flip(1)
        num_relation = self.num_relation
        if num_relation and add_inverse:
            edge_list[:, 2] += num_relation
            num_relation = num_relation * 2
        edge_list = torch.stack([self.edge_list, edge_list], dim=1).flatten(0, 1)

        index = torch.arange(self.num_edge, device=self.device).unsqueeze(-1).expand(-1, 2).flatten()
        data_dict, meta_dict = self.data_mask(edge_index=index)

        return type(self)(edge_list, edge_weight=self.edge_weight[index], num_node=self.num_node,
                          num_relation=num_relation, meta_dict=meta_dict, **data_dict)

    @utils.cached_property
    def adjacency(self):
        """
        Adjacency matrix of this graph.

        If :attr:`num_relation` is specified, a sparse tensor of shape :math:`(|V|, |V|, num\_relation)` will be
        returned.
        Otherwise, a sparse tensor of shape :math:`(|V|, |V|)` will be returned.
        """
        return utils.sparse_coo_tensor(self.edge_list.t(), self.edge_weight, self.shape)

    _tensor_names = ["edge_list", "edge_weight", "num_node", "num_relation", "edge_feature"]

    def to_tensors(self):
        edge_feature = getattr(self, "edge_feature", torch.tensor(0, device=self.device))
        return self.edge_list, self.edge_weight, self.num_node, self.num_relation, edge_feature

    @classmethod
    def from_tensors(cls, tensors):
        edge_list, edge_weight, num_node, num_relation, edge_feature = tensors
        if edge_feature.ndim == 0:
            edge_feature = None
        return cls(edge_list, edge_weight, num_node, num_relation, edge_feature=edge_feature)

    @property
    def node2graph(self):
        """Node id to graph id mapping."""
        return torch.zeros(self.num_node, dtype=torch.long, device=self.device)

    @property
    def edge2graph(self):
        """Edge id to graph id mapping."""
        return torch.zeros(self.num_edge, dtype=torch.long, device=self.device)

    @utils.cached_property
    def degree_out(self):
        """
        Weighted number of edges containing each node as output.

        Note this is the **in-degree** in graph theory.
        """
        return scatter_add(self.edge_weight, self.edge_list[:, 1], dim_size=self.num_node)

    @utils.cached_property
    def degree_in(self):
        """
        Weighted number of edges containing each node as input.

        Note this is the **out-degree** in graph theory.
        """
        return scatter_add(self.edge_weight, self.edge_list[:, 0], dim_size=self.num_node)

    @property
    def edge_list(self):
        """List of edges."""
        return self._edge_list

    @property
    def edge_weight(self):
        """Edge weights."""
        return self._edge_weight

    @property
    def device(self):
        """Device."""
        return self.edge_list.device

    @property
    def requires_grad(self):
        return self.edge_weight.requires_grad

    @property
    def grad(self):
        return self.edge_weight.grad

    @property
    def data(self):
        return self

    def requires_grad_(self):
        self.edge_weight.requires_grad_()
        return self

    def size(self, dim=None):
        if self.num_relation:
            size = torch.Size((self.num_node, self.num_node, self.num_relation))
        else:
            size = torch.Size((self.num_node, self.num_node))
        if dim is None:
            return size
        return size[dim]

    @property
    def shape(self):
        return self.size()

    def copy_(self, src):
        """
        Copy data from ``src`` into ``self`` and return ``self``.

        The ``src`` graph must have the same set of attributes as ``self``.
        """
        self.edge_list.copy_(src.edge_list)
        self.edge_weight.copy_(src.edge_weight)
        self.num_node.copy_(src.num_node)
        self.num_edge.copy_(src.num_edge)
        if self.num_relation is not None:
            self.num_relation.copy_(src.num_relation)

        keys = set(self.data_dict.keys())
        src_keys = set(src.data_dict.keys())
        if keys != src_keys:
            raise RuntimeError("Attributes mismatch. Trying to assign attributes %s, "
                               "but current graph has attributes %s" % (src_keys, keys))
        for k, v in self.data_dict.items():
            v.copy_(src.data_dict[k])

        return self

    def detach(self):
        """
        Detach this graph.
        """
        return type(self)(self.edge_list.detach(), edge_weight=self.edge_weight.detach(),
                          num_node=self.num_node, num_relation=self.num_relation,
                          meta_dict=self.meta_dict, **utils.detach(self.data_dict))

    def clone(self):
        """
        Clone this graph.
        """
        return type(self)(self.edge_list.clone(), edge_weight=self.edge_weight.clone(),
                          num_node=self.num_node, num_relation=self.num_relation,
                          meta_dict=self.meta_dict, **utils.clone(self.data_dict))

    def cuda(self, *args, **kwargs):
        """
        Return a copy of this graph in CUDA memory.

        This is a non-op if the graph is already on the correct device.
        """
        edge_list = self.edge_list.cuda(*args, **kwargs)

        if edge_list is self.edge_list:
            return self
        else:
            return type(self)(edge_list, edge_weight=self.edge_weight,
                              num_node=self.num_node, num_relation=self.num_relation,
                              meta_dict=self.meta_dict, **utils.cuda(self.data_dict, *args, **kwargs))

    def cpu(self):
        """
        Return a copy of this graph in CPU memory.

        This is a non-op if the graph is already in CPU memory.
        """
        edge_list = self.edge_list.cpu()

        if edge_list is self.edge_list:
            return self
        else:
            return type(self)(edge_list, edge_weight=self.edge_weight, num_node=self.num_node,
                              num_relation=self.num_relation, meta_dict=self.meta_dict, **utils.cpu(self.data_dict))

    def to(self, device, *args, **kwargs):
        """
        Return a copy of this graph on the given device.
        """
        device = torch.device(device)
        if device.type == "cpu":
            return self.cpu(*args, **kwargs)
        else:
            return self.cuda(device, *args, **kwargs)

    def __repr__(self):
        fields = ["num_node=%d" % self.num_node, "num_edge=%d" % self.num_edge]
        if self.num_relation is not None:
            fields.append("num_relation=%d" % self.num_relation)
        if self.device.type != "cpu":
            fields.append("device='%s'" % self.device)
        return "%s(%s)" % (self.__class__.__name__, ", ".join(fields))

    def visualize(self, title=None, save_file=None, figure_size=(3, 3), ax=None, layout="spring"):
        """
        Visualize this graph with matplotlib.

        Parameters:
            title (str, optional): title for this graph
            save_file (str, optional): ``png`` or ``pdf`` file to save visualization.
                If not provided, show the figure in window.
            figure_size (tuple of int, optional): width and height of the figure
            ax (matplotlib.axes.Axes, optional): axis to plot the figure
            layout (str, optional): graph layout

        See also:
            `NetworkX graph layout`_

            .. _NetworkX graph layout:
                https://networkx.github.io/documentation/stable/reference/drawing.html#module-networkx.drawing.layout
        """
        is_root = ax is None
        if ax is None:
            fig = plt.figure(figsize=figure_size)
            if title is not None:
                ax = plt.gca()
            else:
                ax = fig.add_axes([0, 0, 1, 1])
        if title is not None:
            ax.set_title(title)

        edge_list = self.edge_list[:, :2].tolist()
        G = nx.DiGraph(edge_list)
        G.add_nodes_from(range(self.num_node))
        if hasattr(nx, "%s_layout" % layout):
            func = getattr(nx, "%s_layout" % layout)
        else:
            raise ValueError("Unknown networkx layout `%s`" % layout)
        if layout == "spring" or layout == "random":
            pos = func(G, seed=0)
        else:
            pos = func(G)
        nx.draw_networkx(G, pos, ax=ax)
        if self.num_relation:
            edge_labels = self.edge_list[:, 2].tolist()
            edge_labels = {tuple(e): l for e, l in zip(edge_list, edge_labels)}
            nx.draw_networkx_edge_labels(G, pos, edge_labels, ax=ax)
        ax.set_frame_on(False)

        if is_root:
            if save_file:
                fig.savefig(save_file)
            else:
                fig.show()

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        return NotImplemented

    def __getstate__(self):
        state = {}
        cls = self.__class__
        for k, v in self.__dict__.items():
            # do not pickle property / cached property
            if hasattr(cls, k) and isinstance(getattr(cls, k), property):
                continue
            state[k] = v
        return state


class PackedGraph(Graph):
    """
    Container for sparse graphs with variadic sizes.

    To create a PackedGraph from Graph objects

        >>> batch = data.Graph.pack(graphs)

    To retrieve Graph objects from a PackedGraph

        >>> graphs = batch.unpack()

    .. warning::
        
        Edges of the same graph are guaranteed to be consecutive in the edge list.
        However, this class doesn't enforce any order on the edges.

    Parameters:
        edge_list (array_like, optional): list of edges of shape :math:`(|E|, 2)` or :math:`(|E|, 3)`.
            Each tuple is (node_in, node_out) or (node_in, node_out, relation).
        edge_weight (array_like, optional): edge weights of shape :math:`(|E|,)`
        num_nodes (array_like, optional): number of nodes in each graph
            By default, it will be inferred from the largest id in `edge_list`
        num_edges (array_like, optional): number of edges in each graph
        num_relation (int, optional): number of relations
        node_feature (array_like, optional): node features of shape :math:`(|V|, ...)`
        edge_feature (array_like, optional): edge features of shape :math:`(|E|, ...)`
        offsets (array_like, optional): node id offsets of shape :math:`(|E|,)`.
            If not provided, nodes in `edge_list` should be relative index, i.e., the index in each graph.
            If provided, nodes in `edge_list` should be absolute index, i.e., the index in the packed graph.
    """

    unpacked_type = Graph

    def __init__(self, edge_list=None, edge_weight=None, num_nodes=None, num_edges=None, num_relation=None,
                 offsets=None, **kwargs):
        edge_list, num_nodes, num_edges, num_cum_nodes, num_cum_edges, offsets = \
            self._get_cumulative(edge_list, num_nodes, num_edges, offsets)

        if offsets is None:
            offsets = self._get_offsets(num_nodes, num_edges, num_cum_nodes)
            edge_list = edge_list.clone()
            edge_list[:, :2] += offsets.unsqueeze(-1)

        num_node = num_nodes.sum()
        if (edge_list[:, :2] >= num_node).any():
            raise ValueError("Sum of `num_nodes` is %d, but found %d in `edge_list`" %
                             (num_node, edge_list[:, :2].max()))

        self._offsets = offsets
        self.num_nodes = num_nodes
        self.num_edges = num_edges
        self.num_cum_nodes = num_cum_nodes
        self.num_cum_edges = num_cum_edges

        super(PackedGraph, self).__init__(edge_list, edge_weight=edge_weight, num_node=num_node,
                                          num_relation=num_relation, **kwargs)

    def _get_offsets(self, num_nodes=None, num_edges=None, num_cum_nodes=None, num_cum_edges=None):
        if num_nodes is None:
            prepend = torch.tensor([0], device=self.device)
            num_nodes = torch.diff(num_cum_nodes, prepend=prepend)
        if num_edges is None:
            prepend = torch.tensor([0], device=self.device)
            num_edges = torch.diff(num_cum_edges, prepend=prepend)
        if num_cum_nodes is None:
            num_cum_nodes = num_nodes.cumsum(0)
        return (num_cum_nodes - num_nodes).repeat_interleave(num_edges)

    def merge(self, graph2graph):
        """
        Merge multiple graphs into a single graph.

        Parameters:
            graph2graph (array_like): ID of the new graph each graph belongs to
        """
        graph2graph = torch.as_tensor(graph2graph, dtype=torch.long, device=self.device)
        # coalesce arbitrary graph IDs to [0, n)
        _, graph2graph = torch.unique(graph2graph, return_inverse=True)

        graph_key = graph2graph * self.batch_size + torch.arange(self.batch_size, device=self.device)
        graph_index = graph_key.argsort()
        graph = self.subbatch(graph_index)
        graph2graph = graph2graph[graph_index]

        num_graph = graph2graph[-1] + 1
        num_nodes = scatter_add(graph.num_nodes, graph2graph, dim_size=num_graph)
        num_edges = scatter_add(graph.num_edges, graph2graph, dim_size=num_graph)
        offsets = self._get_offsets(num_nodes, num_edges)

        data_dict, meta_dict = graph.data_mask(exclude="graph")

        return type(self)(graph.edge_list, edge_weight=graph.edge_weight, num_nodes=num_nodes,
                          num_edges=num_edges, num_relation=graph.num_relation, offsets=offsets,
                          meta_dict=meta_dict, **data_dict)

    def unpack(self):
        """
        Unpack this packed graph into a list of graphs.

        Returns:
            list of Graph
        """
        graphs = []
        for i in range(self.batch_size):
            graphs.append(self.get_item(i))
        return graphs

    def __iter__(self):
        self._iter_index = 0
        return self

    def __next__(self):
        if self._iter_index < self.batch_size:
            item = self[self._iter_index]
            self._iter_index += 1
            return item
        raise StopIteration

    def _check_attribute(self, key, value):
        for type in self._meta_contexts:
            if "reference" in type:
                if value.dtype != torch.long:
                    raise TypeError("Tensors used as reference must be long tensors")
            if type == "node":
                if len(value) != self.num_node:
                    raise ValueError("Expect node attribute `%s` to have shape (%d, *), but found %s" %
                                     (key, self.num_node, value.shape))
            elif type == "edge":
                if len(value) != self.num_edge:
                    raise ValueError("Expect edge attribute `%s` to have shape (%d, *), but found %s" %
                                     (key, self.num_edge, value.shape))
            elif type == "graph":
                if len(value) != self.batch_size:
                    raise ValueError("Expect graph attribute `%s` to have shape (%d, *), but found %s" %
                                     (key, self.batch_size, value.shape))
            elif type == "node reference":
                is_valid = (value >= -1) & (value < self.num_node)
                if not is_valid.all():
                    error_value = value[~is_valid]
                    raise ValueError("Expect node reference in [-1, %d), but found %d" %
                                     (self.num_node, error_value[0]))
            elif type == "edge reference":
                is_valid = (value >= -1) & (value < self.num_edge)
                if not is_valid.all():
                    error_value = value[~is_valid]
                    raise ValueError("Expect edge reference in [-1, %d), but found %d" %
                                     (self.num_edge, error_value[0]))
            elif type == "graph reference":
                is_valid = (value >= -1) & (value < self.batch_size)
                if not is_valid.all():
                    error_value = value[~is_valid]
                    raise ValueError("Expect graph reference in [-1, %d), but found %d" %
                                     (self.batch_size, error_value[0]))

    def unpack_data(self, data, type="auto"):
        """
        Unpack node or edge data according to the packed graph.

        Parameters:
            data (Tensor): data to unpack
            type (str, optional): data type. Can be ``auto``, ``node``, or ``edge``.

        Returns:
            list of Tensor
        """
        if type == "auto":
            if self.num_node == self.num_edge:
                raise ValueError("Ambiguous type. Please specify either `node` or `edge`")
            if len(data) == self.num_node:
                type = "node"
            elif len(data) == self.num_edge:
                type = "edge"
            else:
                raise ValueError("Graph has %d nodes and %d edges, but data has %d entries" %
                                 (self.num_node, self.num_edge, len(data)))
        data_list = []
        if type == "node":
            for i in range(self.batch_size):
                data_list.append(data[self.num_cum_nodes[i] - self.num_nodes[i]: self.num_cum_nodes[i]])
        elif type == "edge":
            for i in range(self.batch_size):
                data_list.append(data[self.num_cum_edges[i] - self.num_edges[i]: self.num_cum_edges[i]])

        return data_list

    def repeat(self, count):
        """
        Repeat this packed graph. This function behaves similarly to `torch.Tensor.repeat`_.

        .. _torch.Tensor.repeat:
            https://pytorch.org/docs/stable/generated/torch.Tensor.repeat.html

        Parameters:
            count (int): number of repetitions

        Returns:
            PackedGraph
        """
        num_nodes = self.num_nodes.repeat(count)
        num_edges = self.num_edges.repeat(count)
        offsets = self._get_offsets(num_nodes, num_edges)
        edge_list = self.edge_list.repeat(count, 1)
        edge_list[:, :2] += (offsets - self._offsets.repeat(count)).unsqueeze(-1)

        data_dict = {}
        for k, v in self.data_dict.items():
            shape = [1] * v.ndim
            shape[0] = count
            length = len(v)
            v = v.repeat(shape)
            for _type in self.meta_dict[k]:
                if _type == "node reference":
                    pack_offsets = torch.arange(count, device=self.device) * self.num_node
                    v = v + pack_offsets.repeat_interleave(length)
                elif _type == "edge reference":
                    pack_offsets = torch.arange(count, device=self.device) * self.num_edge
                    v = v + pack_offsets.repeat_interleave(length)
                elif _type == "graph reference":
                    pack_offsets = torch.arange(count, device=self.device) * self.batch_size
                    v = v + pack_offsets.repeat_interleave(length)
            data_dict[k] = v

        return type(self)(edge_list, edge_weight=self.edge_weight.repeat(count),
                          num_nodes=num_nodes, num_edges=num_edges, num_relation=self.num_relation,
                          offsets=offsets, meta_dict=self.meta_dict, **data_dict)

    def repeat_interleave(self, repeats):
        """
        Repeat this packed graph. This function behaves similarly to `torch.repeat_interleave`_.

        .. _torch.repeat_interleave:
            https://pytorch.org/docs/stable/generated/torch.repeat_interleave.html

        Parameters:
            repeats (Tensor or int): number of repetitions for each graph

        Returns:
            PackedGraph
        """
        repeats = torch.as_tensor(repeats, dtype=torch.long, device=self.device)
        if repeats.numel() == 1:
            repeats = repeats * torch.ones(self.batch_size, dtype=torch.long, device=self.device)
        num_nodes = self.num_nodes.repeat_interleave(repeats)
        num_edges = self.num_edges.repeat_interleave(repeats)
        num_cum_nodes = num_nodes.cumsum(0)
        num_cum_edges = num_edges.cumsum(0)
        num_node = num_nodes.sum()
        num_edge = num_edges.sum()
        batch_size = repeats.sum()
        num_graphs = torch.ones(batch_size, device=self.device)

        # special case 1: graphs[i] may have no node or no edge
        # special case 2: repeats[i] may be 0
        cum_repeats_shifted = repeats.cumsum(0) - repeats
        graph_mask = cum_repeats_shifted < batch_size
        cum_repeats_shifted = cum_repeats_shifted[graph_mask]

        index = num_cum_nodes - num_nodes
        index = torch.cat([index, index[cum_repeats_shifted]])
        value = torch.cat([-num_nodes, self.num_nodes[graph_mask]])
        mask = index < num_node
        node_index = scatter_add(value[mask], index[mask], dim_size=num_node)
        node_index = (node_index + 1).cumsum(0) - 1

        index = num_cum_edges - num_edges
        index = torch.cat([index, index[cum_repeats_shifted]])
        value = torch.cat([-num_edges, self.num_edges[graph_mask]])
        mask = index < num_edge
        edge_index = scatter_add(value[mask], index[mask], dim_size=num_edge)
        edge_index = (edge_index + 1).cumsum(0) - 1

        graph_index = torch.repeat_interleave(repeats)

        offsets = self._get_offsets(num_nodes, num_edges)
        edge_list = self.edge_list[edge_index]
        edge_list[:, :2] += (offsets - self._offsets[edge_index]).unsqueeze(-1)

        node_offsets = None
        edge_offsets = None
        graph_offsets = None
        data_dict = {}
        for k, v in self.data_dict.items():
            num_xs = None
            pack_offsets = None
            for _type in self.meta_dict[k]:
                if _type == "node":
                    v = v[node_index]
                    num_xs = num_nodes
                elif _type == "edge":
                    v = v[edge_index]
                    num_xs = num_edges
                elif _type == "graph":
                    v = v[graph_index]
                    num_xs = num_graphs
                elif _type == "node reference":
                    if node_offsets is None:
                        node_offsets = self._get_repeat_pack_offsets(self.num_nodes, repeats)
                    pack_offsets = node_offsets
                elif _type == "edge reference":
                    if edge_offsets is None:
                        edge_offsets = self._get_repeat_pack_offsets(self.num_edges, repeats)
                    pack_offsets = edge_offsets
                elif _type == "graph reference":
                    if graph_offsets is None:
                        graph_offsets = self._get_repeat_pack_offsets(num_graphs, repeats)
                    pack_offsets = graph_offsets
            # add offsets to make references point to indexes in their own graph
            if num_xs is not None and pack_offsets is not None:
                v = v + pack_offsets.repeat_interleave(num_xs)
            data_dict[k] = v

        return type(self)(edge_list, edge_weight=self.edge_weight[edge_index],
                          num_nodes=num_nodes, num_edges=num_edges, num_relation=self.num_relation,
                          offsets=offsets, meta_dict=self.meta_dict, **data_dict)

    def get_item(self, index):
        """
        Get the i-th graph from this packed graph.

        Parameters:
            index (int): graph index

        Returns:
            Graph
        """
        node_index = torch.arange(self.num_cum_nodes[index] - self.num_nodes[index], self.num_cum_nodes[index],
                                  device=self.device)
        edge_index = torch.arange(self.num_cum_edges[index] - self.num_edges[index], self.num_cum_edges[index],
                                  device=self.device)
        graph_index = index
        edge_list = self.edge_list[edge_index].clone()
        edge_list[:, :2] -= self._offsets[edge_index].unsqueeze(-1)
        data_dict, meta_dict = self.data_mask(node_index, edge_index, graph_index=graph_index)

        return self.unpacked_type(edge_list, edge_weight=self.edge_weight[edge_index], num_node=self.num_nodes[index],
                                  num_relation=self.num_relation, meta_dict=meta_dict, **data_dict)

    def _get_cumulative(self, edge_list, num_nodes, num_edges, offsets):
        if edge_list is None:
            raise ValueError("`edge_list` should be provided")
        if num_edges is None:
            raise ValueError("`num_edges` should be provided")

        edge_list = torch.as_tensor(edge_list)
        num_edges = torch.as_tensor(num_edges, device=edge_list.device)
        num_edge = num_edges.sum()
        if num_edge != len(edge_list):
            raise ValueError("Sum of `num_edges` is %d, but found %d edges in `edge_list`" % (num_edge, len(edge_list)))
        num_cum_edges = num_edges.cumsum(0)

        if offsets is None:
            _edge_list = edge_list
        else:
            offsets = torch.as_tensor(offsets, device=edge_list.device)
            _edge_list = edge_list.clone()
            _edge_list[:, :2] -= offsets.unsqueeze(-1)
        if num_nodes is None:
            num_nodes = []
            for num_edge, num_cum_edge in zip(num_edges, num_cum_edges):
                num_nodes.append(self._maybe_num_node(_edge_list[num_cum_edge - num_edge: num_cum_edge]))
        num_nodes = torch.as_tensor(num_nodes, device=edge_list.device)
        num_cum_nodes = num_nodes.cumsum(0)

        return edge_list, num_nodes, num_edges, num_cum_nodes, num_cum_edges, offsets

    def _get_num_xs(self, index, num_cum_xs):
        x = torch.zeros(num_cum_xs[-1], dtype=torch.long, device=self.device)
        x[index] = 1
        num_cum_indexes = x.cumsum(0)
        num_cum_indexes = torch.cat([torch.zeros(1, dtype=torch.long, device=self.device), num_cum_indexes])
        new_num_cum_xs = num_cum_indexes[num_cum_xs]
        prepend = torch.zeros(1, dtype=torch.long, device=self.device)
        new_num_xs = torch.diff(new_num_cum_xs, prepend=prepend)
        return new_num_xs

    def data_mask(self, node_index=None, edge_index=None, graph_index=None, include=None, exclude=None):
        data_dict, meta_dict = self.data_by_meta(include, exclude)
        node_mapping = None
        edge_mapping = None
        graph_mapping = None
        for k, v in data_dict.items():
            for type in meta_dict[k]:
                if type == "node" and node_index is not None:
                    v = v[node_index]
                elif type == "edge" and edge_index is not None:
                    v = v[edge_index]
                elif type == "graph" and graph_index is not None:
                    v = v[graph_index]
                elif type == "node reference" and node_index is not None:
                    if node_mapping is None:
                        node_mapping = self._get_mapping(node_index, self.num_node)
                    v = node_mapping[v]
                elif type == "edge reference" and edge_index is not None:
                    if edge_mapping is None:
                        edge_mapping = self._get_mapping(edge_index, self.num_edge)
                    v = edge_mapping[v]
                elif type == "graph reference" and graph_index is not None:
                    if graph_mapping is None:
                        graph_mapping = self._get_mapping(graph_index, self.batch_size)
                    v = graph_mapping[v]
            data_dict[k] = v

        return data_dict, meta_dict

    def __getitem__(self, index):
        # why do we check tuple?
        # case 1: x[0, 1] is parsed as (0, 1)
        # case 2: x[[0, 1]] is parsed as [0, 1]
        if not isinstance(index, tuple):
            index = (index,)

        if isinstance(index[0], int):
            item = self.get_item(index[0])
            if len(index) > 1:
                item = item[index[1:]]
            return item
        if len(index) > 1:
            raise ValueError("Complex indexing is not supported for PackedGraph")

        index = self._standarize_index(index[0], self.batch_size)
        count = index.bincount(minlength=self.batch_size)
        if self.batch_size > 0 and count.max() > 1:
            graph = self.repeat_interleave(count)
            index_order = index.argsort()
            order = torch.zeros_like(index)
            order[index_order] = torch.arange(len(index), dtype=torch.long, device=self.device)
            return graph.subbatch(order)

        return self.subbatch(index)

    def __len__(self):
        return len(self.num_nodes)

    def full(self):
        """
        Return a pack of fully connected graphs.

        This is useful for computing node-pair-wise features.
        The computation can be implemented as message passing over a fully connected graph.

        Returns:
            PackedGraph
        """
        # TODO: more efficient implementation?
        graphs = self.unpack()
        graphs = [graph.full() for graph in graphs]
        return graphs[0].pack(graphs)

    @utils.cached_property
    def node2graph(self):
        """Node id to graph id mapping."""
        node2graph = torch.repeat_interleave(self.num_nodes)
        return node2graph

    @utils.cached_property
    def edge2graph(self):
        """Edge id to graph id mapping."""
        edge2graph = torch.repeat_interleave(self.num_edges)
        return edge2graph

    @property
    def batch_size(self):
        """Batch size."""
        return len(self.num_nodes)

    def node_mask(self, index, compact=False):
        """
        Return a masked packed graph based on the specified nodes.

        Note the compact option is only applied to node ids but not graph ids.
        To generate compact graph ids, use :meth:`subbatch`.

        Parameters:
            index (array_like): node index
            compact (bool, optional): compact node ids or not

        Returns:
            PackedGraph
        """
        index = self._standarize_index(index, self.num_node)
        mapping = -torch.ones(self.num_node, dtype=torch.long, device=self.device)
        if compact:
            mapping[index] = torch.arange(len(index), device=self.device)
            num_nodes = self._get_num_xs(index, self.num_cum_nodes)
            offsets = self._get_offsets(num_nodes, self.num_edges)
        else:
            mapping[index] = index
            num_nodes = self.num_nodes
            offsets = self._offsets

        edge_list = self.edge_list.clone()
        edge_list[:, :2] = mapping[edge_list[:, :2]]
        edge_index = (edge_list[:, :2] >= 0).all(dim=-1)
        num_edges = self._get_num_xs(edge_index, self.num_cum_edges)

        if compact:
            data_dict, meta_dict = self.data_mask(index, edge_index)
        else:
            data_dict, meta_dict = self.data_mask(edge_index=edge_index)

        return type(self)(edge_list[edge_index], edge_weight=self.edge_weight[edge_index], num_nodes=num_nodes,
                          num_edges=num_edges, num_relation=self.num_relation, offsets=offsets[edge_index],
                          meta_dict=meta_dict, **data_dict)

    def edge_mask(self, index):
        """
        Return a masked packed graph based on the specified edges.

        Parameters:
            index (array_like): edge index

        Returns:
            PackedGraph
        """
        index = self._standarize_index(index, self.num_edge)
        data_dict, meta_dict = self.data_mask(edge_index=index)
        num_edges = self._get_num_xs(index, self.num_cum_edges)

        return type(self)(self.edge_list[index], edge_weight=self.edge_weight[index], num_nodes=self.num_nodes,
                          num_edges=num_edges, num_relation=self.num_relation, offsets=self._offsets[index],
                          meta_dict=meta_dict, **data_dict)

    def graph_mask(self, index, compact=False):
        """
        Return a masked packed graph based on the specified graphs.

        This function can also be used to re-order the graphs.

        Parameters:
            index (array_like): graph index
            compact (bool, optional): compact graph ids or not

        Returns:
            PackedGraph
        """
        index = self._standarize_index(index, self.batch_size)
        graph_mapping = -torch.ones(self.batch_size, dtype=torch.long, device=self.device)
        graph_mapping[index] = torch.arange(len(index), device=self.device)

        node_index = graph_mapping[self.node2graph] >= 0
        node_index = self._standarize_index(node_index, self.num_node)
        mapping = -torch.ones(self.num_node, dtype=torch.long, device=self.device)
        if compact:
            key = graph_mapping[self.node2graph[node_index]] * self.num_node + node_index
            order = key.argsort()
            node_index = node_index[order]
            mapping[node_index] = torch.arange(len(node_index), device=self.device)
            num_nodes = self.num_nodes[index]
        else:
            mapping[node_index] = node_index
            num_nodes = torch.zeros_like(self.num_nodes)
            num_nodes[index] = self.num_nodes[index]

        edge_list = self.edge_list.clone()
        edge_list[:, :2] = mapping[edge_list[:, :2]]
        edge_index = (edge_list[:, :2] >= 0).all(dim=-1)
        edge_index = self._standarize_index(edge_index, self.num_edge)
        if compact:
            key = graph_mapping[self.edge2graph[edge_index]] * self.num_edge + edge_index
            order = key.argsort()
            edge_index = edge_index[order]
            num_edges = self.num_edges[index]
        else:
            num_edges = torch.zeros_like(self.num_edges)
            num_edges[index] = self.num_edges[index]
        offsets = self._get_offsets(num_nodes, num_edges)

        if compact:
            data_dict, meta_dict = self.data_mask(node_index, edge_index, graph_index=index)
        else:
            data_dict, meta_dict = self.data_mask(edge_index=edge_index)

        return type(self)(edge_list[edge_index], edge_weight=self.edge_weight[edge_index], num_nodes=num_nodes,
                          num_edges=num_edges, num_relation=self.num_relation, offsets=offsets,
                          meta_dict=meta_dict, **data_dict)

    def subbatch(self, index):
        """
        Return a subbatch based on the specified graphs.
        Equivalent to :meth:`graph_mask(index, compact=True) <graph_mask>`.

        Parameters:
            index (array_like): graph index

        Returns:
            PackedGraph

        See also:
            :meth:`PackedGraph.graph_mask`
        """
        return self.graph_mask(index, compact=True)

    def line_graph(self):
        """
        Construct a packed line graph of this packed graph.
        The node features of the line graphs are inherited from the edge features of the original graphs.

        In the line graph, each node corresponds to an edge in the original graph.
        For a pair of edges (a, b) and (b, c) that share the same intermediate node in the original graph,
        there is a directed edge (a, b) -> (b, c) in the line graph.

        Returns:
            PackedGraph
        """
        node_in, node_out = self.edge_list.t()[:2]
        edge_index = torch.arange(self.num_edge, device=self.device)
        edge_in = edge_index[node_out.argsort()]
        edge_out = edge_index[node_in.argsort()]

        degree_in = node_in.bincount(minlength=self.num_node)
        degree_out = node_out.bincount(minlength=self.num_node)
        size = degree_out * degree_in
        starts = (size.cumsum(0) - size).repeat_interleave(size)
        range = torch.arange(size.sum(), device=self.device)
        # each node u has degree_out[u] * degree_in[u] local edges
        local_index = range - starts
        local_inner_size = degree_in.repeat_interleave(size)
        edge_in_offset = (degree_out.cumsum(0) - degree_out).repeat_interleave(size)
        edge_out_offset = (degree_in.cumsum(0) - degree_in).repeat_interleave(size)
        edge_in_index = torch.div(local_index, local_inner_size, rounding_mode="floor") + edge_in_offset
        edge_out_index = local_index % local_inner_size + edge_out_offset

        edge_in = edge_in[edge_in_index]
        edge_out = edge_out[edge_out_index]
        edge_list = torch.stack([edge_in, edge_out], dim=-1)
        node_feature = getattr(self, "edge_feature", None)
        num_nodes = self.num_edges
        num_edges = scatter_add(size, self.node2graph, dim=0, dim_size=self.batch_size)
        offsets = self._get_offsets(num_nodes, num_edges)

        return PackedGraph(edge_list, num_nodes=num_nodes, num_edges=num_edges, offsets=offsets,
                           node_feature=node_feature)

    def undirected(self, add_inverse=False):
        """
        Flip all the edges to create undirected graphs.

        For knowledge graphs, the flipped edges can either have the original relation or an inverse relation.
        The inverse relation for relation :math:`r` is defined as :math:`|R| + r`.

        Parameters:
            add_inverse (bool, optional): whether to use inverse relations for flipped edges
        """
        edge_list = self.edge_list.clone()
        edge_list[:, :2] = edge_list[:, :2].flip(1)
        num_relation = self.num_relation
        if num_relation and add_inverse:
            edge_list[:, 2] += num_relation
            num_relation = num_relation * 2
        edge_list = torch.stack([self.edge_list, edge_list], dim=1).flatten(0, 1)
        offsets = self._offsets.unsqueeze(-1).expand(-1, 2).flatten()

        index = torch.arange(self.num_edge, device=self.device).unsqueeze(-1).expand(-1, 2).flatten()
        data_dict, meta_dict = self.data_mask(edge_index=index, exclude="edge reference")

        return type(self)(edge_list, edge_weight=self.edge_weight[index], num_nodes=self.num_nodes,
                          num_edges=self.num_edges * 2, num_relation=num_relation, offsets=offsets,
                          meta_dict=meta_dict, **data_dict)

    def detach(self):
        """
        Detach this packed graph.
        """
        return type(self)(self.edge_list.detach(), edge_weight=self.edge_weight.detach(),
                          num_nodes=self.num_nodes, num_edges=self.num_edges, num_relation=self.num_relation,
                          offsets=self._offsets, meta_dict=self.meta_dict, **utils.detach(self.data_dict))

    def clone(self):
        """
        Clone this packed graph.
        """
        return type(self)(self.edge_list.clone(), edge_weight=self.edge_weight.clone(),
                          num_nodes=self.num_nodes, num_edges=self.num_edges, num_relation=self.num_relation,
                          offsets=self._offsets, meta_dict=self.meta_dict, **utils.clone(self.data_dict))

    def cuda(self, *args, **kwargs):
        """
        Return a copy of this packed graph in CUDA memory.

        This is a non-op if the graph is already on the correct device.
        """
        edge_list = self.edge_list.cuda(*args, **kwargs)

        if edge_list is self.edge_list:
            return self
        else:
            return type(self)(edge_list, edge_weight=self.edge_weight,
                              num_nodes=self.num_nodes, num_edges=self.num_edges, num_relation=self.num_relation,
                              offsets=self._offsets, meta_dict=self.meta_dict,
                              **utils.cuda(self.data_dict, *args, **kwargs))

    def cpu(self):
        """
        Return a copy of this packed graph in CPU memory.

        This is a non-op if the graph is already in CPU memory.
        """
        edge_list = self.edge_list.cpu()

        if edge_list is self.edge_list:
            return self
        else:
            return type(self)(edge_list, edge_weight=self.edge_weight,
                              num_nodes=self.num_nodes, num_edges=self.num_edges, num_relation=self.num_relation,
                              offsets=self._offsets, meta_dict=self.meta_dict, **utils.cpu(self.data_dict))

    def __repr__(self):
        fields = ["batch_size=%d" % self.batch_size,
                  "num_nodes=%s" % pretty.long_array(self.num_nodes.tolist()),
                  "num_edges=%s" % pretty.long_array(self.num_edges.tolist())]
        if self.num_relation is not None:
            fields.append("num_relation=%d" % self.num_relation)
        if self.device.type != "cpu":
            fields.append("device='%s'" % self.device)
        return "%s(%s)" % (self.__class__.__name__, ", ".join(fields))

    def visualize(self, titles=None, save_file=None, figure_size=(3, 3), layout="spring", num_row=None, num_col=None):
        """
        Visualize the packed graphs with matplotlib.

        Parameters:
            titles (list of str, optional): title for each graph. Default is the ID of each graph.
            save_file (str, optional): ``png`` or ``pdf`` file to save visualization.
                If not provided, show the figure in window.
            figure_size (tuple of int, optional): width and height of the figure
            layout (str, optional): graph layout
            num_row (int, optional): number of rows in the figure
            num_col (int, optional): number of columns in the figure

        See also:
            `NetworkX graph layout`_

            .. _NetworkX graph layout:
                https://networkx.github.io/documentation/stable/reference/drawing.html#module-networkx.drawing.layout
        """
        if titles is None:
            graph = self.get_item(0)
            titles = ["%s %d" % (type(graph).__name__, i) for i in range(self.batch_size)]
        if num_col is None:
            if num_row is None:
                num_col = math.ceil(self.batch_size ** 0.5)
            else:
                num_col = math.ceil(self.batch_size / num_row)
        if num_row is None:
            num_row = math.ceil(self.batch_size / num_col)

        figure_size = (num_col * figure_size[0], num_row * figure_size[1])
        fig = plt.figure(figsize=figure_size)

        for i in range(self.batch_size):
            graph = self.get_item(i)
            ax = fig.add_subplot(num_row, num_col, i + 1)
            graph.visualize(title=titles[i], ax=ax, layout=layout)
        # remove the space of axis labels
        fig.tight_layout()

        if save_file:
            fig.savefig(save_file)
        else:
            fig.show()


Graph.packed_type = PackedGraph


def cat(graphs):
    for i, graph in enumerate(graphs):
        if not isinstance(graph, PackedGraph):
            graphs[i] = graph.pack([graph])

    edge_list = torch.cat([graph.edge_list for graph in graphs])
    pack_num_nodes = torch.stack([graph.num_node for graph in graphs])
    pack_num_edges = torch.stack([graph.num_edge for graph in graphs])
    pack_num_cum_edges = pack_num_edges.cumsum(0)
    graph_index = pack_num_cum_edges < len(edge_list)
    pack_offsets = scatter_add(pack_num_nodes[graph_index], pack_num_cum_edges[graph_index],
                               dim_size=len(edge_list))
    pack_offsets = pack_offsets.cumsum(0)

    edge_list[:, :2] += pack_offsets.unsqueeze(-1)
    offsets = torch.cat([graph._offsets for graph in graphs]) + pack_offsets

    edge_weight = torch.cat([graph.edge_weight for graph in graphs])
    num_nodes = torch.cat([graph.num_nodes for graph in graphs])
    num_edges = torch.cat([graph.num_edges for graph in graphs])
    num_relation = graphs[0].num_relation
    assert all(graph.num_relation == num_relation for graph in graphs)

    # only keep attributes that exist in all graphs
    # TODO: this interface is not safe. re-design the interface
    keys = set(graphs[0].meta_dict.keys())
    for graph in graphs:
        keys = keys.intersection(graph.meta_dict.keys())

    meta_dict = {k: graphs[0].meta_dict[k] for k in keys}
    data_dict = {}
    for k in keys:
        data_dict[k] = torch.cat([graph.data_dict[k] for graph in graphs])

    return type(graphs[0])(edge_list, edge_weight=edge_weight,
                           num_nodes=num_nodes, num_edges=num_edges, num_relation=num_relation, offsets=offsets,
                           meta_dict=meta_dict, **data_dict)