import torch


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

    Generate a 3x3 offset mask:

    off_0 = tl | off_1 | off_2 = tr
    -----------|-------|-----------
    off_3      | off_4 | off_5
    -----------|-------|-----------
    off_6 = bl | off_7 | off_8 = br

    tl = top-left, tr = top-right, bl = bottom-left, br = bottom-right.

    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, 4, H, W]
    Returns:
        torch.Tensor: Offset tensor of shape [N, 18, H, W].
    """
    N, C, H, W = hbox.shape

    xmin, ymin, xmax, ymax = hbox.unbind(1)

    # Compute corners from hbox xmin,ymin,xmax,ymax
    tl = torch.stack((ymin, xmin), dim=1)
    bl = torch.stack((ymax, xmin), dim=1)
    br = torch.stack((ymax, xmax), dim=1)
    tr = torch.stack((ymin, xmax), dim=1)

    # Corners
    off_0 = tl
    off_2 = tr
    off_8 = br
    off_6 = bl

    # Center
    off_4 = (tl + tr + br + bl) / 4

    # Middle-points on edges
    off_1 = (tl + tr) / 2
    off_5 = (tr + br) / 2
    off_7 = (bl + br) / 2
    off_3 = (tl + bl) / 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