import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl.function as fn

"""
    GIN: Graph Isomorphism Networks
    HOW POWERFUL ARE GRAPH NEURAL NETWORKS? (Keyulu Xu, Weihua Hu, Jure Leskovec and Stefanie Jegelka, ICLR 2019)
    https://arxiv.org/pdf/1810.00826.pdf
"""

class GINLayer(nn.Module):
    """
    [!] code adapted from dgl implementation of GINConv

    Parameters
    ----------
    apply_func : callable activation function/layer or None
        If not None, apply this function to the updated node feature,
        the :math:`f_\Theta` in the formula.
    aggr_type :
        Aggregator type to use (``sum``, ``max`` or ``mean``).
    out_dim :
        Rquired for batch norm layer; should match out_dim of apply_func if not None.
    dropout :
        Required for dropout of output features.
    graph_norm : 
        boolean flag for output features normalization w.r.t. graph sizes.
    batch_norm :
        boolean flag for batch_norm layer.
    residual :
        boolean flag for using residual connection.
    init_eps : optional
        Initial :math:`\epsilon` value, default: ``0``.
    learn_eps : bool, optional
        If True, :math:`\epsilon` will be a learnable parameter.
    
    """
    def __init__(self, apply_func, aggr_type, dropout, graph_norm, batch_norm, residual=False, init_eps=0, learn_eps=False):
        super().__init__()
        self.apply_func = apply_func
        
        if aggr_type == 'sum':
            self._reducer = fn.sum
        elif aggr_type == 'max':
            self._reducer = fn.max
        elif aggr_type == 'mean':
            self._reducer = fn.mean
        else:
            raise KeyError('Aggregator type {} not recognized.'.format(aggr_type))
            
        self.graph_norm = graph_norm
        self.batch_norm = batch_norm
        self.residual = residual
        self.dropout = dropout
        
        in_dim = apply_func.mlp.input_dim
        out_dim = apply_func.mlp.output_dim
        
        if in_dim != out_dim:
            self.residual = False
            
        # to specify whether eps is trainable or not.
        if learn_eps:
            self.eps = torch.nn.Parameter(torch.FloatTensor([init_eps]))
        else:
            self.register_buffer('eps', torch.FloatTensor([init_eps]))
            
        self.bn_node_h = nn.BatchNorm1d(out_dim)

    def forward(self, g, h, snorm_n):
        h_in = h # for residual connection
        
        g = g.local_var()
        g.ndata['h'] = h
        g.update_all(fn.copy_u('h', 'm'), self._reducer('m', 'neigh'))
        h = (1 + self.eps) * h + g.ndata['neigh']
        if self.apply_func is not None:
            h = self.apply_func(h)

        if self.graph_norm:
            h = h* snorm_n # normalize activation w.r.t. graph size
        
        if self.batch_norm:
            h = self.bn_node_h(h) # batch normalization  
        
        h = F.relu(h) # non-linear activation
        
        if self.residual:
            h = h_in + h # residual connection
        
        h = F.dropout(h, self.dropout, training=self.training)
        
        return h
    
    
class ApplyNodeFunc(nn.Module):
    """
        This class is used in class GINNet
        Update the node feature hv with MLP
    """
    def __init__(self, mlp):
        super().__init__()
        self.mlp = mlp

    def forward(self, h):
        h = self.mlp(h)
        return h


class MLP(nn.Module):
    """MLP with linear output"""
    def __init__(self, num_layers, input_dim, hidden_dim, output_dim):

        super().__init__()
        self.linear_or_not = True  # default is linear model
        self.num_layers = num_layers
        self.output_dim = output_dim
        self.input_dim = input_dim

        if num_layers < 1:
            raise ValueError("number of layers should be positive!")
        elif num_layers == 1:
            # Linear model
            self.linear = nn.Linear(input_dim, output_dim)
        else:
            # Multi-layer model
            self.linear_or_not = False
            self.linears = torch.nn.ModuleList()
            self.batch_norms = torch.nn.ModuleList()

            self.linears.append(nn.Linear(input_dim, hidden_dim))
            for layer in range(num_layers - 2):
                self.linears.append(nn.Linear(hidden_dim, hidden_dim))
            self.linears.append(nn.Linear(hidden_dim, output_dim))

            for layer in range(num_layers - 1):
                self.batch_norms.append(nn.BatchNorm1d((hidden_dim)))

    def forward(self, x):
        if self.linear_or_not:
            # If linear model
            return self.linear(x)
        else:
            # If MLP
            h = x
            for i in range(self.num_layers - 1):
                h = F.relu(self.batch_norms[i](self.linears[i](h)))
            return self.linears[-1](h)
[((69, 25, 69, 48), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', ({(69, 40, 69, 47): 'out_dim'}, {}), '(out_dim)', True, 'import torch.nn as nn\n'), ((87, 12, 87, 21), 'torch.nn.functional.relu', 'F.relu', ({(87, 19, 87, 20): 'h'}, {}), '(h)', True, 'import torch.nn.functional as F\n'), ((92, 12, 92, 62), 'torch.nn.functional.dropout', 'F.dropout', (), '', True, 'import torch.nn.functional as F\n'), ((76, 21, 76, 40), 'dgl.function.copy_u', 'fn.copy_u', ({(76, 31, 76, 34): '"""h"""', (76, 36, 76, 39): '"""m"""'}, {}), "('h', 'm')", True, 'import dgl.function as fn\n'), ((65, 42, 65, 71), 'torch.FloatTensor', 'torch.FloatTensor', ({(65, 60, 65, 70): '[init_eps]'}, {}), '([init_eps])', False, 'import torch\n'), ((67, 40, 67, 69), 'torch.FloatTensor', 'torch.FloatTensor', ({(67, 58, 67, 68): '[init_eps]'}, {}), '([init_eps])', False, 'import torch\n'), ((125, 26, 125, 58), 'torch.nn.Linear', 'nn.Linear', ({(125, 36, 125, 45): 'input_dim', (125, 47, 125, 57): 'output_dim'}, {}), '(input_dim, output_dim)', True, 'import torch.nn as nn\n'), ((129, 27, 129, 48), 'torch.nn.ModuleList', 'torch.nn.ModuleList', ({}, {}), '()', False, 'import torch\n'), ((130, 31, 130, 52), 'torch.nn.ModuleList', 'torch.nn.ModuleList', ({}, {}), '()', False, 'import torch\n'), ((132, 32, 132, 64), 'torch.nn.Linear', 'nn.Linear', ({(132, 42, 132, 51): 'input_dim', (132, 53, 132, 63): 'hidden_dim'}, {}), '(input_dim, hidden_dim)', True, 'import torch.nn as nn\n'), ((135, 32, 135, 65), 'torch.nn.Linear', 'nn.Linear', ({(135, 42, 135, 52): 'hidden_dim', (135, 54, 135, 64): 'output_dim'}, {}), '(hidden_dim, output_dim)', True, 'import torch.nn as nn\n'), ((134, 36, 134, 69), 'torch.nn.Linear', 'nn.Linear', ({(134, 46, 134, 56): 'hidden_dim', (134, 58, 134, 68): 'hidden_dim'}, {}), '(hidden_dim, hidden_dim)', True, 'import torch.nn as nn\n'), ((138, 40, 138, 68), 'torch.nn.BatchNorm1d', 'nn.BatchNorm1d', ({(138, 56, 138, 66): 'hidden_dim'}, {}), '(hidden_dim)', True, 'import torch.nn as nn\n')]