import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from torch.nn.parameter import Parameter

class BatchAGC(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super(BatchAGC, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(torch.Tensor(in_features, out_features))
        if bias:
            self.bias = Parameter(torch.Tensor(out_features))
            init.constant_(self.bias, 0)
        else:
            self.register_parameter('bias', None)
        init.xavier_uniform_(self.weight)

    def forward(self, x, adj):
        expand_weight = self.weight.expand(x.shape[0], -1, -1)
        support = torch.bmm(x, expand_weight)
        output = torch.bmm(adj, support)
        if self.bias is not None:
            return output + self.bias
        else:
            return output

class BatchFiGNN(nn.Module):
    def __init__(self, f_in, f_out, out_channels):
        super(BatchFiGNN, self).__init__()
        # Edge Weights
        self.a_src = Parameter(torch.Tensor(f_in, 1))
        self.a_dst = Parameter(torch.Tensor(f_in, 1))
        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)
        self.softmax = nn.Softmax(dim=-1)
        # Transformation
        self.w = Parameter(torch.Tensor(f_in, f_out))
        self.bias = Parameter(torch.Tensor(f_out))
        # State Update by GRU
        self.rnn = torch.nn.GRUCell(f_out, f_out, bias=True)
        # Attention Pooling
        self.mlp_1 = nn.Linear(f_out, out_channels, bias=True)
        self.mlp_2 = nn.Linear(f_out, 1, bias=True)

        init.xavier_uniform_(self.w)
        init.constant_(self.bias, 0)
        init.xavier_uniform_(self.a_src)
        init.xavier_uniform_(self.a_dst)

    def forward(self, h, adj, steps):
        bs, n = h.size()[:2]
        ## Edge Weights  
        attn_src = torch.matmul(h, self.a_src)
        attn_dst = torch.matmul(h, self.a_dst)
        attn = attn_src.expand(-1, -1, n) + \
            attn_dst.expand(-1, -1, n).permute(0, 2, 1)
        attn = self.leaky_relu(attn)
        mask = torch.eye(adj.size()[-1]).unsqueeze(0).cuda()
        mask = 1 - mask
        attn = attn * mask
        attn = self.softmax(attn)
        ## State Update
        s = h
        for _ in range(steps):
            ## Transformation
            a = torch.matmul(s, self.w)
            a = torch.matmul(attn, a) + self.bias
            ## GRU
            s = self.rnn(s.view(-1, s.size()[-1]), a.view(-1, a.size()[-1]))
            s = h.view(h.size()) + h
        ## Attention Pooling
        output = self.mlp_1(s)
        weight = self.mlp_2(s).permute(0, 2, 1)
        output = torch.matmul(weight, output).squeeze()
        return output

class BatchGAT(nn.Module):
    def __init__(self, n_head, f_in, f_out, attn_dropout, bias=True):
        super(BatchGAT, self).__init__()
        self.n_head = n_head
        self.w = Parameter(torch.Tensor(n_head, f_in, f_out))
        self.a_src = Parameter(torch.Tensor(n_head, f_out, 1))
        self.a_dst = Parameter(torch.Tensor(n_head, f_out, 1))

        self.leaky_relu = nn.LeakyReLU(negative_slope=0.2)
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(attn_dropout)
        if bias:
            self.bias = Parameter(torch.Tensor(f_out))
            init.constant_(self.bias, 0)
        else:
            self.register_parameter('bias', None)

        init.xavier_uniform_(self.w)
        init.xavier_uniform_(self.a_src)
        init.xavier_uniform_(self.a_dst)

    def forward(self, h, adj):
        bs, n = h.size()[:2]
        h_prime = torch.matmul(h.unsqueeze(1), self.w)
        attn_src = torch.matmul(torch.tanh(h_prime), self.a_src)
        attn_dst = torch.matmul(torch.tanh(h_prime), self.a_dst)
        attn = attn_src.expand(-1, -1, -1, n) + \
            attn_dst.expand(-1, -1, -1, n).permute(0, 1, 3, 2)
        attn = self.leaky_relu(attn)
        attn = self.softmax(attn)
        attn = self.dropout(attn)
        output = torch.matmul(attn, h_prime)
        if self.bias is not None:
            return output + self.bias
        else:
            return output