"""User-defined function related data structures."""
from __future__ import absolute_import


class EdgeBatch(object):
    """The class that can represent a batch of edges.

    Parameters
    ----------
    graph : DGLGraph
        Graph object.
    eid : Tensor
        Edge IDs.
    etype : (str, str, str)
        Edge type.
    src_data : dict[str, Tensor]
        Src node features.
    edge_data : dict[str, Tensor]
        Edge features.
    dst_data : dict[str, Tensor]
        Dst node features.
    """

    def __init__(self, graph, eid, etype, src_data, edge_data, dst_data):
        self._graph = graph
        self._eid = eid
        self._etype = etype
        self._src_data = src_data
        self._edge_data = edge_data
        self._dst_data = dst_data

    @property
    def src(self):
        """Return a view of the source node features for the edges in the batch.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set a node feature 'h'.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.ndata['h'] = torch.ones(2, 1)

        >>> # Define a UDF that retrieves the source node features for edges.
        >>> def edge_udf(edges):
        >>>     # edges.src['h'] is a tensor of shape (E, 1),
        >>>     # where E is the number of edges in the batch.
        >>>     return {'src': edges.src['h']}

        >>> # Copy features from source nodes to edges.
        >>> g.apply_edges(edge_udf)
        >>> g.edata['src']
        tensor([[1.],
                [1.],
                [1.]])

        >>> # Use edge UDF in message passing, which is equivalent to
        >>> # dgl.function.copy_u.
        >>> import dgl.function as fn
        >>> g.update_all(edge_udf, fn.sum('src', 'h'))
        >>> g.ndata['h']
        tensor([[1.],
                [2.]])
        """
        return self._src_data

    @property
    def dst(self):
        """Return a view of the destination node features for the edges in the batch.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set a node feature 'h'.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.ndata['h'] = torch.tensor([[0.], [1.]])

        >>> # Define a UDF that retrieves the destination node features for
        >>> # edges.
        >>> def edge_udf(edges):
        >>>     # edges.dst['h'] is a tensor of shape (E, 1),
        >>>     # where E is the number of edges in the batch.
        >>>     return {'dst': edges.dst['h']}

        >>> # Copy features from destination nodes to edges.
        >>> g.apply_edges(edge_udf)
        >>> g.edata['dst']
        tensor([[1.],
                [1.],
                [1.]])

        >>> # Use edge UDF in message passing.
        >>> import dgl.function as fn
        >>> g.update_all(edge_udf, fn.sum('dst', 'h'))
        >>> g.ndata['h']
        tensor([[0.],
                [2.]])
        """
        return self._dst_data

    @property
    def data(self):
        """Return a view of the edge features for the edges in the batch.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set an edge feature 'h'.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.edata['h'] = torch.tensor([[1.], [1.], [1.]])

        >>> # Define a UDF that retrieves the feature 'h' for all edges.
        >>> def edge_udf(edges):
        >>>     # edges.data['h'] is a tensor of shape (E, 1),
        >>>     # where E is the number of edges in the batch.
        >>>     return {'data': edges.data['h']}

        >>> # Make a copy of the feature with name 'data'.
        >>> g.apply_edges(edge_udf)
        >>> g.edata['data']
        tensor([[1.],
                [1.],
                [1.]])

        >>> # Use edge UDF in message passing, which is equivalent to
        >>> # dgl.function.copy_e.
        >>> import dgl.function as fn
        >>> g.update_all(edge_udf, fn.sum('data', 'h'))
        >>> g.ndata['h']
        tensor([[1.],
                [2.]])
        """
        return self._edge_data

    def edges(self):
        """Return the edges in the batch.

        Returns
        -------
        (U, V, EID) : (Tensor, Tensor, Tensor)
            The edges in the batch. For each :math:`i`, :math:`(U[i], V[i])` is
            an edge from :math:`U[i]` to :math:`V[i]` with ID :math:`EID[i]`.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))

        >>> # Define a UDF that retrieves and concatenates the end nodes of the
        >>> # edges.
        >>> def edge_udf(edges):
        >>>     src, dst, _ = edges.edges()
        >>>     return {'uv': torch.stack([src, dst], dim=1).float()}

        >>> # Create a feature 'uv' with the end nodes of the edges.
        >>> g.apply_edges(edge_udf)
        >>> g.edata['uv']
        tensor([[0., 1.],
                [1., 1.],
                [1., 0.]])

        >>> # Use edge UDF in message passing.
        >>> import dgl.function as fn
        >>> g.update_all(edge_udf, fn.sum('uv', 'h'))
        >>> g.ndata['h']
        tensor([[1., 0.],
                [1., 2.]])
        """
        u, v = self._graph.find_edges(self._eid, etype=self.canonical_etype)
        return u, v, self._eid

    def batch_size(self):
        """Return the number of edges in the batch.

        Returns
        -------
        int

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))

        >>> # Define a UDF that returns one for each edge.
        >>> def edge_udf(edges):
        >>>     return {'h': torch.ones(edges.batch_size(), 1)}

        >>> # Creates a feature 'h'.
        >>> g.apply_edges(edge_udf)
        >>> g.edata['h']
        tensor([[1.],
                [1.],
                [1.]])

        >>> # Use edge UDF in message passing.
        >>> import dgl.function as fn
        >>> g.update_all(edge_udf, fn.sum('h', 'h'))
        >>> g.ndata['h']
        tensor([[1.],
                [2.]])
        """
        return len(self._eid)

    def __len__(self):
        """Return the number of edges in this edge batch.

        Returns
        -------
        int
        """
        return self.batch_size()

    @property
    def canonical_etype(self):
        """Return the canonical edge type (i.e. triplet of source, edge, and
        destination node type) for this edge batch."""
        return self._etype


class NodeBatch(object):
    """The class to represent a batch of nodes.

    Parameters
    ----------
    graph : DGLGraph
        Graph object.
    nodes : Tensor
        Node ids.
    ntype : str, optional
        The node type of this node batch,
    data : dict[str, Tensor]
        Node feature data.
    msgs : dict[str, Tensor], optional
        Messages data.
    """

    def __init__(self, graph, nodes, ntype, data, msgs=None):
        self._graph = graph
        self._nodes = nodes
        self._ntype = ntype
        self._data = data
        self._msgs = msgs

    @property
    def data(self):
        """Return a view of the node features for the nodes in the batch.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set a feature 'h'.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.ndata['h'] = torch.ones(2, 1)

        >>> # Define a UDF that computes the sum of the messages received and
        >>> # the original feature for each node.
        >>> def node_udf(nodes):
        >>>     # nodes.data['h'] is a tensor of shape (N, 1),
        >>>     # nodes.mailbox['m'] is a tensor of shape (N, D, 1),
        >>>     # where N is the number of nodes in the batch, D is the number
        >>>     # of messages received per node for this node batch.
        >>>     return {'h': nodes.data['h'] + nodes.mailbox['m'].sum(1)}

        >>> # Use node UDF in message passing.
        >>> import dgl.function as fn
        >>> g.update_all(fn.copy_u('h', 'm'), node_udf)
        >>> g.ndata['h']
        tensor([[2.],
                [3.]])
        """
        return self._data

    @property
    def mailbox(self):
        """Return a view of the messages received.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set a feature 'h'.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.ndata['h'] = torch.ones(2, 1)

        >>> # Define a UDF that computes the sum of the messages received and
        >>> # the original feature for each node.
        >>> def node_udf(nodes):
        >>>     # nodes.data['h'] is a tensor of shape (N, 1),
        >>>     # nodes.mailbox['m'] is a tensor of shape (N, D, 1),
        >>>     # where N is the number of nodes in the batch, D is the number
        >>>     # of messages received per node for this node batch.
        >>>     return {'h': nodes.data['h'] + nodes.mailbox['m'].sum(1)}

        >>> # Use node UDF in message passing.
        >>> import dgl.function as fn
        >>> g.update_all(fn.copy_u('h', 'm'), node_udf)
        >>> g.ndata['h']
        tensor([[2.],
                [3.]])
        """
        return self._msgs

    def nodes(self):
        """Return the nodes in the batch.

        Returns
        -------
        NID : Tensor
            The IDs of the nodes in the batch. :math:`NID[i]` gives the ID of
            the i-th node.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph and set a feature 'h'.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.ndata['h'] = torch.ones(2, 1)

        >>> # Define a UDF that computes the sum of the messages received and
        >>> # the original ID for each node.
        >>> def node_udf(nodes):
        >>>     # nodes.nodes() is a tensor of shape (N),
        >>>     # nodes.mailbox['m'] is a tensor of shape (N, D, 1),
        >>>     # where N is the number of nodes in the batch, D is the number
        >>>     # of messages received per node for this node batch.
        >>>     return {'h': nodes.nodes().unsqueeze(-1).float()
        >>>         + nodes.mailbox['m'].sum(1)}

        >>> # Use node UDF in message passing.
        >>> import dgl.function as fn
        >>> g.update_all(fn.copy_u('h', 'm'), node_udf)
        >>> g.ndata['h']
        tensor([[1.],
                [3.]])
        """
        return self._nodes

    def batch_size(self):
        """Return the number of nodes in the batch.

        Returns
        -------
        int

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> # Instantiate a graph.
        >>> g = dgl.graph((torch.tensor([0, 1, 1]), torch.tensor([1, 1, 0])))
        >>> g.ndata['h'] = torch.ones(2, 1)

        >>> # Define a UDF that computes the sum of the messages received for
        >>> # each node and increments the result by 1.
        >>> def node_udf(nodes):
        >>>     return {'h': torch.ones(nodes.batch_size(), 1)
        >>>         + nodes.mailbox['m'].sum(1)}

        >>> # Use node UDF in message passing.
        >>> import dgl.function as fn
        >>> g.update_all(fn.copy_u('h', 'm'), node_udf)
        >>> g.ndata['h']
        tensor([[2.],
                [3.]])
        """
        return len(self._nodes)

    def __len__(self):
        """Return the number of nodes in this node batch.

        Returns
        -------
        int
        """
        return self.batch_size()

    @property
    def ntype(self):
        """Return the node type of this node batch, if available."""
        return self._ntype