"""
This file contains the self attention and the multi attention modules

Description :
The Network being implemented here produces context aware representation by applying attention to
each pair of tokens from the input sequence

The idea is to use this network to predict the top n goals to use
for resubstitution in the Hindsight experience replay

"""
import torch
import torch.nn as nn
import torch.nn.functional as F


def multiplicative_fn( x, q, linear_x, linear_q):
    """

    :param x: Input token
    :param q: Query vector
    :param linear_x: Weight matrix for x
    :param linear_q: Weight matrix for query
    :return: compatibility function
    """

    # Return the cosine similarity
    x_ = linear_x(x)
    q_ = linear_q(q)
    f = F.cosine_similarity(x_, q_)
    return f


def additive_fn( x, q, linear_x, linear_q, out_linear, num_hidden, activation):
    """

    :param x: Input token
    :param q: Query vector
    :param linear_x: Weight matrix for x
    :param linear_q: Weight matrix for query
    :param num_hidden: Number of hidden layers
    :param activation: The Non linear activation function to use
    :return: compatibility function
    """

    # Additive attention achieves better empirical performance
    x_ = linear_x(x)
    q_ = linear_q(q)
    o = out_linear(activation(x_ + q_))
    return o


class VanillaAttention(nn.Module):
    """
    The standard attention module
    Vanilla Attention computes alignments scores of the query to each
    of the input tokens in an input sequence.
    """
    def __init__(self, input_dim, query_dim, num_hidden, embedding_dim,
                 activation, output_features,
                 use_additive=True, save_attention=False,
                 attention_dict=None, name=None):
        """

        :param input_dim: The input dimension of the sequence or length of the sequence
        :param query_dim: The dimension of the query embedding - Current State || Desired Goal Vector
        :param num_hidden: The number of hidden layers in the network
        :param embedding_dim: The dimension of each embedding - State || Achieved Goal Vector
        :param activation: The non-linear activation function to use
        :param use_additive_fn: Boolean to determine whether to use additive or multiplicative attention
        """
        super(VanillaAttention, self).__init__()
        self.input_dim = input_dim # n
        self.query_dim = query_dim # q
        self.num_hidden = num_hidden # h
        self.embedding_dim = embedding_dim # e
        self.output = output_features
        self.additive = use_additive
        self.activation = activation
        self.save_attention = save_attention
        self.attention_dict = attention_dict
        self.name = name

        # Define the linear layers
        self.linear_x = nn.Linear(in_features=embedding_dim, out_features=num_hidden)
        self.linear_q = nn.Linear(in_features=query_dim, out_features=num_hidden)
        self.out_linear = nn.Linear(in_features=num_hidden, out_features=output_features)
        self.output_softmax_scores = nn.Softmax()

    def forward(self, input_sequence, query_vector):
        # Returns the alignment scores
        if self.additive:
            s = additive_fn(x=input_sequence, q=query_vector, linear_x=self.linear_x,
                        linear_q=self.linear_q, num_hidden=self.num_hidden, activation=self.activation,
                            out_linear=self.out_linear)
        else:
            s = multiplicative_fn(x=input_sequence, q=query_vector, linear_x=self.linear_x,
                              linear_q=self.linear_q)

        scores = self.output_softmax_scores(s)
        # A large score here for a particular x (embedding) means that it contributes
        # important information to the given query vector
        # Dimension of scores -> b x n
        expectations_of_sampling = torch.sum(torch.mul(scores, input_sequence))
        if self.save_attention:
            if self.name is not None and self.attention_dict is not None:
                self.attention_dict[self.name] = scores
        return scores, expectations_of_sampling

    # Weights Initialization
    def init_weights(self, init_range=0.1):
        self.linear_x.weight.data.uniform_(-init_range, init_range)
        self.linear_q.weight.data.uniform_(-init_range, init_range)
        self.out_linear.weight.data.uniform(-init_range, init_range)


class MultiAttention(nn.Module):
    """
    The Mutli Attention module

    Alignment Score computed for each feature

    Score of a token pair is vector rather than a scalar

    Has embedding_dim indicators for embedding_dim features
    Each indicator has a probability
    distribution that is generated by applying softmax to
    the n alignment scores of the corresponding feature.


    """
    def __init__(self, input_dim, embedding_dim, query_dim, num_hidden,
                 activation, output_features,
                 use_additive=True, save_attention=False,
                 attention_dict=None, name=None):
        """
        :param input_dim: The input dimension of the sequence or length of the sequence
        :param query_dim: The dimension of the query embedding - Current State || Desired Goal Vector
        :param num_hidden: The number of hidden layers in the network
        :param embedding_dim: The dimension of each embedding - State || Achieved Goal Vector
        :param activation: The non-linear activation function to use
        :param use_additive_fn: Boolean to determine whether to use additive or multiplicative attention
        """
        super(MultiAttention, self).__init__()
        self.input_dim = input_dim
        self.query_dim = query_dim
        self.embedding_dim = embedding_dim
        self.num_hidden = num_hidden
        self.activation = activation
        self.additive = use_additive
        self.out_features = output_features
        self.save_attention = save_attention
        self.attention_dict = attention_dict
        self.name = name

        # Define the linear layers
        self.linear_x = nn.Linear(in_features=embedding_dim, out_features=num_hidden)
        self.linear_q = nn.Linear(in_features=query_dim, out_features=num_hidden)
        self.out_linear = nn.Linear(in_features=num_hidden, out_features=embedding_dim)
        self.output_softmax_scores = nn.Softmax()

    def forward(self, input_sequence, query_vector):
        if self.additive:
            s = additive_fn(x=input_sequence, q=query_vector, linear_x=self.linear_x,
                            linear_q=self.linear_q, num_hidden=self.num_hidden, activation=self.activation,
                            out_linear=self.out_linear)
        else:
            s = multiplicative_fn(x=input_sequence, q=query_vector, linear_x=self.linear_x,
                                  linear_q=self.linear_q)

        prob_distribution = self.output_softmax_scores(s)
        score_vectors = torch.sum(torch.bmm(prob_distribution, input_sequence))
        if self.save_attention:
            if self.attention_dict is not None and self.name is not None:
                self.attention_dict[self.name] = prob_distribution
        return score_vectors, prob_distribution

    # Weights Initialization
    def init_weights(self, init_range=0.1):
        self.linear_x.weight.data.uniform_(-init_range, init_range)
        self.linear_q.weight.data.uniform_(-init_range, init_range)
        self.out_linear.weight.data.uniform(-init_range, init_range)


class SelfAttention(nn.Module):
    """
    The Self Attention module produces context-aware representations by
    exploring the dependency between two tokens xi and xj from the same sequence x.

    """
    def __init__(self, input_dim, embedding_dim, query_dim, num_hidden,
                 output_features, activation, use_additive=True,
                 token2token=True, seq2token=False, attention_dict=None,
                 save_attention=False, name=None):
        """

        :param input_dim:
        :param embedding_dim:
        :param query_dim:
        :param num_hidden:
        :param activation:
        :param use_additive:
        :param token2token:
        :param seq2token:
        """
        super(SelfAttention, self).__init__()
        self.token2token = token2token
        self.seq2token = seq2token
        self.input_dim = input_dim
        self.embedding_dim = embedding_dim
        self.query_dim = query_dim
        self.num_hidden = num_hidden
        self.out_features = output_features
        self.activation = activation
        self.additive = use_additive
        self.attention_dict = attention_dict
        self.save_attention = save_attention
        self.name = name

        # Define the linear layers
        self.linear_x = nn.Linear(in_features=self.embedding_dim, out_features=self.num_hidden)
        self.linear_q = nn.Linear(in_features=self.query_dim, out_features=self.num_hidden)
        self.out_linear = nn.Linear(in_features=self.num_hidden, out_features=self.out_features)
        self.softmax_probs = nn.Softmax()

    def forward(self, input_sequence):
        # In this case the query vector itself is the input sequence

        # Dimensions of the vectors
        # Input Dimension -> B x Seq Length x embedding_dim
        # Query Vector Dimension -> B x Seq Length x embedding_dim

        x = input_sequence
        query_vector = input_sequence

        # Dimension of x and query after linear_x and linear_q respectively
        # B x Seq Length x hidden_dim
        x_ = self.linear_x(x)
        q_ = self.linear_q(query_vector)

        t = nn.ReLU(x_ + q_)

        # Dimension of t after out_linear
        # B x Seq Length x embedding_dim
        o = self.out_linear(t)

        # Softmax Scores
        scores = self.softmax_probs(o)
        # Dimension of expectation of the sampling
        # B x seq length (Contrary to the paper, we are summing along the embedding dimension)
        expectation_of_sampling = torch.sum(torch.mul(scores), dim=-1)

        if self.save_attention:
            if self.attention_dict is not None and self.name is not None:
                self.attention_dict[self.name] = scores

        return scores, expectation_of_sampling

    # Weights Initialization
    def init_weights(self, init_range=0.1):
        self.linear_x.weight.data.uniform_(-init_range, init_range)
        self.linear_q.weight.data.uniform_(-init_range, init_range)
        self.out_linear.weight.data.uniform(-init_range, init_range)


class GoalNetwork(nn.Module):
    """
    This network uses the self, multi and vanilla attention modules and returns the
    top n vectors according to the softmax probabilities.
    """

    def __init__(self, input_dim, query_dim, embedding_dim, num_hidden,
                 output_features, activation, use_additive, use_token2token,
                 use_self_attn, save_attention=False, attention_dict=None,
                 use_multi_attn=False):
        """

        :param input_dim:
        :param query_dim:
        :param embedding_dim:
        :param num_hidden:
        :param output_features:
        :param activation:
        :param use_additive:
        :param use_token2token:
        :param use_self_attn:
        :param save_attention:
        :param attention_dict:
        """
        super(GoalNetwork, self).__init__()
        self.input_dim = input_dim
        self.query_dim = query_dim
        self.embedding_dim = embedding_dim
        self.num_hidden = num_hidden
        self.output_features = output_features
        self.activation = activation
        self.additive = use_additive
        self.use_self_attn = use_self_attn
        self.token2token = use_token2token
        self.save_attn = save_attention
        self.attention_dict = attention_dict
        self.multi_attn = use_multi_attn
        self.self_attn = None
        self.self_linear = None

        if self.use_self_attn:
            self.self_attn = SelfAttention(input_dim=self.input_dim, embedding_dim=self.embedding_dim,
                                           query_dim=self.query_dim, num_hidden=self.num_hidden,
                                           output_features=self.embedding_dim, activation=self.activation,
                                           use_additive=self.additive, token2token=self.token2token,
                                           save_attention=self.save_attn, attention_dict=self.attention_dict,
                                           name='SelfAttention')
            self.self_linear = nn.Linear(in_features=1, out_features=self.num_hidden)

        if self.multi_attn:
            # Use multi dimensional attention
            self.attn = MultiAttention(input_dim=self.input_dim, embedding_dim=self.embedding_dim,
                                           query_dim=self.query_dim, num_hidden=self.num_hidden,
                                           output_features=self.embedding_dim, activation=self.activation,
                                           use_additive=self.additive,
                                           save_attention=self.save_attn, attention_dict=self.attention_dict,
                                           name='MultiAttention')

        else:
            # Use the traditional vanilla attention
            self.attn = VanillaAttention(input_dim=self.input_dim, embedding_dim=self.embedding_dim,
                                           query_dim=self.query_dim, num_hidden=self.num_hidden,
                                           output_features=1, activation=self.activation,
                                           use_additive=self.additive,
                                           save_attention=self.save_attn, attention_dict=self.attention_dict,
                                           name='VanillaAttention')

        self.linear_attn = nn.Linear(in_features=1, out_features=self.num_hidden)
        self.output_linear = nn.Linear(in_features=num_hidden, out_features=1)

    def forward(self, input_sequence, current_embedding):
        scores_self, expectations_self = None, None
        if self.self_attn is not None:
            scores_self, expectations_self = self.self_attn(input_sequence)
        scores, expectation = self.attn(input_sequence, current_embedding)
        if expectations_self is not None and self.self_linear is not None:
            expectations_self = self.self_linear(expectations_self)
        expectation = self.linear_attn(expectation)
        if expectations_self is not None:
            e = expectation + expectations_self
        else:
            e = expectation
        e = self.activation(e)
        o = self.output_linear(e)
        return o





