import torch


def get_relative_position_matrix(length, max_relative_position, direction, offset=True):
    """ Generate matrix of relative positions between inputs ([..., length])."""
    range_vec = torch.arange(length).long()
    if torch.cuda.is_available():
        range_vec = range_vec.cuda()
    range_mat = range_vec[:, None].expand(length, length)
    distance_mat = range_mat - range_mat.transpose(0, 1)
    if max_relative_position is None:
        distance_mat_clipped = distance_mat
    else:
        distance_mat_clipped = torch.clamp(distance_mat,
                                           -max_relative_position, max_relative_position)
    if direction:
        # Shift values to be >= 0. Each integer still uniquely identifies a relative
        # position difference.
        if offset and max_relative_position is not None:
            final_mat = distance_mat_clipped + max_relative_position
        else:
            final_mat = distance_mat_clipped
    else:
        # Do not distinguish the forward and backward positions.
        # Just leave the absolute relative position representation.
        final_mat = distance_mat_clipped.abs()
    return final_mat