import torch
import torch.nn as nn
import torch.nn.functional as F
from layers import MultiHeadAttention


class MultiHeadAttentionLayer(nn.Module):
    """Feed-Forward Sublayer: fully-connected Feed-Forward network,
    built based on MHA vectors from MultiHeadAttention layer with skip-connections

        Args:
            num_heads: number of attention heads in MHA layers.
            input_dim: embedding size that will be used as d_model in MHA layers.
            feed_forward_hidden: number of neuron units in each FF layer.

        Call arguments:
            x: batch of shape (batch_size, n_nodes, node_embedding_size).
            mask: mask for MHA layer

        Returns:
               outputs of shape (batch_size, n_nodes, input_dim)

    """

    def __init__(self, input_dim, num_heads, feed_forward_hidden=512, **kwargs):
        super().__init__(**kwargs)
        self.mha = MultiHeadAttention(n_heads=num_heads, d_model=input_dim)

        self.ff1 = nn.Linear(input_dim, feed_forward_hidden)
        self.ff2 = nn.Linear(feed_forward_hidden, input_dim)

    def forward(self, x, mask=None):
        mha_out = self.mha(x, x, x, mask)
        sc1_out = torch.add(x, mha_out)
        tanh1_out = torch.tanh(sc1_out)

        ff1_out = self.ff1(tanh1_out)
        relu1_out = F.relu(ff1_out)
        ff2_out = self.ff2(relu1_out)
        sc2_out = torch.add(tanh1_out, ff2_out)
        tanh2_out = torch.tanh(sc2_out)

        return tanh2_out


class GraphAttentionEncoder(nn.Module):
    """Graph Encoder, which uses MultiHeadAttentionLayer sublayer.

    Args:
        input_dim: embedding size that will be used as d_model in MHA layers.
        num_heads: number of attention heads in MHA layers.
        num_layers: number of attention layers that will be used in encoder.
        feed_forward_hidden: number of neuron units in each FF layer.

    Call arguments:
        x: tuples of 3 tensors:  (batch_size, 2), (batch_size, n_nodes-1, 2), (batch_size, n_nodes-1)
        First tensor contains coordinates for depot, second one is for coordinates of other nodes,
        Last tensor is for normalized demands for nodes except depot

        mask: mask for MHA layer

    Returns:
           Embedding for all nodes + mean embedding for graph.
           Tuples ((batch_size, n_nodes, input_dim), (batch_size, input_dim))
    """

    def __init__(self, input_dim, num_heads, num_layers, feed_forward_hidden=512):
        super().__init__()

        self.input_dim = input_dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.feed_forward_hidden = feed_forward_hidden

        # initial embeddings (batch_size, n_nodes-1, 2) --> (batch-size, input_dim), separate for depot and other nodes
        self.init_embed_depot = nn.Linear(
            2, self.input_dim
        )  # nn.Linear(2, embedding_dim)
        self.init_embed = nn.Linear(3, self.input_dim)

        self.mha_layers = [
            MultiHeadAttentionLayer(
                self.input_dim, self.num_heads, self.feed_forward_hidden
            )
            for _ in range(self.num_layers)
        ]
        self.mha_layers = nn.ModuleList(self.mha_layers)

    def forward(self, x, mask=None, cur_num_nodes=None):

        x = torch.cat(
            (
                self.init_embed_depot(x[0])[
                    :, None, :
                ],  # (batch_size, 2) --> (batch_size, 1, 2)
                self.init_embed(
                    torch.cat((x[1], x[2][:, :, None]), -1)
                ),  # (batch_size, n_nodes-1, 2) + (batch_size, n_nodes-1)
            ),
            1,
        )  # (batch_size, n_nodes, input_dim)

        # stack attention layers
        for i in range(self.num_layers):
            x = self.mha_layers[i](x, mask)

        if mask is not None:
            output = (x, torch.sum(x, 1) / cur_num_nodes)
        else:
            output = (x, torch.mean(x, 1))

        return output  # (embeds of nodes, avg graph embed)=((batch_size, n_nodes, input), (batch_size, input_dim))
