#
# 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
from .convolve import SparseDropoutGraphConvActivation
from .data import Data
from typing import List, \
    Union, \
    Callable
from collections import defaultdict


class Layer(torch.nn.Module):
    def __init__(self, output_dim: Union[int, List[int]], **kwargs) -> None:
        super().__init__(**kwargs)
        self.output_dim = output_dim


class InputLayer(Layer):
    def __init__(self, data: Data, output_dim: Union[int, List[int]]= None, **kwargs) -> None:
        output_dim = output_dim or \
            list(map(lambda a: a.count, data.node_types))
        if not isinstance(output_dim, list):
            output_dim = [output_dim,] * len(data.node_types)

        super().__init__(output_dim, **kwargs)
        self.data = data
        self.node_reps = None
        self.build()

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

    def forward(self) -> List[torch.nn.Parameter]:
        return self.node_reps

    def __repr__(self) -> str:
        s = ''
        s += 'GNN input layer with output_dim: %s\n' % self.output_dim
        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(Layer):
    def __init__(self,
        data: Data,
        previous_layer: Layer,
        output_dim: Union[int, List[int]],
        keep_prob: float = 1.,
        rel_activation: Callable[[torch.Tensor], torch.Tensor] = lambda x: x,
        layer_activation: Callable[[torch.Tensor], torch.Tensor] = torch.nn.functional.relu,
        **kwargs):
        if not isinstance(output_dim, list):
            output_dim = [ output_dim ] * len(data.node_types)
        super().__init__(output_dim, **kwargs)
        self.data = data
        self.previous_layer = previous_layer
        self.input_dim = previous_layer.output_dim
        self.keep_prob = keep_prob
        self.rel_activation = rel_activation
        self.layer_activation = layer_activation
        self.convolutions = None
        self.build()

    def build(self):
        self.convolutions = {}
        for (node_type_row, node_type_column) in self.data.relation_types.keys():
            adjacency_matrices = \
                self.data.get_adjacency_matrices(node_type_row, node_type_column)
            self.convolutions[node_type_row, node_type_column] = SparseMultiDGCA(self.input_dim,
                self.output_dim, adjacency_matrices,
                self.keep_prob, self.rel_activation)

        # for node_type_row, node_type_col in enumerate(self.data.node_
        #    if rt.node_type_row == i or rt.node_type_col == i:

    def __call__(self):
        prev_layer_repr = self.previous_layer()
        next_layer_repr = defaultdict(list)

        for (nt_row, nt_col), rel in self.data.relation_types.items():
            conv = SparseDropoutGraphConvActivation(self.input_dim[nt_col],
                self.output_dim[nt_row], rel.adjacency_matrix,
                self.keep_prob, self.rel_activation)
            next_layer_repr[nt_row].append(conv)

            conv = SparseDropoutGraphConvActivation(self.input_dim[nt_row],
                self.output_dim[nt_col], rel.adjacency_matrix.transpose(0, 1),
                self.keep_prob, self.rel_activation)
            next_layer_repr[nt_col].append(conv)

        next_layer_repr = list(map(sum, next_layer_repr))
        return next_layer_repr


        #for i, nt in enumerate(self.data.node_types):
        #    new_repr = []
        #    for nt_row, nt_col in self.data.relation_types.keys():
        #        if nt_row != i and nt_col != i:
        #            continue
        #        if nt_row == i:
        #            x = prev_layer_repr[nt_col]
        #        else:
        #            x = prev_layer_repr[nt_row]
        #        conv = self.convolutions[key]
        #        new_repr.append(conv(x))
        #    new_repr = sum(new_repr)
        #    new_layer_repr.append(new_repr)
        # return new_layer_repr