import torch
import torch.nn.functional as F
from torchvision.transforms import v2


class ResizeAndPadBatch:
    """
    Vectorized resize + zero pad for batched tensors.

    Args:
        imgs (Tensor): shape [B, C, H, W], pixel values in [0,1] or [0,255].
        target_size (tuple): (target_h, target_w) final image shape.

    Returns:
        Tensor: shape [B, C, target_h, target_w]
    """

    def __init__(self, target_size=(640, 640), fill_value=0):
        self.target_size = target_size
        self.fill_value = fill_value
        self.transform_stats = None

    def __call__(self, imgs: torch.Tensor):

        b, c, h, w = imgs.shape
        target_h, target_w = self.target_size

        # Compute scaling factors: preserve aspect ratio by same scale factor per image
        scale_factors = float(
            torch.minimum(torch.tensor(target_h / h), torch.tensor(target_w / w))
        )

        # New intermediate size
        new_h = int(h * scale_factors)
        new_w = int(w * scale_factors)

        # Resize with bilinear interpolation
        resized = F.interpolate(
            imgs, size=(new_h, new_w), mode="bilinear", align_corners=False
        )

        # Calculate padding amounts (left, right, top, bottom)
        pad_h = target_h - new_h
        pad_w = target_w - new_w
        pad_top = pad_h // 2
        pad_bottom = pad_h - pad_top
        pad_left = pad_w // 2
        pad_right = pad_w - pad_left

        # Apply padding (pad order in F.pad is [left, right, top, bottom])
        padded = F.pad(
            resized,
            (pad_left, pad_right, pad_top, pad_bottom),
            mode="constant",
            value=self.fill_value,
        )

        self.transform_stats = {
            "scale_factor": scale_factors,
            'hw': [h,w],
            "new_hw": [new_h, new_w],
            "pad_TBLR": [pad_top,pad_bottom, pad_left, pad_right]
        }

        return padded


clip_transforms = v2.Compose(
    [
        v2.Resize(size=(512, 512)),
        v2.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    ]
)

det_transforms = v2.Compose([ResizeAndPadBatch(target_size=(640, 640))])