import math
import torch


def rand_bbox(img_shape, lam, margin=0.0, count=1):
    """Standard CutMix bounding-box
    Generates a random square bbox based on lambda value. This impl includes
    support for enforcing a border margin as percent of bbox dimensions.

    Args:
        img_shape (tuple): Image shape as tuple
        lam (float): Cutmix lambda value
        margin (float): Percentage of bbox dimension to enforce as margin (reduce amount of box outside image)
        count (int): Number of bbox to generate
    """
    ratio = math.sqrt(1 - lam)
    img_h, img_w = img_shape[-2:]
    cut_h, cut_w = int(img_h * ratio), int(img_w * ratio)
    margin_y, margin_x = int(margin * cut_h), int(margin * cut_w)
    cy = torch.randint(0 + margin_y, img_h - margin_y, (count,))
    cx = torch.randint(0 + margin_x, img_w - margin_x, (count,))
    yl = torch.clamp(cy - cut_h // 2, 0, img_h)
    yh = torch.clamp(cy + cut_h // 2, 0, img_h)
    xl = torch.clamp(cx - cut_w // 2, 0, img_w)
    xh = torch.clamp(cx + cut_w // 2, 0, img_w)
    return yl, yh, xl, xh