import torch


def laplace_attention(queries, keys, values, scale, normalise=True):
    """
    Here we compute the Laplace exponential attention. Each value is weighted by an amount that depends
    on the distance of the query from the corresponding key (specifically, w_i ~ exp(-||q-k_i||/scale))
    :param queries: e.g. query corresponding to x_target: [batch_size, N_target, key_size]
    :param keys: e.g. x_context: [batch_size, N_context, key_size]
    :param values: e.g. values corresponding to the aggregated embeddings r_i [batch_size, N_context, value_size]
    :param scale: float value which scales the L1 distance.
    :param normalise: Boolean, determines whether we should normalise s.t. sum of weights = 1.
    :return: A tensor [batch_size, N_target, value_size].
    """

    keys = torch.unsqueeze(keys, dim=1)  # [batch_size, 1, N_context, key_size]
    queries = torch.unsqueeze(queries, dim=1)  # [batch_size, N_target, 1, key_size]

    unnorm_weights = -torch.abs((keys - queries)/scale)  # [batch_size, N_target, N_context, key_size]
    unnorm_weights = torch.sum(unnorm_weights, dim=-1, keepdim=False) # [batch_size, N_target, N_context]

    if normalise:
        attention = torch.softmax(unnorm_weights, dim=-1)  # [batch_size, N_target, N_context]
    else:
        attention = 1 + torch.tanh(unnorm_weights)  # [batch_size, N_target, N_context]

    # Einstein summation over weights and values
    output= torch.matmul(attention, values)  # [batch_size, N_target, value_size]

    return output