import torch


def coordinate_embeddings(boxes, dim):
    """
    Coordinate embeddings of bounding boxes
    :param boxes: [BS, K, 4] ([x1, y1, x2, y2])
    :param dim: sin/cos embedding dimension
    :return: [BS, K, 4, 2 * dim]
    """

    batch_size, num_boxes, num_loc = boxes.shape

    # transform to (x_c, y_c, w, h) format
    pos = boxes.new_zeros((batch_size, num_boxes, 4))
    pos[:, :, 0] = (boxes[:, :, 0] + boxes[:, :, 2]) / 2 * 100
    pos[:, :, 1] = (boxes[:, :, 1] + boxes[:, :, 3]) / 2 * 100
    pos[:, :, 2] = (boxes[:, :, 2] - boxes[:, :, 0]) * 100
    pos[:, :, 3] = (boxes[:, :, 3] - boxes[:, :, 1]) * 100

    # sin/cos embedding
    dim_mat = 1000 ** (torch.arange(dim, dtype=boxes.dtype, device=boxes.device) / float(dim))
    sin_embedding = (pos.view((batch_size, num_boxes, 4, 1)) / dim_mat.view((1, 1, 1, -1))).sin()
    cos_embedding = (pos.view((batch_size, num_boxes, 4, 1)) / dim_mat.view((1, 1, 1, -1))).cos()

    return torch.cat((sin_embedding, cos_embedding), dim=-1)