import tensorflow as tf
from kgcnn.layers.base import GraphBaseLayer
from kgcnn.layers.gather import GatherState, GatherEmbeddingSelection
from kgcnn.layers.pooling import PoolingLocalEdgesAttention, PoolingNodes, PoolingNodesAttention
from kgcnn.layers.modules import LazySubtract, Dense, Dropout, LazyConcatenate, Activation
from kgcnn.layers.update import GRUUpdate

# import tensorflow.keras as ks
# import tensorflow.python.keras as ks
ks = tf.keras

# Gated recurrent unit update. See kgcnn.layers.conv.mpnn_conv for details.
HamNetGRUUnion = GRUUpdate


@ks.utils.register_keras_serializable(package='kgcnn', name='HamNetNaiveUnion')
class HamNetNaiveUnion(GraphBaseLayer):
    r"""Simple union that concatenates a feature tensor :math:`\mathbf{x}` and its updates :math:`\mathbf{x}_u`
    and applies a fully connected dense layer,
    i.e. a linear transformation with weights :math:`\mathbf{W}^{\top}`, :math:`\mathbf{b}` plus activation
    :math:`\sigma`.

    .. math::

        \mathbf{x}^{\prime} = \sigma \left[ \left( \mathbf{x} \; || \; \mathbf{x}_u \right) \mathbf{W}^{\top} +
        \mathbf{b} \right]

    """
    def __init__(self,
                 units: int,
                 activation="kgcnn>leaky_relu",
                 use_bias: bool = True,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 **kwargs):
        r"""Initialize layer with arguments of :obj:`ks.layers.Dense`.

        Args:
            units (int): Positive integer, dimensionality of the output space.
            activation: Activation function to use.
                If you don't specify anything, no activation is applied
                (ie. "linear" activation: `a(x) = x`). Default is "kgcnn>leaky_relu".
            use_bias (bool): Boolean, whether the layer uses a bias vector. Default is True.
            kernel_initializer: Initializer for the `kernel` weights matrix. Default is "glorot_uniform".
            bias_initializer: Initializer for the bias vector. Default is "zeros".
            kernel_regularizer: Regularizer function applied to
                the `kernel` weights matrix. Default is None.
            bias_regularizer: Regularizer function applied to the bias vector. Default is None.
            activity_regularizer: Regularizer function applied to
                the output of the layer (its "activation"). Default is None.
            kernel_constraint: Constraint function applied to
                the `kernel` weights matrix. Default is None.
            bias_constraint: Constraint function applied to the bias vector. Default is None.
        """
        super(HamNetNaiveUnion, self).__init__(**kwargs)
        self.units = int(units)
        self.use_bias = use_bias
        kernel_args = {"kernel_regularizer": kernel_regularizer,
                       "activity_regularizer": activity_regularizer, "bias_regularizer": bias_regularizer,
                       "kernel_constraint": kernel_constraint, "bias_constraint": bias_constraint,
                       "kernel_initializer": kernel_initializer, "bias_initializer": bias_initializer}
        self.lay_dense = Dense(units=units, activation=activation, use_bias=use_bias, **kernel_args)
        self.lay_concat = LazyConcatenate()

    def build(self, input_shape):
        """Build layer."""
        super(HamNetNaiveUnion, self).build(input_shape)

    def call(self, inputs, **kwargs):
        """Forward pass.

        Args:
            inputs (list): [nodes, node_updates]

                - nodes (tf.RaggedTensor): Node features of shape `(batch, [N], F)`
                - node_updates (tf.RaggedTensor): Node features of shape `(batch, [N], F)`

        Returns:
            tf.RaggedTensor: Embedding tensor of updated node features of shape `(batch, [N], F)`.
        """
        n, nu = inputs
        nnu = self.lay_concat([n, nu], **kwargs)
        n_out = self.lay_dense(nnu, **kwargs)
        return n_out

    def get_config(self):
        """Update layer config."""
        config = super(HamNetNaiveUnion, self).get_config()
        config.update({"units": self.units, "use_bias": self.use_bias})
        conf_dense = self.lay_dense.get_config()
        for x in ["kernel_regularizer", "activity_regularizer", "bias_regularizer", "kernel_constraint",
                  "bias_constraint", "kernel_initializer", "bias_initializer", "use_bias", "activation"]:
            config.update({x: conf_dense[x]})
        return config


@ks.utils.register_keras_serializable(package='kgcnn', name='HamNetGlobalReadoutAttend')
class HamNetGlobalReadoutAttend(GraphBaseLayer):
    r"""Computes attentive updates for fingerprint generation according to `HamNet <https://arxiv.org/abs/2105.03688>`_.
    The naming convention follows the authors `implementation <https://github.com/PKUterran/MoleculeClub>`_.
    The layer is used in :obj:`HamNetFingerprintGenerator` and computes the attentive state updates.
    The node features are first transformed by a :obj:`Dense` layer:
    :math:`\mathbf{h}' = \sigma\;(\mathbf{h} \mathbf{W}^T)` which yields the attention coefficients from state
    :math:`\mathbf{s}`:

    .. math::
        a_i = w^T [\mathbf{h}_i' \; || \; \mathbf{s}]

    with :math:`\alpha_i = \text{softmax}({a_i \; | \; i \in V})` the final state update :math:`\mathbf{m}`:

    .. math::
        \mathbf{m} = \sigma \; \sum_i \alpha_i \mathbf{h}'_i

    Update :math:`\mathbf{m}` is returned by the layer. Here, :math:`\sigma` denotes an activation function.

    """

    def __init__(self,
                 units,
                 activation="kgcnn>leaky_relu",
                 activation_last="elu",
                 use_bias=True,
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 use_dropout=False,
                 rate=None, noise_shape=None, seed=None,
                 **kwargs):
        """Initialize layer.

        Args:
            units (int): Units for the linear transformation of node features before attention.
            activation (str, dict): Activation. Default is "kgcnn>leaky_relu".
            activation_last (str, dict): Last activation for messages. Default is "elu".
            use_bias (bool): Boolean, whether the layer uses a bias vector. Default is True.
            kernel_initializer: Initializer for the `kernel` weights matrix. Default is "glorot_uniform".
            bias_initializer: Initializer for the bias vector. Default is "zeros".
            kernel_regularizer: Regularizer function applied to
                the `kernel` weights matrix. Default is None.
            bias_regularizer: Regularizer function applied to the bias vector. Default is None.
            activity_regularizer: Regularizer function applied to
                the output of the layer (its "activation"). Default is None.
            kernel_constraint: Constraint function applied to
                the `kernel` weights matrix. Default is None.
            bias_constraint: Constraint function applied to the bias vector. Default is None.
            use_dropout (bool): Whether to use dropout on input features. Default is False.
            rate (float): Float between 0 and 1. Fraction of the input units to drop.
            noise_shape: 1D integer tensor representing the shape of the
                binary dropout mask that will be multiplied with the input.
            seed (int): A Python integer to use as random seed.
        """
        super(HamNetGlobalReadoutAttend, self).__init__(**kwargs)
        self.units = int(units)
        self.use_bias = use_bias
        self.use_dropout = use_dropout
        kernel_args = {"kernel_regularizer": kernel_regularizer,
                       "activity_regularizer": activity_regularizer, "bias_regularizer": bias_regularizer,
                       "kernel_constraint": kernel_constraint, "bias_constraint": bias_constraint,
                       "kernel_initializer": kernel_initializer, "bias_initializer": bias_initializer}
        self.gather_state = GatherState()
        if self.use_dropout:
            self.dropout_layer = Dropout(rate=rate, noise_shape=noise_shape, seed=seed)
        self.dense_attend = Dense(units=units, activation=activation, use_bias=use_bias, **kernel_args)
        self.dense_align = Dense(1, activation="linear", use_bias=use_bias, **kernel_args)
        self.lay_concat = LazyConcatenate(axis=-1)
        self.pool_attention = PoolingNodesAttention()
        self.final_activ = Activation(activation=activation_last,
                                      activity_regularizer=activity_regularizer)

    def build(self, input_shape):
        """Build layer."""
        super(HamNetGlobalReadoutAttend, self).build(input_shape)

    def call(self, inputs, **kwargs):
        """Forward pass.

        Args:
            inputs: [state, nodes]

                - state (tf.Tensor): Molecular embedding of shape `(batch, F)`
                - nodes (tf.RaggedTensor): Node features of shape `(batch, [N], F)`

        Returns:
            tf.RaggedTensor: Embedding tensor of pooled node attentions of shape (batch, F)
        """
        hm_ftr, hv_ftr = inputs
        hm_v_ftr = self.gather_state([hm_ftr, hv_ftr], **kwargs)

        attend_ftr = hv_ftr
        if self.use_dropout:
            attend_ftr = self.dropout_layer(attend_ftr, **kwargs)
        attend_ftr = self.dense_attend(attend_ftr, **kwargs)
        align_ftr = self.lay_concat([hm_v_ftr, hv_ftr], **kwargs)
        if self.use_dropout:
            align_ftr = self.dropout_layer(align_ftr, **kwargs)
        align_ftr = self.dense_align(align_ftr, **kwargs)
        mm_ftr = self.pool_attention([attend_ftr, align_ftr], **kwargs)
        mm_ftr = self.final_activ(mm_ftr, **kwargs)
        return mm_ftr, align_ftr

    def get_config(self):
        """Update layer config."""
        config = super(HamNetGlobalReadoutAttend, self).get_config()
        config.update({"use_bias": self.use_bias, "units": self.units, "use_dropout": self.use_dropout})
        conf_sub = self.dense_attend.get_config()
        for x in ["kernel_regularizer", "activity_regularizer", "bias_regularizer", "kernel_constraint",
                  "bias_constraint", "kernel_initializer", "bias_initializer", "activation"]:
            config.update({x: conf_sub[x]})
        if self.use_dropout:
            conf_drop = self.dropout_layer.get_config()
            for x in ["rate", "noise_shape", "seed"]:
                config.update({x: conf_drop[x]})
        conf_last = self.final_activ.get_config()
        config.update({"activation_last": conf_last["activation"]})
        return config


@ks.utils.register_keras_serializable(package='kgcnn', name='HamNetFingerprintGenerator')
class HamNetFingerprintGenerator(GraphBaseLayer):
    r"""Computes readout or fingerprint generation according to `HamNet <https://arxiv.org/abs/2105.03688>`_.
    The naming convention follows the authors `implementation <https://github.com/PKUterran/MoleculeClub>`_.
    The layer generates a molecular or global message by iteratively updating from node embeddings. Initial state
    :math:`\mathbf{s}^0 = \frac{1}{n} \sum_i \sigma (\mathbf{h} W^T)` is updated :math:`l=1\dots L` times from
    messages :math:`\mathbf{m}^l` via a gated recurrent unit and subsequent activation :math:`\sigma`:

    .. math::
        \mathbf{s}^{l+1} = \sigma \left[\; \text{GRU}(\mathbf{s}^l, \mathbf{m}^{l}) \;\right]

    The message is obtained from an attentive readout function :math:`f` which is implemented here in
    :obj:`HamNetGlobalReadoutAttend`:

    .. math::
        \mathbf{m}^{l+1} = f(\mathbf{h}, \mathbf{m}^l)

    The final embedding :math:`\mathbf{s}^L` is used as output or molecular state.

    """

    def __init__(self,
                 units: int,
                 units_attend: int,
                 activation="kgcnn>leaky_relu",
                 use_bias: bool = True,
                 depth=4,
                 pooling_method="mean",
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 recurrent_activation='sigmoid',
                 recurrent_initializer='orthogonal',
                 recurrent_regularizer=None,
                 recurrent_constraint=None,
                 dropout=0.0,
                 recurrent_dropout=0.0,
                 reset_after=True,
                 use_dropout=False,
                 rate=None, noise_shape=None, seed=None,
                 **kwargs):
        """Initialize layer.

        Args:
            units (int): Units for the linear transformation of node features before attention.
            units_attend (int): Units for attention attributes.
            activation (str, dict): Activation. Default is "kgcnn>leaky_relu".
            use_bias (bool): Boolean, whether the layer uses a bias vector. Default is True.
            depth (int): Number of iterations. Default is 4.
            pooling_method(str): Initial pooling before iteration. Default is "mean".
            kernel_initializer: Initializer for the `kernel` weights matrix. Default is "glorot_uniform".
            bias_initializer: Initializer for the bias vector. Default is "zeros".
            kernel_regularizer: Regularizer function applied to
                the `kernel` weights matrix. Default is None.
            bias_regularizer: Regularizer function applied to the bias vector. Default is None.
            activity_regularizer: Regularizer function applied to
                the output of the layer (its "activation"). Default is None.
            kernel_constraint: Constraint function applied to
                the `kernel` weights matrix. Default is None.
            bias_constraint: Constraint function applied to the bias vector. Default is None.
            recurrent_activation: Activation function to use for the recurrent step.
                Default: sigmoid (`sigmoid`). If you pass `None`, no activation is
                applied (ie. "linear" activation: `a(x) = x`).
            recurrent_initializer: Initializer for the `recurrent_kernel`
                weights matrix, used for the linear transformation of the recurrent state.
                Default: `orthogonal`.
            recurrent_regularizer: Regularizer function applied to the
                `recurrent_kernel` weights matrix. Default: `None`.
            recurrent_constraint: Constraint function applied to the `recurrent_kernel`
                weights matrix. Default: `None`.
            dropout: Float between 0 and 1. Fraction of the units to drop for the
                linear transformation of the inputs. Default: 0.
            recurrent_dropout: Float between 0 and 1. Fraction of the units to drop for
                the linear transformation of the recurrent state. Default: 0.
            reset_after: GRU convention (whether to apply reset gate after or
                before matrix multiplication). False = "before",
                True = "after" (default and cuDNN compatible).
            use_dropout (bool): Whether to use dropout on input features. Default is False.
            rate (float): Float between 0 and 1. Fraction of the input units to drop.
            noise_shape: 1D integer tensor representing the shape of the
                binary dropout mask that will be multiplied with the input.
            seed (int): A Python integer to use as random seed.
        """
        super(HamNetFingerprintGenerator, self).__init__(**kwargs)
        self.units = int(units)
        self.units_attend = int(units_attend)
        self.use_bias = bool(use_bias)
        self.use_dropout = use_dropout
        self.depth = int(depth)
        self.pooling_method = pooling_method
        kernel_args = {"kernel_regularizer": kernel_regularizer, "activity_regularizer": activity_regularizer,
                       "bias_regularizer": bias_regularizer,
                       "kernel_constraint": kernel_constraint, "bias_constraint": bias_constraint,
                       "kernel_initializer": kernel_initializer, "bias_initializer": bias_initializer}
        gru_args = {"recurrent_activation": recurrent_activation,
                    "use_bias": use_bias, "kernel_initializer": kernel_initializer,
                    "recurrent_initializer": recurrent_initializer, "bias_initializer": bias_initializer,
                    "kernel_regularizer": kernel_regularizer, "recurrent_regularizer": recurrent_regularizer,
                    "bias_regularizer": bias_regularizer, "kernel_constraint": kernel_constraint,
                    "recurrent_constraint": recurrent_constraint, "bias_constraint": bias_constraint,
                    "dropout": dropout, "recurrent_dropout": recurrent_dropout, "reset_after": reset_after}
        self.pool_nodes = PoolingNodes(pooling_method=self.pooling_method)
        self.vertex2mol = Dense(
            units=units, activation=activation, use_bias=use_bias, **kernel_args)

        self.readouts = [HamNetGlobalReadoutAttend(
            units=units_attend, activation=activation, activation_last="elu", use_bias=use_bias,
            use_dropout=use_dropout, rate=rate, noise_shape=noise_shape, seed=seed,
            **kernel_args) for _ in range(self.depth)]

        self.unions = [ks.layers.GRUCell(
            units=units, activation="tanh", **gru_args) for _ in range(self.depth)]
        self.final_activ = Activation(activation=activation,
                                      activity_regularizer=activity_regularizer)

    def build(self, input_shape):
        """Build layer."""
        super(HamNetFingerprintGenerator, self).build(input_shape)

    def call(self, inputs, **kwargs):
        """Forward pass.

        Args:
            inputs (tf.RaggedTensor): Node embeddings of shape `(batch, [N], F)`

        Returns:
            tf.RaggedTensor: Embedding tensor of pooled node attentions of shape `(batch, F)`
        """
        self.assert_ragged_input_rank(inputs)
        hv_ftr = inputs
        hm_ftr = self.vertex2mol(hv_ftr, **kwargs)
        hm_ftr = self.pool_nodes(hm_ftr, **kwargs)
        alignments = []
        for i in range(self.depth):
            mm_ftr, align = self.readouts[i]([hm_ftr, hv_ftr], **kwargs)
            # alignments.append(align)
            hm_ftr, _ = self.unions[i](mm_ftr, hm_ftr, **kwargs)
            hm_ftr = self.final_activ(hm_ftr, **kwargs)
        return hm_ftr

    def get_config(self):
        """Update layer config."""
        config = super(HamNetFingerprintGenerator, self).get_config()
        config.update({"use_bias": self.use_bias, "units": self.units, "units_attend": self.units_attend,
                       "use_dropout": self.use_dropout, "depth": self.depth, "pooling_method": self.pooling_method})
        conf_sub = self.vertex2mol.get_config()
        for x in ["kernel_regularizer", "activity_regularizer", "bias_regularizer", "kernel_constraint",
                  "bias_constraint", "kernel_initializer", "bias_initializer", "activation"]:
            config.update({x: conf_sub[x]})
        if len(self.unions) > 0:
            conf_gru = self.unions[0].get_config()
            for x in ["recurrent_activation", "recurrent_initializer", "recurrent_regularizer", "recurrent_constraint",
                      "dropout", "recurrent_dropout", "reset_after"]:
                config.update({x: conf_gru[x]})
        if len(self.readouts) > 0:
            conf_read = self.readouts[0].get_config()
            for x in ["use_dropout", "seed", "rate", "noise_shape", ]:
                config.update({x: conf_read[x]})
        return config


@ks.utils.register_keras_serializable(package='kgcnn', name='HamNaiveDynMessage')
class HamNaiveDynMessage(GraphBaseLayer):
    r"""Message passing block from `HamNet <https://arxiv.org/abs/2105.03688>`_ which makes use of attention.
    The naming convention follows the authors `implementation <https://github.com/PKUterran/MoleculeClub>`_.
    The layer computes the following, let :math:`\mathbf{h}`, :math:`\mathbf{\epsilon}_{ij}` be node, edge features
    and :math:`\mathbf{q}`, :math:`\mathbf{p}` be (generalized) node coordinates and momentum. With
    :math:`\mathbf{p}_{ij} = \mathbf{p}_{j} - \mathbf{p}_{i}` and
    :math:`\mathbf{q}_{ij} = \mathbf{q}_{j} - \mathbf{q}_{i}` the attention coefficients read:

    .. math::
        \mathbf{a}_{ij} = \mathbf{w}^T \left(\mathbf{p}_{ij} \; || \; \mathbf{q}_{ij} \; ||
        \mathbf{\epsilon}_{ij} \right)

    and the new node update or message, using the attention coefficients
    :math:`\alpha_{ij} = \, \text{softmax}(\{\mathbf{a}_{ij} \; | \; j \in \mathcal{N}(i)\})`:

    .. math::
        \mathbf{m}_{v} = \sigma \; \sum_{j \in \mathcal{N}(i)} \; \alpha_{ij} \;  \sigma
        \left[ \; \mathbf{h}_j \; \mathbf{W}^T \;  \right]

    and edge updates:

    .. math::
        \mathbf{m}_{e} = \sigma \left[ \left(\; \mathbf{h}_i \; || \; \mathbf{q}_{ij} \; || \;
        \mathbf{h}_j \; \right) \mathbf{W}^T \right]

    the layer returns :math:`\mathbf{m}_{v}` and :math:`\mathbf{m}_{e}`.

    """
    def __init__(self,
                 units,
                 units_edge,
                 activation="kgcnn>leaky_relu",
                 activation_last="elu",
                 use_bias=True,
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 use_dropout=False,
                 rate=None, noise_shape=None, seed=None,
                 **kwargs):
        r"""Initialize layer.

        Args:
            units (int): Units for the linear transformation of node features before attention.
            units_edge (int): Units for :obj:`Dense` layer for edge updates.
            activation (str, dict): Activation. Default is "kgcnn>leaky_relu".
            activation_last (str, dict): Last activation for messages. Default is "elu".
            use_bias (bool): Boolean, whether the layer uses a bias vector. Default is True.
            kernel_initializer: Initializer for the `kernel` weights matrix. Default is "glorot_uniform".
            bias_initializer: Initializer for the bias vector. Default is "zeros".
            kernel_regularizer: Regularizer function applied to
                the `kernel` weights matrix. Default is None.
            bias_regularizer: Regularizer function applied to the bias vector. Default is None.
            activity_regularizer: Regularizer function applied to
                the output of the layer (its "activation"). Default is None.
            kernel_constraint: Constraint function applied to
                the `kernel` weights matrix. Default is None.
            bias_constraint: Constraint function applied to the bias vector. Default is None.
            use_dropout (bool): Whether to use dropout on input features. Default is False.
            rate (float): Float between 0 and 1. Fraction of the input units to drop.
            noise_shape: 1D integer tensor representing the shape of the
                binary dropout mask that will be multiplied with the input.
            seed (int): A Python integer to use as random seed.
        """
        super(HamNaiveDynMessage, self).__init__(**kwargs)
        self.units = int(units)
        self.units_edge = int(units_edge)
        self.use_bias = use_bias
        self.use_dropout = use_dropout
        kernel_args = {"kernel_regularizer": kernel_regularizer,
                       "activity_regularizer": activity_regularizer, "bias_regularizer": bias_regularizer,
                       "kernel_constraint": kernel_constraint, "bias_constraint": bias_constraint,
                       "kernel_initializer": kernel_initializer, "bias_initializer": bias_initializer}
        self.gather_v = GatherEmbeddingSelection(selection_index= [0, 1])
        self.gather_p = GatherEmbeddingSelection(selection_index= [0, 1])
        self.gather_q = GatherEmbeddingSelection(selection_index= [0, 1])
        self.lazy_sub_p = LazySubtract()
        self.lazy_sub_q = LazySubtract()
        # self.lay_concat = LazyConcatenate(axis=-1)
        self.lay_concat_align = LazyConcatenate(axis=-1)
        self.lay_concat_edge = LazyConcatenate(axis=-1)
        if self.use_dropout:
            self.dropout_layer = Dropout(rate=rate, noise_shape=noise_shape, seed=seed)
        self.dense_attend = Dense(units=units, use_bias=use_bias, activation=activation, **kernel_args)
        self.dense_align = Dense(1, activation="linear", use_bias=use_bias, **kernel_args)
        self.dense_e = Dense(units=units_edge, activation=activation, use_bias=use_bias, **kernel_args)
        self.pool_attention = PoolingLocalEdgesAttention()
        self.final_activ = Activation(activation=activation_last,
                                      activity_regularizer=activity_regularizer)

    def build(self, input_shape):
        """Build layer."""
        super(HamNaiveDynMessage, self).build(input_shape)

    def call(self, inputs, **kwargs):
        """Forward pass.

        Args:
            inputs: [hv_ftr, he_ftr, p_ftr, q_ftr, edge_index]

                - hv_ftr (tf.RaggedTensor): Node features of shape `(batch, [N], F)`
                - he_ftr (tf.RaggedTensor): Edge features of shape `(batch, [M], F)`
                - p_ftr (tf.RaggedTensor): Momentum node features of shape `(batch, [N], F)`
                - q_ftr (tf.RaggedTensor): Positional node features of shape `(batch, [N], F)`
                - edge_index (tf.RaggedTensor): Edge connection index list of shape `(batch, [M], 2)`

        Returns:
            list: [mv_ftr, me_ftr]

                - mv_ftr (tf.RaggedTensor): Node feature updates of shape `(batch, [N], F)`
                - me_ftr (tf.RaggedTensor): Edge feature updates of shape `(batch, [M], F)`
        """
        hv_ftr, he_ftr, p_ftr, q_ftr, edi = inputs
        if self.use_dropout:
            hv_ftr = self.dropout_layer(hv_ftr, **kwargs)
            he_ftr = self.dropout_layer(he_ftr, **kwargs)
        hv_u_ftr, hv_v_ftr = self.gather_v([hv_ftr, edi], **kwargs)
        q_u_ftr, q_v_ftr = self.gather_p([q_ftr, edi], **kwargs)
        p_u_ftr, p_v_ftr = self.gather_q([p_ftr, edi], **kwargs)
        p_uv_ftr = self.lazy_sub_p([p_v_ftr, p_u_ftr], **kwargs)
        q_uv_ftr = self.lazy_sub_p([q_v_ftr, q_u_ftr], **kwargs)

        attend_ftr = self.dense_attend(hv_v_ftr, **kwargs)

        align_ftr = self.lay_concat_align([p_uv_ftr, q_uv_ftr, he_ftr], **kwargs)
        align_ftr = self.dense_align(align_ftr, **kwargs)
        mv_ftr = self.pool_attention([hv_ftr, attend_ftr, align_ftr, edi], **kwargs)
        mv_ftr = self.final_activ(mv_ftr, **kwargs)

        me_ftr = self.lay_concat_edge([hv_u_ftr, p_uv_ftr, q_uv_ftr, hv_v_ftr], **kwargs)
        me_ftr = self.dense_e(me_ftr)

        return mv_ftr, me_ftr

    def get_config(self):
        """Update layer config."""
        config = super(HamNaiveDynMessage, self).get_config()
        config.update({"use_bias": self.use_bias, "units": self.units, "use_dropout": self.use_dropout})
        conf_sub = self.dense_attend.get_config()
        for x in ["kernel_regularizer", "activity_regularizer", "bias_regularizer", "kernel_constraint",
                  "bias_constraint", "kernel_initializer", "bias_initializer", "activation"]:
            config.update({x: conf_sub[x]})
        if self.use_dropout:
            conf_drop = self.dropout_layer.get_config()
            for x in ["rate", "noise_shape", "seed"]:
                config.update({x: conf_drop[x]})
        conf_last = self.final_activ.get_config()
        config.update({"activation_last": conf_last["activation"]})
        return config
