import torch
import torch.nn.functional as F
import math


def make_bool_mask(*dims: int) -> torch.BoolTensor:
    """Make a boolean mask for the given dimensions.

    Example:
        _make_bool_mask(2, -2, 2) == (True, True, False, False, True, True)
        _make_bool_mask(2, 0, 2) == (True, True, True, True)

    Args:
        dims: The dimensions to make the mask for.

    Returns:
        A tensor of booleans.
    """
    lens = [abs(d) for d in dims]
    flags = [d > 0 for d in dims]
    parts = [torch.full((l,), f, dtype=torch.bool) for l, f in zip(lens, flags)]
    return torch.cat(parts) if parts else torch.empty(0, dtype=torch.bool)



def resize_with_pad(
    image: torch.Tensor,
    target_h: int,
    target_w: int,
    mode: str = "bilinear",
) -> torch.Tensor:
    """
    Resize a tensor to target size while keeping aspect ratio.
    Supports both [3, H, W] and [B, 3, H, W] inputs.
    Vectorized over the batch dimension (no Python loop).

    Args:
        image: torch.Tensor, shape [3, H, W] or [B, 3, H, W], values in [0, 1].
        target_h: target height.
        target_w: target width.
        mode: interpolation mode ('bilinear', 'nearest', etc.).
    Returns:
        torch.Tensor of shape [3, target_h, target_w] or [B, 3, target_h, target_w].
    """
    if image.ndim == 3:
        # [3, H, W] -> [1, 3, H, W]
        image = image.unsqueeze(0)
        squeeze_back = True
    elif image.ndim == 4:
        squeeze_back = False
    else:
        raise ValueError(f"Expected tensor of shape [3,H,W] or [B,3,H,W], got {image.shape}")

    B, C, H, W = image.shape

    if H == target_h and W == target_w:
        return image.squeeze(0) if squeeze_back else image

    scale = min(target_h / H, target_w / W)
    new_h, new_w = int(round(H * scale)), int(round(W * scale))

    image_resized = F.interpolate(
        image,
        size=(new_h, new_w),
        mode=mode,
        align_corners=False if mode == "bilinear" else None,
    )  # [B, 3, new_h, new_w]

    pad_top = (target_h - new_h) // 2
    pad_bottom = target_h - new_h - pad_top
    pad_left = (target_w - new_w) // 2
    pad_right = target_w - new_w - pad_left

    image_padded = F.pad(image_resized, (pad_left, pad_right, pad_top, pad_bottom), value=0.0)
    image_padded = image_padded.clamp(0.0, 1.0)

    # sanity check
    assert image_padded.shape[-2:] == (target_h, target_w)

    return image_padded.squeeze(0) if squeeze_back else image_padded


def resize_center_crop(
    image: torch.Tensor,
    target_h: int,
    target_w: int,
    mode: str = "bilinear",
) -> torch.Tensor:
    """
    Resize an image or a batch of images so that the shortest side is scaled
    to at least the target size, then apply a center crop to exactly
    [C, target_h, target_w] or [B, C, target_h, target_w].

    Supports:
        - [C, H, W]
        - [B, C, H, W]
    and performs vectorized operations over the batch dimension.

    Args:
        image: Input tensor of shape [C,H,W] or [B,C,H,W].
        target_h: Target height after cropping.
        target_w: Target width after cropping.
        mode: Interpolation mode for resizing.

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

    # ------------------------------------------------------
    # Normalize input shape to batch mode
    # ------------------------------------------------------
    if image.ndim == 3:
        # Convert [C, H, W] → [1, C, H, W]
        assert image.shape[0] in (1, 3), f"Expected [C,H,W], got {image.shape}"
        image = image.unsqueeze(0)
        squeeze_back = True
    elif image.ndim == 4:
        squeeze_back = False
    else:
        raise ValueError(f"Expected [C,H,W] or [B,C,H,W], got {image.shape}")

    B, C, H, W = image.shape

    # If the image is already the correct size, return early
    if (H, W) == (target_h, target_w):
        return image.squeeze(0) if squeeze_back else image

    # ------------------------------------------------------
    # Compute scale factor so the shortest side >= target size
    # ------------------------------------------------------
    # Use max() to ensure both resized dimensions are >= target dimensions.
    scale = max(target_h / H, target_w / W)

    # Ceil prevents rounding down to a size smaller than required
    new_h = int(math.ceil(H * scale))
    new_w = int(math.ceil(W * scale))

    # ------------------------------------------------------
    # Resize the batch in a single vectorized operation
    # ------------------------------------------------------
    align_corners = False if mode in ("bilinear", "bicubic") else None
    x = F.interpolate(
        image,
        size=(new_h, new_w),
        mode=mode,
        align_corners=align_corners,
    )  # [B, C, new_h, new_w]

    # ------------------------------------------------------
    # Center crop to the target size
    # ------------------------------------------------------
    top = max((new_h - target_h) // 2, 0)
    left = max((new_w - target_w) // 2, 0)

    x = x[:, :, top : top + target_h, left : left + target_w]

    # ------------------------------------------------------
    # (Rare) If rounding left us with slightly smaller output,
    # add minimal padding to reach the exact target size
    # ------------------------------------------------------
    cur_h, cur_w = x.shape[-2:]
    pad_h = target_h - cur_h
    pad_w = target_w - cur_w

    if pad_h > 0 or pad_w > 0:
        # F.pad padding order: (left, right, top, bottom)
        x = F.pad(
            x,
            (0, max(pad_w, 0), 0, max(pad_h, 0)),  # Only pad right/bottom
        )

    # Clamp values to valid range
    x = x.clamp(0.0, 1.0)

    # Sanity check
    assert x.shape[-2:] == (target_h, target_w)

    # ------------------------------------------------------
    # Return a squeezed tensor if input was a single image
    # ------------------------------------------------------
    return x.squeeze(0) if squeeze_back else x