import torch


def calc_additive_attention(W1, W2, v, x):
    """
    Calculates additive attention.

    Args:
        W1: weight matrix for projecting x
        W2: weight matrix for projecting x
        v: vector used in attention mechanism
        x: input vectors

    Returns:
        attention output
    """
    batch_size, seq_len, input_size = x.size()
    W1x = W1(x)  # (batch_size, seq_len, att_dim)
    W2x = W2(x)  # (batch_size, seq_len, att_dim)

    s = torch.zeros((batch_size, seq_len, seq_len)).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))

    # process one sequence at a time to save memory
    for idx in range(seq_len):
        s[:, idx, :] = v(torch.tanh(W1x + W2x[:, idx, :].unsqueeze(1))).transpose(1, 2).squeeze()  # (batch_size, 1, seq_len)
    mask = (torch.eye(seq_len) == 1).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
    s.masked_fill_(mask, float('-inf'))  # mask out similarity between the same tokens

    a = torch.softmax(s, dim=2)  # (batch_size, seq_len, seq_len)
    c = torch.bmm(a, x)  # (batch_size, seq_len, input_size)

    return c