"""Submodule providing graph convolutional layer.

# References
The layer is implemented as described in [Semi-Supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907).

In this version of the implementation, we allow for batch sizes of arbitrary size.
"""
from typing import Tuple, Union, Dict, List
import tensorflow as tf
from tensorflow.python.ops import embedding_ops  # pylint: disable=import-error,no-name-in-module
from tensorflow.keras.layers import Flatten, Layer  # pylint: disable=import-error,no-name-in-module


class EmbeddingLookup(Layer):
    """Layer implementing simple embedding lookup layer."""

    def __init__(
        self,
        **kwargs: Dict
    ):
        """Create new GraphConvolution layer.

        Parameters
        ----------------------
        **kwargs: Dict,
            Kwargs to pass to the parent Layer class.
        """
        super().__init__(**kwargs)
        self._flatten_layer = None

    def build(self, input_shape) -> None:
        """Build the embedding lookup layer.

        Parameters
        ------------------------------
        input_shape
            Shape of the output of the previous layer.
        """
        self._flatten_layer = Flatten()
        super().build(input_shape)

    def call(
        self,
        inputs: Tuple[tf.Tensor],
    ) -> tf.Tensor:
        """Returns called embeddingg lookup.

        Parameters
        ---------------------------
        inputs: Tuple[tf.Tensor],
        """
        node_ids, node_features = inputs
        return self._flatten_layer(tf.nn.embedding_lookup(
            node_features,
            ids=node_ids
        ))