import tensorflow as tf
from kgcnn.layers.base import GraphBaseLayer
from kgcnn.layers.modules import Dense

ks = tf.keras


@tf.keras.utils.register_keras_serializable(package='kgcnn', name='TrafoEdgeNetMessages')
class TrafoEdgeNetMessages(GraphBaseLayer):
    """Make message from edges for a linear transformation, i.e. matrix multiplication.
    The actual matrix is not a trainable weight of this layer but generated by a dense layer.
    This was proposed by `NMPNN <http://arxiv.org/abs/1704.01212>`_ .

    Args:
        target_shape (int): Target shape for message matrix.
    """

    def __init__(self, target_shape,
                 activation="linear",
                 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',
                 **kwargs):
        """Initialize layer."""
        super(TrafoEdgeNetMessages, self).__init__(**kwargs)
        self.target_shape = target_shape
        self._units_out = int(target_shape[0])
        self._units_in = int(target_shape[1])
        self.lay_dense = Dense(units=self._units_out * self._units_in,
                               activation=activation, use_bias=use_bias,
                               kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer,
                               activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint,
                               bias_constraint=bias_constraint, kernel_initializer=kernel_initializer,
                               bias_initializer=bias_initializer)

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

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

        Args:
            inputs (tf.RaggedTensor): Messages embeddings or messages (batch, [M], F)

        Returns:
            tf.RaggedTensor: Messages in matrix for multiplication of shape (batch, [M], F, F)
        """
        self.assert_ragged_input_rank(inputs)
        up_scale = self.lay_dense(inputs, **kwargs)
        dens_trafo, trafo_part = up_scale.values, up_scale.row_splits
        dens_m = tf.reshape(dens_trafo, (ks.backend.shape(dens_trafo)[0], self._units_out, self._units_in))
        out = tf.RaggedTensor.from_row_splits(dens_m, trafo_part, validate=self.ragged_validate)
        return out

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


@tf.keras.utils.register_keras_serializable(package='kgcnn', name='MatMulMessages')
class MatMulMessages(GraphBaseLayer):
    """Linear transformation of edges or messages, i.e. matrix multiplication.

    The message dimension must be suitable for matrix multiplication. The actual matrix is not a trainable weight of
    this layer but passed as input.
    This was proposed by `NMPNN <http://arxiv.org/abs/1704.01212>`_ .

    """

    def __init__(self, **kwargs):
        """Initialize layer."""
        super(MatMulMessages, self).__init__(**kwargs)

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

    def call(self, inputs, **kwargs):
        """Forward pass. Operates on values without checking splits of the ragged dimension.

        Args:
            inputs (list): [trafo_mat, edges]

                - trafo_mat (tf.RaggedTensor): Transformation matrix for each message of shape (batch, [M], F', F).
                - edges (tf.RaggedTensor): Edge embeddings or messages (batch, [M], F)
            
        Returns:
            tf.RaggedTensor: Transformation of messages by matrix multiplication of shape (batch, [M], F')
        """
        self.assert_ragged_input_rank(inputs)
        dens_trafo, trafo_part = inputs[0].values, inputs[0].row_splits
        dens_e, epart = inputs[1].values, inputs[1].row_splits
        out = tf.keras.backend.batch_dot(dens_trafo, dens_e)
        out = tf.RaggedTensor.from_row_splits(out, epart, validate=self.ragged_validate)
        return out

    def get_config(self):
        """Update layer config."""
        config = super(MatMulMessages, self).get_config()
        return config
