"""Torch modules for graph global pooling."""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import torch as th
import torch.nn as nn
import numpy as np

from ...backend import pytorch as F
from ...readout import sum_nodes, mean_nodes, max_nodes, broadcast_nodes,\
    softmax_nodes, topk_nodes


__all__ = ['SumPooling', 'AvgPooling', 'MaxPooling', 'SortPooling',
           'GlobalAttentionPooling', 'Set2Set',
           'SetTransformerEncoder', 'SetTransformerDecoder', 'WeightAndSum']

class SumPooling(nn.Module):
    r"""

    Description
    -----------
    Apply sum pooling over the nodes in a graph .

    .. math::
        r^{(i)} = \sum_{k=1}^{N_i} x^{(i)}_k

    Notes
    -----
        Input: Could be one graph, or a batch of graphs. If using a batch of graphs,
        make sure nodes in all graphs have the same feature size, and concatenate
        nodes' feature together as the input.

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

    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import SumPooling
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature size is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)  # feature size is 16
    >>>
    >>> sumpool = SumPooling()  # create a sum pooling layer

    Case 1: Input a single graph

    >>> sumpool(g1, g1_node_feats)
        tensor([[10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10.,
                 10., 10.]])

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats])
    >>>
    >>> sumpool(batch_g, batch_f)
        tensor([[10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10., 10.,
                 10., 10.],
                [20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20., 20.,
                 20., 20.]])
    """
    def __init__(self):
        super(SumPooling, self).__init__()

    def forward(self, graph, feat):
        r"""

        Compute sum pooling.

        Parameters
        ----------
        graph : DGLGraph
            a DGLGraph or a batch of DGLGraphs
        feat : torch.Tensor
            The input feature with shape :math:`(N, D)`, where :math:`N` is the number
            of nodes in the graph, and :math:`D` means the size of features.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(B, D)`, where :math:`B` refers to the
            batch size of input graphs.
        """
        with graph.local_scope():
            graph.ndata['h'] = feat
            readout = sum_nodes(graph, 'h')
            return readout


class AvgPooling(nn.Module):
    r"""

    Description
    -----------
    Apply average pooling over the nodes in a graph.

    .. math::
        r^{(i)} = \frac{1}{N_i}\sum_{k=1}^{N_i} x^{(i)}_k

    Notes
    -----
        Input: Could be one graph, or a batch of graphs. If using a batch of graphs,
        make sure nodes in all graphs have the same feature size, and concatenate
        nodes' feature together as the input.

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

    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import AvgPooling
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature size is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)  # feature size is 16
    >>>
    >>> avgpool = AvgPooling()  # create an average pooling layer

    Case 1: Input single graph

    >>> avgpool(g1, g1_node_feats)
        tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' note features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats])
    >>>
    >>> avgpool(batch_g, batch_f)
        tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
    """
    def __init__(self):
        super(AvgPooling, self).__init__()

    def forward(self, graph, feat):
        r"""

        Compute average pooling.

        Parameters
        ----------
        graph : DGLGraph
            A DGLGraph or a batch of DGLGraphs.
        feat : torch.Tensor
            The input feature with shape :math:`(N, D)`, where :math:`N` is the number
            of nodes in the graph, and :math:`D` means the size of features.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(B, D)`, where
            :math:`B` refers to the batch size of input graphs.
        """
        with graph.local_scope():
            graph.ndata['h'] = feat
            readout = mean_nodes(graph, 'h')
            return readout


class MaxPooling(nn.Module):
    r"""

    Description
    -----------
    Apply max pooling over the nodes in a graph.

    .. math::
        r^{(i)} = \max_{k=1}^{N_i}\left( x^{(i)}_k \right)

    Notes
    -----
        Input: Could be one graph, or a batch of graphs. If using a batch of graphs,
        make sure nodes in all graphs have the same feature size, and concatenate
        nodes' feature together as the input.

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

    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import MaxPooling
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature size is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)  # feature size is 16
    >>>
    >>> maxpool = MaxPooling()  # create a max pooling layer

    Case 1: Input a single graph

    >>> maxpool(g1, g1_node_feats)
        tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats])
    >>>
    >>> maxpool(batch_g, batch_f)
        tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
    """
    def __init__(self):
        super(MaxPooling, self).__init__()

    def forward(self, graph, feat):
        r"""Compute max pooling.

        Parameters
        ----------
        graph : DGLGraph
            A DGLGraph or a batch of DGLGraphs.
        feat : torch.Tensor
            The input feature with shape :math:`(N, *)`, where
            :math:`N` is the number of nodes in the graph.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(B, *)`, where
            :math:`B` refers to the batch size.
        """
        with graph.local_scope():
            graph.ndata['h'] = feat
            readout = max_nodes(graph, 'h')
            return readout


class SortPooling(nn.Module):
    r"""

    Description
    -----------
    Apply Sort Pooling (`An End-to-End Deep Learning Architecture for Graph Classification
    <https://www.cse.wustl.edu/~ychen/public/DGCNN.pdf>`__) over the nodes in a graph.

    Parameters
    ----------
    k : int
        The number of nodes to hold for each graph.

    Notes
    -----
        Input: Could be one graph, or a batch of graphs. If using a batch of graphs,
        make sure nodes in all graphs have the same feature size, and concatenate
        nodes' feature together as the input.

    Examples
    --------

    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import SortPooling
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature size is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)  # feature size is 16
    >>>
    >>> sortpool = SortPooling(k=2)  # create a sort pooling layer

    Case 1: Input a single graph

    >>> sortpool(g1, g1_node_feats)
        tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
                 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats])
    >>>
    >>> sortpool(batch_g, batch_f)
        tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
                 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
                [1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
                 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
    """
    def __init__(self, k):
        super(SortPooling, self).__init__()
        self.k = k

    def forward(self, graph, feat):
        r"""

        Compute sort pooling.

        Parameters
        ----------
        graph : DGLGraph
            A DGLGraph or a batch of DGLGraphs.
        feat : torch.Tensor
            The input feature with shape :math:`(N, D)`, where :math:`N` is the
            number of nodes in the graph, and :math:`D` means the size of features.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(B, k * D)`, where :math:`B` refers
            to the batch size of input graphs.
        """
        with graph.local_scope():
            # Sort the feature of each node in ascending order.
            feat, _ = feat.sort(dim=-1)
            graph.ndata['h'] = feat
            # Sort nodes according to their last features.
            ret = topk_nodes(graph, 'h', self.k, sortby=-1)[0].view(
                -1, self.k * feat.shape[-1])
            return ret


class GlobalAttentionPooling(nn.Module):
    r"""

    Description
    -----------
    Apply Global Attention Pooling (`Gated Graph Sequence Neural Networks
    <https://arxiv.org/abs/1511.05493.pdf>`__) over the nodes in a graph.

    .. math::
        r^{(i)} = \sum_{k=1}^{N_i}\mathrm{softmax}\left(f_{gate}
        \left(x^{(i)}_k\right)\right) f_{feat}\left(x^{(i)}_k\right)

    Parameters
    ----------
    gate_nn : torch.nn.Module
        A neural network that computes attention scores for each feature.
    feat_nn : torch.nn.Module, optional
        A neural network applied to each feature before combining them with attention
        scores.

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

    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import GlobalAttentionPooling
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature dimension is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)
    >>>
    >>> gate_nn = th.nn.Linear(16, 1)  # the gate layer that maps node feature to scalar
    >>> gap = GlobalAttentionPooling(gate_nn)  # create a Global Attention Pooling layer

    Case 1: Input a single graph

    >>> gap(g1, g1_node_feats)
        tensor([[1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
                 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]],
               grad_fn=<SegmentReduceBackward>)

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats], 0)
    >>>
    >>> gap(batch_g, batch_f)
        tensor([[1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
                 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
                [1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000,
                 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]],
               grad_fn=<SegmentReduceBackward>)

    Notes
    -----
    See our `GGNN example <https://github.com/dmlc/dgl/tree/master/examples/pytorch/ggnn>`_
    on how to use GatedGraphConv and GlobalAttentionPooling layer to build a Graph Neural
    Networks that can solve Soduku.
    """
    def __init__(self, gate_nn, feat_nn=None):
        super(GlobalAttentionPooling, self).__init__()
        self.gate_nn = gate_nn
        self.feat_nn = feat_nn

    def forward(self, graph, feat):
        r"""

        Compute global attention pooling.

        Parameters
        ----------
        graph : DGLGraph
            A DGLGraph or a batch of DGLGraphs.
        feat : torch.Tensor
            The input feature with shape :math:`(N, D)` where :math:`N` is the
            number of nodes in the graph, and :math:`D` means the size of features.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(B, D)`, where :math:`B` refers
            to the batch size.
        """
        with graph.local_scope():
            gate = self.gate_nn(feat)
            assert gate.shape[-1] == 1, "The output of gate_nn should have size 1 at the last axis."
            feat = self.feat_nn(feat) if self.feat_nn else feat

            graph.ndata['gate'] = gate
            gate = softmax_nodes(graph, 'gate')
            graph.ndata.pop('gate')

            graph.ndata['r'] = feat * gate
            readout = sum_nodes(graph, 'r')
            graph.ndata.pop('r')

            return readout


class Set2Set(nn.Module):
    r"""

    Description
    -----------
    For each individual graph in the batch, set2set computes

    .. math::
        q_t &= \mathrm{LSTM} (q^*_{t-1})

        \alpha_{i,t} &= \mathrm{softmax}(x_i \cdot q_t)

        r_t &= \sum_{i=1}^N \alpha_{i,t} x_i

        q^*_t &= q_t \Vert r_t

    for this graph.

    Parameters
    ----------
    input_dim : int
        The size of each input sample.
    n_iters : int
        The number of iterations.
    n_layers : int
        The number of recurrent layers.

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

    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import Set2Set
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature size is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)  # feature size is 16
    >>>
    >>> s2s = Set2Set(16, 2, 1)  # create a Set2Set layer(n_iters=2, n_layers=1)

    Case 1: Input a single graph

    >>> s2s(g1, g1_node_feats)
        tensor([[ 0.0162, -0.0080, -0.0864, -0.0488,  0.0408,  0.1140, -0.1426, -0.0379,
                 -0.1094,  0.2622,  0.1352,  0.0011,  0.1970, -0.1837, -0.0377, -0.0360,
                  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,
                  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]],
               grad_fn=<CatBackward>)

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats], 0)
    >>>
    >>> s2s(batch_g, batch_f)
        tensor([[ 0.0162, -0.0080, -0.0864, -0.0488,  0.0408,  0.1140, -0.1426, -0.0379,
                 -0.1094,  0.2622,  0.1352,  0.0011,  0.1970, -0.1837, -0.0377, -0.0360,
                  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,
                  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000],
                [ 0.0162, -0.0080, -0.0864, -0.0488,  0.0408,  0.1140, -0.1426, -0.0379,
                 -0.1094,  0.2622,  0.1352,  0.0011,  0.1970, -0.1837, -0.0377, -0.0360,
                  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,
                  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]],
               grad_fn=<CatBackward>)

    Notes
    -----
    Set2Set is widely used in molecular property predictions, see
    `dgl-lifesci's MPNN example <https://github.com/awslabs/dgl-lifesci/blob/
    ecd95c905479ec048097777039cf9a19cfdcf223/python/dgllife/model/model_zoo/
    mpnn_predictor.py>`__
    on how to use DGL's Set2Set layer in graph property prediction applications.
    """
    def __init__(self, input_dim, n_iters, n_layers):
        super(Set2Set, self).__init__()
        self.input_dim = input_dim
        self.output_dim = 2 * input_dim
        self.n_iters = n_iters
        self.n_layers = n_layers
        self.lstm = th.nn.LSTM(self.output_dim, self.input_dim, n_layers)
        self.reset_parameters()

    def reset_parameters(self):
        """Reinitialize learnable parameters."""
        self.lstm.reset_parameters()

    def forward(self, graph, feat):
        r"""
        Compute set2set pooling.

        Parameters
        ----------
        graph : DGLGraph
            The input graph.
        feat : torch.Tensor
            The input feature with shape :math:`(N, D)` where  :math:`N` is the
            number of nodes in the graph, and :math:`D` means the size of features.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(B, D)`, where :math:`B` refers to
            the batch size, and :math:`D` means the size of features.
        """
        with graph.local_scope():
            batch_size = graph.batch_size

            h = (feat.new_zeros((self.n_layers, batch_size, self.input_dim)),
                 feat.new_zeros((self.n_layers, batch_size, self.input_dim)))

            q_star = feat.new_zeros(batch_size, self.output_dim)

            for _ in range(self.n_iters):
                q, h = self.lstm(q_star.unsqueeze(0), h)
                q = q.view(batch_size, self.input_dim)
                e = (feat * broadcast_nodes(graph, q)).sum(dim=-1, keepdim=True)
                graph.ndata['e'] = e
                alpha = softmax_nodes(graph, 'e')
                graph.ndata['r'] = feat * alpha
                readout = sum_nodes(graph, 'r')
                q_star = th.cat([q, readout], dim=-1)

            return q_star

    def extra_repr(self):
        """Set the extra representation of the module.
        which will come into effect when printing the model.
        """
        summary = 'n_iters={n_iters}'
        return summary.format(**self.__dict__)


class MultiHeadAttention(nn.Module):
    r"""Multi-Head Attention block, used in Transformer, Set Transformer and so on.

    Parameters
    ----------
    d_model : int
        The feature size (input and output) in Multi-Head Attention layer.
    num_heads : int
        The number of heads.
    d_head : int
        The hidden size per head.
    d_ff : int
        The inner hidden size in the Feed-Forward Neural Network.
    dropouth : float
        The dropout rate of each sublayer.
    dropouta : float
        The dropout rate of attention heads.

    Notes
    -----
    This module was used in SetTransformer layer.
    """
    def __init__(self, d_model, num_heads, d_head, d_ff, dropouth=0., dropouta=0.):
        super(MultiHeadAttention, self).__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_head
        self.d_ff = d_ff
        self.proj_q = nn.Linear(d_model, num_heads * d_head, bias=False)
        self.proj_k = nn.Linear(d_model, num_heads * d_head, bias=False)
        self.proj_v = nn.Linear(d_model, num_heads * d_head, bias=False)
        self.proj_o = nn.Linear(num_heads * d_head, d_model, bias=False)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropouth),
            nn.Linear(d_ff, d_model)
        )
        self.droph = nn.Dropout(dropouth)
        self.dropa = nn.Dropout(dropouta)
        self.norm_in = nn.LayerNorm(d_model)
        self.norm_inter = nn.LayerNorm(d_model)
        self.reset_parameters()

    def reset_parameters(self):
        """Reinitialize learnable parameters."""
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, x, mem, lengths_x, lengths_mem):
        """
        Compute multi-head self-attention.

        Parameters
        ----------
        x : torch.Tensor
            The input tensor used to compute queries.
        mem : torch.Tensor
            The memory tensor used to compute keys and values.
        lengths_x : list
            The array of node numbers, used to segment x.
        lengths_mem : list
            The array of node numbers, used to segment mem.
        """
        batch_size = len(lengths_x)
        max_len_x = max(lengths_x)
        max_len_mem = max(lengths_mem)

        queries = self.proj_q(x).view(-1, self.num_heads, self.d_head)
        keys = self.proj_k(mem).view(-1, self.num_heads, self.d_head)
        values = self.proj_v(mem).view(-1, self.num_heads, self.d_head)

        # padding to (B, max_len_x/mem, num_heads, d_head)
        queries = F.pad_packed_tensor(queries, lengths_x, 0)
        keys = F.pad_packed_tensor(keys, lengths_mem, 0)
        values = F.pad_packed_tensor(values, lengths_mem, 0)

        # attention score with shape (B, num_heads, max_len_x, max_len_mem)
        e = th.einsum('bxhd,byhd->bhxy', queries, keys)
        # normalize
        e = e / np.sqrt(self.d_head)

        # generate mask
        mask = th.zeros(batch_size, max_len_x, max_len_mem).to(e.device)
        for i in range(batch_size):
            mask[i, :lengths_x[i], :lengths_mem[i]].fill_(1)
        mask = mask.unsqueeze(1)
        e.masked_fill_(mask == 0, -float('inf'))

        # apply softmax
        alpha = th.softmax(e, dim=-1)
        # sum of value weighted by alpha
        out = th.einsum('bhxy,byhd->bxhd', alpha, values)
        # project to output
        out = self.proj_o(
            out.contiguous().view(batch_size, max_len_x, self.num_heads * self.d_head))
        # pack tensor
        out = F.pack_padded_tensor(out, lengths_x)

        # intra norm
        x = self.norm_in(x + out)

        # inter norm
        x = self.norm_inter(x + self.ffn(x))

        return x


class SetAttentionBlock(nn.Module):
    r"""SAB block introduced in Set-Transformer paper.

    Parameters
    ----------
    d_model : int
        The feature size (input and output) in Multi-Head Attention layer.
    num_heads : int
        The number of heads.
    d_head : int
        The hidden size per head.
    d_ff : int
        The inner hidden size in the Feed-Forward Neural Network.
    dropouth : float
        The dropout rate of each sublayer.
    dropouta : float
        The dropout rate of attention heads.

    Notes
    -----
    This module was used in SetTransformer layer.
    """
    def __init__(self, d_model, num_heads, d_head, d_ff, dropouth=0., dropouta=0.):
        super(SetAttentionBlock, self).__init__()
        self.mha = MultiHeadAttention(d_model, num_heads, d_head, d_ff,
                                      dropouth=dropouth, dropouta=dropouta)

    def forward(self, feat, lengths):
        """
        Compute a Set Attention Block.

        Parameters
        ----------
        feat : torch.Tensor
            The input feature.
        lengths : list
            The array of node numbers, used to segment feat tensor.
        """
        return self.mha(feat, feat, lengths, lengths)


class InducedSetAttentionBlock(nn.Module):
    r"""ISAB block introduced in Set-Transformer paper.

    Parameters
    ----------
    d_model : int
        The feature size (input and output) in Multi-Head Attention layer.
    num_heads : int
        The number of heads.
    d_head : int
        The hidden size per head.
    d_ff : int
        The inner hidden size in the Feed-Forward Neural Network.
    dropouth : float
        The dropout rate of each sublayer.
    dropouta : float
        The dropout rate of attention heads.

    Notes
    -----
    This module was used in SetTransformer layer.
    """
    def __init__(self, m, d_model, num_heads, d_head, d_ff, dropouth=0., dropouta=0.):
        super(InducedSetAttentionBlock, self).__init__()
        self.m = m
        self.d_model = d_model
        self.inducing_points = nn.Parameter(
            th.FloatTensor(m, d_model)
        )
        self.mha = nn.ModuleList([
            MultiHeadAttention(d_model, num_heads, d_head, d_ff,
                               dropouth=dropouth, dropouta=dropouta) for _ in range(2)])
        self.reset_parameters()

    def reset_parameters(self):
        """Reinitialize learnable parameters."""
        nn.init.xavier_uniform_(self.inducing_points)

    def forward(self, feat, lengths):
        """
        Compute an Induced Set Attention Block.

        Parameters
        ----------
        feat : torch.Tensor
            The input feature.
        lengths : list
            The array of node numbers, used to segment feat tensor.

        Returns
        -------
        torch.Tensor
            The output feature
        """
        batch_size = len(lengths)
        query = self.inducing_points.repeat(batch_size, 1)
        memory = self.mha[0](query, feat, [self.m] * batch_size, lengths)
        return self.mha[1](feat, memory, lengths, [self.m] * batch_size)

    def extra_repr(self):
        """Set the extra representation of the module.
        which will come into effect when printing the model.
        """
        shape_str = '({}, {})'.format(self.inducing_points.shape[0], self.inducing_points.shape[1])
        return 'InducedVector: ' + shape_str


class PMALayer(nn.Module):
    r"""Pooling by Multihead Attention, used as the Decoder Module in Set Transformer.

    Parameters
    ----------
    d_model : int
        The feature size (input and output) in Multi-Head Attention layer.
    num_heads : int
        The number of heads.
    d_head : int
        The hidden size per head.
    d_ff : int
        The kernel size in FFN (Positionwise Feed-Forward Network) layer.
    dropouth : float
        The dropout rate of each sublayer.
    dropouta : float
        The dropout rate of attention heads.

    Notes
    -----
    This module was used in SetTransformer layer.
    """
    def __init__(self, k, d_model, num_heads, d_head, d_ff, dropouth=0., dropouta=0.):
        super(PMALayer, self).__init__()
        self.k = k
        self.d_model = d_model
        self.seed_vectors = nn.Parameter(
            th.FloatTensor(k, d_model)
        )
        self.mha = MultiHeadAttention(d_model, num_heads, d_head, d_ff,
                                      dropouth=dropouth, dropouta=dropouta)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropouth),
            nn.Linear(d_ff, d_model)
        )
        self.reset_parameters()

    def reset_parameters(self):
        """Reinitialize learnable parameters."""
        nn.init.xavier_uniform_(self.seed_vectors)

    def forward(self, feat, lengths):
        """
        Compute Pooling by Multihead Attention.

        Parameters
        ----------
        feat : torch.Tensor
            The input feature.
        lengths : list
            The array of node numbers, used to segment feat tensor.

        Returns
        -------
        torch.Tensor
            The output feature
        """
        batch_size = len(lengths)
        query = self.seed_vectors.repeat(batch_size, 1)
        return self.mha(query, self.ffn(feat), [self.k] * batch_size, lengths)

    def extra_repr(self):
        """Set the extra representation of the module.
        which will come into effect when printing the model.
        """
        shape_str = '({}, {})'.format(self.seed_vectors.shape[0], self.seed_vectors.shape[1])
        return 'SeedVector: ' + shape_str


class SetTransformerEncoder(nn.Module):
    r"""

    Description
    -----------
    The Encoder module in `Set Transformer: A Framework for Attention-based
    Permutation-Invariant Neural Networks <https://arxiv.org/pdf/1810.00825.pdf>`__.

    Parameters
    ----------
    d_model : int
        The hidden size of the model.
    n_heads : int
        The number of heads.
    d_head : int
        The hidden size of each head.
    d_ff : int
        The kernel size in FFN (Positionwise Feed-Forward Network) layer.
    n_layers : int
        The number of layers.
    block_type : str
        Building block type: 'sab' (Set Attention Block) or 'isab' (Induced
        Set Attention Block).
    m : int or None
        The number of induced vectors in ISAB Block. Set to None if block type
        is 'sab'.
    dropouth : float
        The dropout rate of each sublayer.
    dropouta : float
        The dropout rate of attention heads.

    Examples
    --------
    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import SetTransformerEncoder
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature size is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)  # feature size is 16
    >>>
    >>> set_trans_enc = SetTransformerEncoder(16, 4, 4, 64)  # create a settrans encoder.

    Case 1: Input a single graph

    >>> set_trans_enc(g1, g1_node_feats)
        tensor([[ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921]],
               grad_fn=<NativeLayerNormBackward>)

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats])
    >>>
    >>> set_trans_enc(batch_g, batch_f)
        tensor([[ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921],
                [ 0.1910, -0.7224, -1.7265,  1.8916,  0.4857,  0.8494, -1.4954, -0.2460,
                  1.8127, -0.8369, -0.8274,  0.4267, -0.5167,  0.6989,  0.1073, -0.0921]],
               grad_fn=<NativeLayerNormBackward>)

    See Also
    --------
    SetTransformerDecoder

    Notes
    -----
    SetTransformerEncoder is not a readout layer, the tensor it returned is nodewise
    representation instead out graphwise representation, and the SetTransformerDecoder
    would return a graph readout tensor.
    """
    def __init__(self, d_model, n_heads, d_head, d_ff,
                 n_layers=1, block_type='sab', m=None, dropouth=0., dropouta=0.):
        super(SetTransformerEncoder, self).__init__()
        self.n_layers = n_layers
        self.block_type = block_type
        self.m = m
        layers = []
        if block_type == 'isab' and m is None:
            raise KeyError('The number of inducing points is not specified in ISAB block.')

        for _ in range(n_layers):
            if block_type == 'sab':
                layers.append(
                    SetAttentionBlock(d_model, n_heads, d_head, d_ff,
                                      dropouth=dropouth, dropouta=dropouta))
            elif block_type == 'isab':
                layers.append(
                    InducedSetAttentionBlock(m, d_model, n_heads, d_head, d_ff,
                                             dropouth=dropouth, dropouta=dropouta))
            else:
                raise KeyError("Unrecognized block type {}: we only support sab/isab")

        self.layers = nn.ModuleList(layers)

    def forward(self, graph, feat):
        """
        Compute the Encoder part of Set Transformer.

        Parameters
        ----------
        graph : DGLGraph
            The input graph.
        feat : torch.Tensor
            The input feature with shape :math:`(N, D)`, where :math:`N` is the
            number of nodes in the graph.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(N, D)`.
        """
        lengths = graph.batch_num_nodes()
        for layer in self.layers:
            feat = layer(feat, lengths)
        return feat


class SetTransformerDecoder(nn.Module):
    r"""

    Description
    -----------
    The Decoder module in `Set Transformer: A Framework for Attention-based
    Permutation-Invariant Neural Networks <https://arxiv.org/pdf/1810.00825.pdf>`__.

    Parameters
    ----------
    d_model : int
        Hidden size of the model.
    num_heads : int
        The number of heads.
    d_head : int
        Hidden size of each head.
    d_ff : int
        Kernel size in FFN (Positionwise Feed-Forward Network) layer.
    n_layers : int
        The number of layers.
    k : int
        The number of seed vectors in PMA (Pooling by Multihead Attention) layer.
    dropouth : float
        Dropout rate of each sublayer.
    dropouta : float
        Dropout rate of attention heads.

    Examples
    --------
    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import SetTransformerDecoder
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature size is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)  # feature size is 16
    >>>
    >>> set_trans_dec = SetTransformerDecoder(16, 4, 4, 64, 1, 5)  # define the layer

    Case 1: Input a single graph

    >>> set_trans_dec(g1, g1_node_feats)
        tensor([[ 0.4635,  0.0275, -0.2637,  0.7168,  1.7655,  0.7687, -0.0031, -0.0562,
                  0.8125,  1.3546, -1.1208, -0.0198, -0.6820, -1.7057, -2.0887,  0.0310,
                  0.7714, -0.6568, -0.3391, -0.3792,  1.3881,  1.0091, -0.1747, -1.5370,
                  1.3287,  1.7640,  0.6081, -0.4512, -0.5583, -1.3908, -1.2931, -0.0894,
                  2.0826, -0.3916,  0.9458, -0.0952,  0.6316, -1.0485, -0.1104,  0.1100,
                  1.6364,  1.0246, -0.2355, -0.9597, -1.2427, -0.4817, -1.5894, -0.2764,
                  0.9830, -0.2319, -0.3492, -0.7830,  1.4185, -0.1799,  0.2063, -0.7108,
                  1.1052,  2.3187,  0.5359, -0.2413, -0.5357, -1.3557, -1.4125, -0.7675,
                 -0.0231, -0.2948, -0.3586,  0.6925,  0.6982,  1.1432, -0.5939, -1.6942,
                  1.6847,  1.5113, -0.7235,  0.0262, -1.4526, -0.0706, -1.3626,  0.8179]],
               grad_fn=<ViewBackward>)

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats])
    >>>
    >>> set_trans_dec(batch_g, batch_f)
        tensor([[ 0.4635,  0.0275, -0.2637,  0.7168,  1.7655,  0.7687, -0.0031, -0.0562,
                  0.8125,  1.3546, -1.1208, -0.0198, -0.6820, -1.7057, -2.0887,  0.0310,
                  0.7714, -0.6568, -0.3391, -0.3792,  1.3881,  1.0091, -0.1747, -1.5370,
                  1.3287,  1.7640,  0.6081, -0.4512, -0.5583, -1.3908, -1.2931, -0.0894,
                  2.0826, -0.3916,  0.9458, -0.0952,  0.6316, -1.0485, -0.1104,  0.1100,
                  1.6364,  1.0246, -0.2355, -0.9597, -1.2427, -0.4817, -1.5894, -0.2764,
                  0.9830, -0.2319, -0.3492, -0.7830,  1.4185, -0.1799,  0.2063, -0.7108,
                  1.1052,  2.3187,  0.5359, -0.2413, -0.5357, -1.3557, -1.4125, -0.7675,
                 -0.0231, -0.2948, -0.3586,  0.6925,  0.6982,  1.1432, -0.5939, -1.6942,
                  1.6847,  1.5113, -0.7235,  0.0262, -1.4526, -0.0706, -1.3626,  0.8179],
                [ 0.4635,  0.0275, -0.2637,  0.7168,  1.7655,  0.7687, -0.0031, -0.0562,
                  0.8125,  1.3546, -1.1208, -0.0198, -0.6820, -1.7057, -2.0887,  0.0310,
                  0.7714, -0.6568, -0.3391, -0.3792,  1.3881,  1.0091, -0.1747, -1.5370,
                  1.3287,  1.7640,  0.6081, -0.4512, -0.5583, -1.3908, -1.2931, -0.0894,
                  2.0826, -0.3916,  0.9458, -0.0952,  0.6316, -1.0485, -0.1104,  0.1100,
                  1.6364,  1.0246, -0.2355, -0.9597, -1.2427, -0.4817, -1.5894, -0.2764,
                  0.9830, -0.2319, -0.3492, -0.7830,  1.4185, -0.1799,  0.2063, -0.7108,
                  1.1052,  2.3187,  0.5359, -0.2413, -0.5357, -1.3557, -1.4125, -0.7675,
                 -0.0231, -0.2948, -0.3586,  0.6925,  0.6982,  1.1432, -0.5939, -1.6942,
                  1.6847,  1.5113, -0.7235,  0.0262, -1.4526, -0.0706, -1.3626,  0.8179]],
               grad_fn=<ViewBackward>)

    See Also
    --------
    SetTransformerEncoder
    """
    def __init__(self, d_model, num_heads, d_head, d_ff, n_layers, k, dropouth=0., dropouta=0.):
        super(SetTransformerDecoder, self).__init__()
        self.n_layers = n_layers
        self.k = k
        self.d_model = d_model
        self.pma = PMALayer(k, d_model, num_heads, d_head, d_ff,
                            dropouth=dropouth, dropouta=dropouta)
        layers = []
        for _ in range(n_layers):
            layers.append(
                SetAttentionBlock(d_model, num_heads, d_head, d_ff,
                                  dropouth=dropouth, dropouta=dropouta))

        self.layers = nn.ModuleList(layers)

    def forward(self, graph, feat):
        """
        Compute the decoder part of Set Transformer.

        Parameters
        ----------
        graph : DGLGraph
            The input graph.
        feat : torch.Tensor
            The input feature with shape :math:`(N, D)`, where :math:`N` is the
            number of nodes in the graph, and :math:`D` means the size of features.

        Returns
        -------
        torch.Tensor
            The output feature with shape :math:`(B, D)`, where :math:`B` refers to
            the batch size.
        """
        len_pma = graph.batch_num_nodes()
        len_sab = [self.k] * graph.batch_size
        feat = self.pma(feat, len_pma)
        for layer in self.layers:
            feat = layer(feat, len_sab)
        return feat.view(graph.batch_size, self.k * self.d_model)


class WeightAndSum(nn.Module):
    """Compute importance weights for atoms and perform a weighted sum.

    Parameters
    ----------
    in_feats : int
        Input atom feature size

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

    >>> import dgl
    >>> import torch as th
    >>> from dgl.nn.pytorch.glob import WeightAndSum
    >>>
    >>> g1 = dgl.rand_graph(10, 20)  # g1 is a random graph with 10 nodes and 20 edges
    >>> g1_node_feats = th.ones(10, 16)  # feature size is 16
    >>>
    >>> g2 = dgl.rand_graph(20, 50)  # g2 is a random graph with 20 nodes and 50 edges
    >>> g2_node_feats = th.ones(20, 16)  # feature size is 16
    >>>
    >>> weight_and_sum = WeightAndSum(16)  # create a weight and sum layer(in_feats=16)

    Case 1: Input a single graph

    >>> weight_and_sum(g1, g1_node_feats)
        tensor([[5.1436, 5.1436, 5.1436, 5.1436, 5.1436, 5.1436, 5.1436, 5.1436, 5.1436,
                 5.1436, 5.1436, 5.1436, 5.1436, 5.1436, 5.1436, 5.1436]],
               grad_fn=<SegmentReduceBackward>)

    Case 2: Input a batch of graphs

    Build a batch of DGL graphs and concatenate all graphs' node features into one tensor.

    >>> batch_g = dgl.batch([g1, g2])
    >>> batch_f = th.cat([g1_node_feats, g2_node_feats])
    >>>
    >>> sumpool(batch_g, batch_f)
        tensor([[ 5.1436,  5.1436,  5.1436,  5.1436,  5.1436,  5.1436,  5.1436,  5.1436,
                  5.1436,  5.1436,  5.1436,  5.1436,  5.1436,  5.1436,  5.1436,  5.1436],
                [10.2872, 10.2872, 10.2872, 10.2872, 10.2872, 10.2872, 10.2872, 10.2872,
                 10.2872, 10.2872, 10.2872, 10.2872, 10.2872, 10.2872, 10.2872, 10.2872]],
               grad_fn=<SegmentReduceBackward>)

    Notes
    -----
    WeightAndSum module was commonly used in molecular property prediction networks,
    see the GCN predictor in `dgl-lifesci <https://github.com/awslabs/dgl-lifesci/blob/
    ae0491431804611ba466ff413f69d435789dbfd5/python/dgllife/model/model_zoo/
    gcn_predictor.py>`__
    to understand how to use WeightAndSum layer to get the graph readout output.
    """
    def __init__(self, in_feats):
        super(WeightAndSum, self).__init__()
        self.in_feats = in_feats
        self.atom_weighting = nn.Sequential(
            nn.Linear(in_feats, 1),
            nn.Sigmoid()
        )

    def forward(self, g, feats):
        """Compute molecule representations out of atom representations

        Parameters
        ----------
        g : DGLGraph
            DGLGraph with batch size B for processing multiple molecules in parallel
        feats : FloatTensor of shape (N, self.in_feats)
            Representations for all atoms in the molecules
            * N is the total number of atoms in all molecules

        Returns
        -------
        FloatTensor of shape (B, self.in_feats)
            Representations for B molecules
        """
        with g.local_scope():
            g.ndata['h'] = feats
            g.ndata['w'] = self.atom_weighting(g.ndata['h'])
            h_g_sum = sum_nodes(g, 'h', 'w')

        return h_g_sum
