import math

import torch
import torch.nn as nn
from torch.nn.modules.module import Module
from GNN.GCN_layer  import GraphConvolution

class GraphResConvolution(Module):
    """
    Simple GCN layer, similar to https://arxiv.org/abs/1609.02907
    """

    def __init__(self, state_dim, name=''):
        super(GraphResConvolution, self).__init__()
        self.state_dim = state_dim

        self.gcn_1 = GraphConvolution(state_dim, '%s_1' % name)
        self.gcn_2 = GraphConvolution(state_dim, '%s_2' % name)

        self.relu1 = nn.ReLU()
        self.relu2 = nn.ReLU()

        self.name = name

    def forward(self, input, adj):

        output_1 = self.gcn_1(input, adj)
        output_1_relu = self.relu1(output_1)

        output_2 =  self.gcn_2(output_1_relu, adj)

        output_2_res = output_2 + input

        output = self.relu2(output_2_res)

        return output



    def __repr__(self):
        return self.__class__.__name__ + ' (' +  self.name + ')'
[((17, 21, 17, 63), 'GNN.GCN_layer.GraphConvolution', 'GraphConvolution', ({(17, 38, 17, 47): 'state_dim', (17, 49, 17, 62): "'%s_1' % name"}, {}), "(state_dim, '%s_1' % name)", False, 'from GNN.GCN_layer import GraphConvolution\n'), ((18, 21, 18, 63), 'GNN.GCN_layer.GraphConvolution', 'GraphConvolution', ({(18, 38, 18, 47): 'state_dim', (18, 49, 18, 62): "'%s_2' % name"}, {}), "(state_dim, '%s_2' % name)", False, 'from GNN.GCN_layer import GraphConvolution\n'), ((20, 21, 20, 30), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((21, 21, 21, 30), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', True, 'import torch.nn as nn\n')]