#
# This module implements a single layer of the Decagon
# model. This is going to be already quite complex, as
# we will be using all the graph convolutional building
# blocks.
#
# h_{i}^(k+1) = ϕ(∑_r ∑_{j∈N{r}^{i}} c_{r}^{ij} * \
#   W_{r}^(k) h_{j}^{k} + c_{r}^{i} h_{i}^(k))
#
# N{r}^{i} - set of neighbors of node i under relation r
# W_{r}^(k) - relation-type specific weight matrix
# h_{i}^(k) - hidden state of node i in layer k
#   h_{i}^(k)∈R^{d(k)} where d(k) is the dimensionality
#   of the representation in k-th layer
# ϕ - activation function
# c_{r}^{ij} - normalization constants
#   c_{r}^{ij} = 1/sqrt(|N_{r}^{i}| |N_{r}^{j}|)
# c_{r}^{i} - normalization constants
#   c_{r}^{i} = 1/|N_{r}^{i}|
#


import torch


class InputLayer(torch.nn.Module):
    def __init__(self, data, dimensionality=32, **kwargs):
        super().__init__(**kwargs)
        self.data = data
        self.dimensionality = dimensionality
        self.node_reps = None
        self.build()

    def build(self):
        self.node_reps = []
        for i, nt in enumerate(self.data.node_types):
            reps = torch.rand(nt.count, self.dimensionality)
            reps = torch.nn.Parameter(reps)
            self.register_parameter('node_reps[%d]' % i, reps)
            self.node_reps.append(reps)

    def forward(self):
        return self.node_reps

    def __repr__(self):
        s = ''
        s += 'GNN input layer with dimensionality: %d\n' % self.dimensionality
        s += '  # of node types: %d\n' % len(self.data.node_types)
        for nt in self.data.node_types:
            s += '    - %s (%d)\n' % (nt.name, nt.count)
        return s.strip()


class DecagonLayer(torch.nn.Module):
    def __init__(self, data, **kwargs):
        super().__init__(**kwargs)
        self.data = data

    def __call__(self, previous_layer):
        pass