import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class GraphAttentionLayer(nn.Module):
    """
    Simple GAT layer, similar to https://arxiv.org/abs/1710.10903
    """

    def __init__(self, in_features, out_features, dropout, alpha, concat=True):
        super(GraphAttentionLayer, self).__init__()
        self.dropout = dropout
        self.in_features = in_features
        self.out_features = out_features
        self.alpha = alpha
        self.concat = concat

        self.W = nn.Parameter(torch.empty(size=(in_features, out_features)))
        nn.init.xavier_uniform_(self.W.data, gain=1.414)
        self.a = nn.Parameter(torch.empty(size=(2 * out_features, 1)))
        nn.init.xavier_uniform_(self.a.data, gain=1.414)

        self.leakyrelu = nn.LeakyReLU(self.alpha)

    def forward(self, h, adj):
        # h.shape: (N, in_features), Wh.shape: (N, out_features)
        # Wh = torch.mm(h, self.W)
        # i:batch, j:horizon, k: in_features, l: num_sensors, m: out_features
        Wh = torch.einsum('ijkl, km -> ijml', h, self.W)
        e = self._prepare_attentional_mechanism_input(Wh)

        attention = torch.where(adj > 0, e, -9e15*torch.ones_like(e))
        attention = F.softmax(attention, dim=-1)
        attention = F.dropout(attention, self.dropout, training=self.training)
        h_prime = torch.einsum('ijkk, ijlk -> ijlk', attention, Wh)

        if self.concat:
            return F.elu(h_prime)
        else:
            return h_prime

    def _prepare_attentional_mechanism_input(self, Wh):
        Wh1 = torch.einsum('ijkl, km -> ijml', Wh,
                           self.a[:self.out_features, :])
        Wh2 = torch.einsum('ijkl, km -> ijml', Wh,
                           self.a[self.out_features:, :])
        # broadcast add
        e = Wh1 + Wh2.transpose(-1, -2)
        return self.leakyrelu(e)

    def __repr__(self):
        return self.__class__.__name__ + ' (' + str(
            self.in_features) + ' -> ' + str(self.out_features) + ')'
