import torch as pt
import torch.nn.functional as F
import numbers

def resize_down_up(x, scale=0.5, mode="bilinear"):
    # shrinks image then blows it back up to see what info was lost
    if isinstance(scale, pt.Tensor):
        scale = scale.item()
    _val_invars(x, scale=scale, mode=mode, type="resize")
    B, C, H, W = x.shape
    h_new, w_new = max(1, int(H * scale)), max(1, int(W * scale))

    kwargs = {}
    if mode in {"bilinear", "bicubic"}:
        kwargs["align_corners"] = False

    x_down = F.interpolate(x, size=(h_new, w_new), mode=mode, **kwargs)
    return F.interpolate(x_down, size=(H,W), mode=mode, **kwargs)

def decimate(x, factor=2):
    # decimation: keep every factor-th pixel, then upsample back up.
    _val_invars(x, factor=factor, type="decimate")
    B, C, H, W = x.shape
    x_dec = x[:, :, ::factor, ::factor]
    return F.interpolate(x_dec, size=(H,W), mode="nearest")

def blur_decimate(x, blur):
    B, C, H, W = x.shape
    x_down = blur(x)
    return F.interpolate(x_down, size=(H,W), mode="bilinear", align_corners=False)

def checkerboard_alias_attack(x, epsilon=0.5):
    # injects high-frequency noise that mimics sampling artifacts
    _val_invars(x, epsilon=epsilon, type="attack")
    B, C, H, W = x.shape
    device = x.device
    dtype = x.dtype

    rows = pt.arange(H, device=device).view(H,1)
    cols = pt.arange(W, device=device).view(1,W)
    checker = ((rows + cols) % 2).float() * 2.0 - 1.0 # values in {-1, +1}
    checker = checker.view( 1, 1, H, W).expand(B, C, H, W).to(dtype)

    x_adv = (x + epsilon * checker).clamp(-1.0, 1.0)
    return x_adv

def _val_invars(x, **kwargs):
    # validate invariants
    if x.ndim != 4:
        raise ValueError(f"x needs shape (B,C,H,W), got {x.shape}")
    op_type = kwargs.get("type")

    if op_type == "resize":
        scale = kwargs.get("scale")
        if not isinstance(scale, numbers.Real):
            raise ValueError(f"scale must be a real number, got {scale}")
        if scale <= 0.0 or scale >= 1.0:
            raise ValueError(f"scale must be in (0,1), got {scale}")
        mode = kwargs.get("mode")
        valid_modes = {"bilinear", "bicubic", "nearest", "area"}
        if mode not in valid_modes:
            raise ValueError(f"mode must be one of {valid_modes}, got {mode}")
    elif op_type == "decimate":
        factor = kwargs.get("factor")
        if factor < 2:
            raise ValueError(f"decimation factor must be >=2, got {factor}")
    elif op_type == "attack":
        epsilon = kwargs.get("epsilon")
        if not (0.0 < epsilon < 1.0):
            raise ValueError(f"epsilon must be (0,1), got {epsilon}")