import torch


def corners_to_offset_mask(corners: torch.Tensor) -> torch.Tensor:
    """
    Convert the corner regression to an offset mask.

    Generate a 3x3 offset mask:

    off_0 = c0 | off_1 | off_2 = c3
    -----------|-------|-----------
    off_3      | off_4 | off_5
    -----------|-------|-----------
    off_6 = c1 | off_7 | off_8 = c2

    Where the corners off_0, off_2, off_6 and off_8 are equal to the original corners,
    off_4 is the mean of all corners (center), and off_1, off_3, off_5, off_7 are
    middle-points on the edge between the corners accordingly.

    Args:
        corners: Tensor of shape [N, 8, H, W]
    Returns:
        torch.Tensor: Offset tensor of shape [N, 18, H, W].
    """
    N, C, H, W = corners.shape
    # Swap (x, y) coordinates since the offset arrays expects (y, x) format
    corners = corners[:, [1, 0, 3, 2, 5, 4, 7, 6], :, :]
    c0, c1, c2, c3 = corners.view(N, 4, 2, H, W).unbind(1)

    # Corners
    off_0 = c0
    off_2 = c3
    off_8 = c2
    off_6 = c1

    # Center
    off_4 = (c0 + c1 + c2 + c3) / 4

    # Middle-points on edges
    off_1 = (off_0 + off_2) / 2
    off_5 = (off_2 + off_8) / 2
    off_7 = (off_6 + off_8) / 2
    off_3 = (off_0 + off_6) / 2

    # Cat in column-major order
    offset = torch.cat(
        (off_0, off_1, off_2, off_3, off_4, off_5, off_6, off_7, off_8), dim=1
    )
    return offset