"""Expectation-over-transformation ops applied inside the PGD loop so the
perturbation survives what a scraper/trainer does: JPEG, resize, crop, blur.

JPEG is non-differentiable, so we use BPDA / straight-through: forward pass uses
real PIL JPEG, backward pass is identity. That makes the optimizer account for
recompression in expectation instead of finding a fragile solution.
"""

from __future__ import annotations

import io
import math
import random

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image


def _to_pil_batch(x: torch.Tensor) -> list[Image.Image]:
    a = (x.clamp(0, 1).permute(0, 2, 3, 1).detach().cpu().float().numpy() * 255)
    a = a.round().astype(np.uint8)
    return [Image.fromarray(im) for im in a]


def jpeg_ste(x: torch.Tensor, quality: int) -> torch.Tensor:
    imgs = _to_pil_batch(x)
    out = []
    for im in imgs:
        buf = io.BytesIO()
        im.save(buf, format="JPEG", quality=int(quality))
        buf.seek(0)
        out.append(np.asarray(Image.open(buf).convert("RGB"), dtype=np.float32) / 255.0)
    xj = torch.from_numpy(np.stack(out)).permute(0, 3, 1, 2).to(x.device, x.dtype)
    return x + (xj - x).detach()  # straight-through


def _gauss_kernel(sigma: float, device, dtype) -> torch.Tensor:
    r = max(1, int(math.ceil(3 * sigma)))
    xs = torch.arange(-r, r + 1, device=device, dtype=torch.float32)
    k = torch.exp(-(xs ** 2) / (2 * sigma * sigma))
    k = (k / k.sum()).to(dtype)
    return k


def gaussian_blur(x: torch.Tensor, sigma: float) -> torch.Tensor:
    k = _gauss_kernel(sigma, x.device, x.dtype)
    ck = k.view(1, 1, 1, -1).repeat(x.shape[1], 1, 1, 1)
    pad = (k.numel() // 2)
    x = F.conv2d(F.pad(x, (pad, pad, 0, 0), mode="reflect"), ck, groups=x.shape[1])
    ck = k.view(1, 1, -1, 1).repeat(x.shape[1], 1, 1, 1)
    x = F.conv2d(F.pad(x, (0, 0, pad, pad), mode="reflect"), ck, groups=x.shape[1])
    return x


def resize_roundtrip(x: torch.Tensor, scale: float) -> torch.Tensor:
    _, _, h, w = x.shape
    nh, nw = max(8, int(h * scale)), max(8, int(w * scale))
    down = F.interpolate(x, size=(nh, nw), mode="bilinear", align_corners=False, antialias=True)
    return F.interpolate(down, size=(h, w), mode="bilinear", align_corners=False)


def random_crop_pad(x: torch.Tensor, keep: float) -> torch.Tensor:
    _, _, h, w = x.shape
    ch, cw = int(h * keep), int(w * keep)
    top = random.randint(0, h - ch)
    left = random.randint(0, w - cw)
    return x[:, :, top:top + ch, left:left + cw]


def apply_eot(x: torch.Tensor, rng: random.Random, strong: bool = True) -> torch.Tensor:
    """Sample a random subset of transforms. Always returns a valid [0,1] image."""
    if rng.random() < 0.85:
        x = jpeg_ste(x, rng.randint(55, 92))
    if rng.random() < 0.6:
        x = resize_roundtrip(x, rng.uniform(0.55, 1.0))
    if strong and rng.random() < 0.4:
        x = gaussian_blur(x, rng.uniform(0.4, 1.4))
    if strong and rng.random() < 0.35:
        x = random_crop_pad(x, rng.uniform(0.85, 0.98))
    return x.clamp(0, 1)