import torch


def generate_sorted_element_mask(previous_actions, input_length: int) -> torch.Tensor:
    """
    Generates a mask that prevents actions from attending to elements of the unordered set that have already been
    placed into the ordered sequence.
    Args:
        previous_actions: List of previous actions (in order) that we need to mask
        input_length: Number of elements in the unordered sequence

    Returns:
        Memory mask of shape (nb of previous actions + 1, input sequence length) suitable for use in transformer

    """
    # Generate lower triangular matrix (creates columns for masked input elements)
    # i_th column of masked_cols is equal to the {a_i}'th column of the mask
    masked_cols = torch.tril(torch.ones(len(previous_actions) + 1, len(previous_actions)) * float('-inf'), diagonal=-1)
    # Create empty mask
    mask = torch.zeros(len(previous_actions) + 1, input_length)
    # For each previous action, prevent further actions from attending to its corresponding input element
    mask[:, previous_actions] = masked_cols
    return mask