from neuralogic.core.constructs.metadata import Metadata
from neuralogic.core.constructs.function import Activation, Aggregation
from neuralogic.core.constructs.factories import R, V
from neuralogic.nn.module.module import Module


class ResGatedGraphConv(Module):
    r"""
    Residual Gated Graph Convolutional layer from `"Residual Gated Graph ConvNets" <https://arxiv.org/abs/1711.07553>`_.
    Which can be expressed as:

    .. math::

        \mathbf{x}^{\prime}_i = act(\mathbf{W}_1 \mathbf{x}_i +
         {agg}_{j \in \mathcal{N}(i)}(\eta_{i,j} \odot \mathbf{W}_2 \mathbf{x}_j))

    .. math::

        \mathbf{\eta}_{i,j} = gating\_act(\mathbf{W}_3 \mathbf{x}_i + \mathbf{W}_4 \mathbf{x}_j)

    Where *act* is an activation function, *agg* aggregation function, *gating_act* is a gating activation function and
    :math:`W_n` are learnable parameters. This equation is translated into the logic form as:

    .. code:: logtalk

        (R.<output_name>__gate(V.I, V.J) <= (R.<feature_name>(V.I)[<W>], R.<feature_name>(V.J)[<W>])) | [Activation.IDENTITY]
        R.<output_name>__gate / 2 | [<activation>]

        (R.<output_name>(V.I) <= R.<feature_name>(V.I)[<W>]) | [Activation.IDENTITY]
        (R.<output_name>(V.I) <= (
            R.<output_name>__gate(V.I, V.J), R.<feature_name>(V.J)[<W>], R.<edge_name>(V.J, V.I))
        ) | Metadata(activation="elementproduct-identity", aggregation=<aggregation>)

        R.<output_name> / 1 | [<activation>]

    Examples
    --------

    The whole computation of this module (parametrized as :code:`ResGatedGraphConv(1, 2, "h1", "h0", "_edge")`)
    is as follows:

    .. code:: logtalk

        metadata = Metadata(activation="elementproduct-identity", aggregation=Aggregation.SUM)

        (R.h1__gate(V.I, V.J) <= (R.h0(V.I)[2, 1], R.h0(V.J)[2, 1])) | [Activation.IDENTITY]
        R.h1__gate / 2 | [Activation.SIGMOID]

        (R.h1(V.I) <= R.h0(V.I)[2, 1]) | [Activation.IDENTITY]
        (R.h1(V.I) <= (R.h1__gate(V.I, V.J), R.h0(V.J)[2, 1], R._edge(V.J, V.I))) | metadata
        R.h1 / 1 | [Activation.IDENTITY]

    Parameters
    ----------

    in_channels : int
        Input feature size.
    out_channels : int
        Output feature size.
    output_name : str
        Output (head) predicate name of the module.
    feature_name : str
        Feature predicate name to get features from.
    edge_name : str
        Edge predicate name to use for neighborhood relations.
    gating_activation : Activation
        Gating activation function.
        Default: ``Activation.SIGMOID``
    activation : Activation
        Activation function of the output.
        Default: ``Activation.IDENTITY``
    aggregation : Aggregation
        Aggregation function of nodes' neighbors.
        Default: ``Aggregation.SUM``

    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        output_name: str,
        feature_name: str,
        edge_name: str,
        gating_activation: Activation = Activation.SIGMOID,
        activation: Activation = Activation.IDENTITY,
        aggregation: Aggregation = Aggregation.SUM,
    ):
        self.output_name = output_name
        self.feature_name = feature_name
        self.edge_name = edge_name

        self.in_channels = in_channels
        self.out_channels = out_channels

        self.aggregation = aggregation
        self.activation = activation
        self.gating_activation = gating_activation

    def __call__(self):
        head = R.get(self.output_name)(V.I)
        feature = R.get(self.feature_name)
        gate = R.get(f"{self.output_name}__gate")

        prod_metadata = Metadata(activation="elementproduct-identity", aggregation=self.aggregation)
        w = self.out_channels, self.in_channels

        return [
            (gate(V.I, V.J) <= (feature(V.I)[w], feature(V.J)[w])) | [Activation.IDENTITY],
            gate / 2 | Metadata(activation=self.gating_activation),
            (head <= feature(V.I)[w]) | [Activation.IDENTITY],
            (head <= (gate(V.I, V.J), feature(V.J)[w], R.get(self.edge_name)(V.J, V.I))) | prod_metadata,
            R.get(self.output_name) / 1 | Metadata(activation=self.activation),
        ]
