from typing import Optional
from typing import Tuple
import torch


def decode_boxes(anchors, delta, clip_delta=5.0, image_shape: Optional[Tuple[int, int]] = None):
    """Decode boxes according to z_where.

    Notes:
        Make use of broadcast to deal with batch dimension.

    References: https://github.com/pytorch/vision/blob/master/torchvision/models/detection/_utils.py

    Args:
        anchors (torch.Tensor): (..., 4). [x_min, y_min, x_max, y_max]
        delta (torch.Tensor): (..., 4)
        clip_delta (float): clip dw and dh
        image_shape (optional, tuple): [H, W]

    Returns:
        boxes (torch.Tensor): (..., 4). [x_min, y_min, x_max, y_max]
    """
    x_min, y_min, x_max, y_max = torch.split(anchors, 1, dim=-1)
    ctr_x = 0.5 * (x_min + x_max)
    ctr_y = 0.5 * (y_min + y_max)
    w = x_max - x_min
    h = y_max - y_min

    # decode z_where
    dx, dy, dw, dh = torch.split(delta, 1, dim=-1)
    # Prevent sending too large values into torch.exp()
    dw = torch.clamp(dw, min=-clip_delta, max=clip_delta)
    dh = torch.clamp(dh, min=-clip_delta, max=clip_delta)

    pred_x_ctr = ctr_x + dx * w
    pred_y_ctr = ctr_y + dy * h
    pred_w = torch.exp(dw) * w
    pred_h = torch.exp(dh) * h

    # decode new boxes
    x_min = (pred_x_ctr - 0.5 * pred_w)
    x_max = (pred_x_ctr + 0.5 * pred_w)
    y_min = (pred_y_ctr - 0.5 * pred_h)
    y_max = (pred_y_ctr + 0.5 * pred_h)

    # clip to image
    if image_shape is not None:
        x_min = x_min.clamp(min=0.0)
        x_max = x_max.clamp(max=image_shape[1])
        y_min = y_min.clamp(min=0.0)
        y_max = y_max.clamp(max=image_shape[0])

    return torch.cat([x_min, y_min, x_max, y_max], dim=-1)