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

def gaussian_kernel1d(size: int, sigma: float, device, dtype):
    """
    Generate a 1D Gaussian kernel.
    
    Args:
        size: Size of the kernel.
        sigma: Standard deviation for the kernel.
        device: Device for the kernel.
        dtype: Data type for the kernel.
    """
    x = torch.arange(size, device=device, dtype=dtype) - (size - 1) / 2
    k = torch.exp(-0.5 * (x / sigma) ** 2)
    k = k / k.sum()
    return k


def gaussian_blur(img: torch.Tensor, ksize: int, sigma: float):
    """
    Apply Gaussian blur to an image tensor.
    
    Args:
        img: Input image tensor of shape (N, C, H, W).
        ksize: Kernel size for blurring.
        sigma: Standard deviation for blurring.
    """
    if ksize <= 1 or sigma <= 0:
        return img
    dtype = img.dtype
    device = img.device
    k1d = gaussian_kernel1d(ksize | 1, sigma, device, dtype)
    kx = k1d.view(1, 1, 1, -1)
    ky = k1d.view(1, 1, -1, 1)
    c = img.shape[1]
    img = F.conv2d(img, kx.expand(c, 1, 1, -1), padding=(0, (ksize | 1) // 2), groups=c)
    img = F.conv2d(img, ky.expand(c, 1, -1, 1), padding=(((ksize | 1) // 2), 0), groups=c)
    return img


def unsharp_mask(img: torch.Tensor, sigma: float, amount: float, threshold: float = 0.0):
    """
    Unsharp mask for tensors shaped (N,C,H,W) or (1,1,H,W), using Gaussian blur.

    Args:
        img: Input tensor in [0,1] (expected), shape (N,C,H,W).
        sigma: Gaussian sigma (radius proxy). If <= 0, no sharpening is applied.
        amount: Sharpening amount (gain) typically in [0..2].
        threshold: Only enhance details where |img - blur| > threshold. Range [0..1].

    Returns:
        Sharpened tensor, clamped to [0,1].
    """
    if sigma <= 0.0 or amount == 0.0:
        return img
    # Compute a reasonable odd kernel size from sigma
    # k ~ 2*ceil(3*sigma)+1 ensures sufficient support
    if sigma < 0.2:
        ksize = 1
    else:
        k_est = int(2 * math.ceil(3.0 * float(sigma)) + 1)
        ksize = max(3, k_est | 1)
    base = img
    # Ensure shape is NCHW
    need_unsqueeze = False
    if base.dim() == 3:
        base = base.unsqueeze(0)
        need_unsqueeze = True
    if base.dim() != 4:
        # Fallback: do nothing for unsupported shapes
        return img
    blurred = gaussian_blur(base, ksize=ksize, sigma=float(sigma))
    high = base - blurred
    if threshold > 0.0:
        mask = (high.abs() > float(threshold)).to(base.dtype)
        high = high * mask
    out = base + float(amount) * high
    if need_unsqueeze:
        out = out.squeeze(0)
    return out


def _depthwise_conv(img: torch.Tensor, kernel: torch.Tensor):
    """Depthwise conv2d for NCHW with per-channel same kernel."""
    c = img.shape[1]
    k = kernel.to(dtype=img.dtype, device=img.device)
    k = k.view(1, 1, *k.shape).expand(c, 1, k.shape[-2], k.shape[-1])
    pad_h = k.shape[-2] // 2
    pad_w = k.shape[-1] // 2
    return F.conv2d(img, k, padding=(pad_h, pad_w), groups=c)


def sobel_magnitude_nchw(img: torch.Tensor):
    """
    Sobel gradient magnitude for NCHW tensor. Returns (N,1,H,W) normalized to [0,1].
    """
    if img.dim() != 4:
        raise ValueError("sobel_magnitude_nchw expects NCHW tensor")
    # Kernels
    kx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
    ky = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
    gx = _depthwise_conv(img, kx)
    gy = _depthwise_conv(img, ky)
    mag = torch.sqrt(gx * gx + gy * gy + 1e-12)
    mag = mag.mean(dim=1, keepdim=True)
    mag = (mag - mag.amin(dim=(-2, -1), keepdim=True)) / (mag.amax(dim=(-2, -1), keepdim=True) - mag.amin(dim=(-2, -1), keepdim=True) + 1e-12)
    return mag


def laplacian_magnitude_nchw(img: torch.Tensor):
    """Laplacian-of-Gaussian-lite: plain Laplacian magnitude per-pixel, normalized to [0,1], shape (N,1,H,W)."""
    if img.dim() != 4:
        raise ValueError("laplacian_magnitude_nchw expects NCHW tensor")
    k = torch.tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
    lo = _depthwise_conv(img, k)
    mag = lo.abs().mean(dim=1, keepdim=True)
    mag = (mag - mag.amin(dim=(-2, -1), keepdim=True)) / (mag.amax(dim=(-2, -1), keepdim=True) - mag.amin(dim=(-2, -1), keepdim=True) + 1e-12)
    return mag


def local_variance_nchw(img: torch.Tensor, ksize: int = 7):
    """
    Local variance over a sliding window per channel, averaged across channels, normalized to [0,1].
    """
    if img.dim() != 4:
        raise ValueError("local_variance_nchw expects NCHW tensor")
    k = max(3, int(ksize) | 1)
    c = img.shape[1]
    device, dtype = img.device, img.dtype
    w = torch.ones((1, 1, k, k), device=device, dtype=dtype) / float(k * k)
    w = w.expand(c, 1, k, k)
    pad = (k // 2, k // 2)
    mean = F.conv2d(img, w, padding=pad, groups=c)
    mean2 = F.conv2d(img * img, w, padding=pad, groups=c)
    var = (mean2 - mean * mean).clamp_min(0.0)
    var = var.mean(dim=1, keepdim=True)
    var = (var - var.amin(dim=(-2, -1), keepdim=True)) / (var.amax(dim=(-2, -1), keepdim=True) - var.amin(dim=(-2, -1), keepdim=True) + 1e-12)
    return var

def perlin_noise(h, w, scale, octaves, persistence, lacunarity, seed, device, dtype):
    """
    Generate Perlin noise.
    
    Args:
        h: Height of the noise.
        w: Width of the noise.
        scale: Scale of the noise.
        octaves: Number of octaves for the noise.
        persistence: Persistence of the noise.
        lacunarity: Lacunarity of the noise.
        seed: Seed for the noise.
        device: Device for the noise.
        dtype: Data type for the noise.
    """
    if scale <= 0:
        scale = 1.0
    rng = torch.Generator(device=device)
    if seed is not None:
        rng.manual_seed(seed)

    def _gradients(gh, gw):
        theta = torch.rand((gh + 1, gw + 1), generator=rng, device=device, dtype=dtype) * 2 * math.pi
        return torch.stack((torch.cos(theta), torch.sin(theta)), dim=-1)

    def _perlin2d(size_y, size_x, freq_y, freq_x):
        gh = int(math.floor(freq_y)) + 1
        gw = int(math.floor(freq_x)) + 1
        g = _gradients(gh, gw)
        y = torch.linspace(0, freq_y, steps=size_y, device=device, dtype=dtype)
        x = torch.linspace(0, freq_x, steps=size_x, device=device, dtype=dtype)
        yi = y.floor().long()
        xi = x.floor().long()
        yf = (y - yi).unsqueeze(1).expand(size_y, size_x)
        xf = (x - xi).unsqueeze(0).expand(size_y, size_x)

        def _g(ix, iy):
            return g[iy.clamp_max(gh), ix.clamp_max(gw)]

        x0 = xi.unsqueeze(0).expand(size_y, size_x)
        y0 = yi.unsqueeze(1).expand(size_y, size_x)
        g00 = _g(x0, y0)
        g10 = _g(x0 + 1, y0)
        g01 = _g(x0, y0 + 1)
        g11 = _g(x0 + 1, y0 + 1)

        d00 = torch.stack((xf, yf), dim=-1)
        d10 = torch.stack((xf - 1.0, yf), dim=-1)
        d01 = torch.stack((xf, yf - 1.0), dim=-1)
        d11 = torch.stack((xf - 1.0, yf - 1.0), dim=-1)

        n00 = (g00 * d00).sum(dim=-1)
        n10 = (g10 * d10).sum(dim=-1)
        n01 = (g01 * d01).sum(dim=-1)
        n11 = (g11 * d11).sum(dim=-1)

        def _fade(t): return t * t * t * (t * (t * 6 - 15) + 10)
        u = _fade(xf)
        v = _fade(yf)
        nx0 = n00 * (1 - u) + n10 * u
        nx1 = n01 * (1 - u) + n11 * u
        nxy = nx0 * (1 - v) + nx1 * v
        return nxy

    yfreq = max(h / scale, 1.0)
    xfreq = max(w / scale, 1.0)
    amp = 1.0
    freq_mul = 1.0
    total = torch.zeros((h, w), device=device, dtype=dtype)
    maxamp = 0.0
    for _ in range(max(1, octaves)):
        total += amp * _perlin2d(h, w, yfreq * freq_mul, xfreq * freq_mul)
        maxamp += amp
        amp *= persistence
        freq_mul *= lacunarity
    total = total / maxamp if maxamp > 0 else total
    total = (total - total.min()) / (total.max() - total.min() + 1e-12)
    return total


def bayer_matrix(size: int, device, dtype):
    """
    Generate a simple Bayer matrix.
    
    Args:
        size: Size of the matrix.
        device: Device for the matrix.
        dtype: Data type for the matrix.
    """
    base2 = torch.tensor([[0, 2], [3, 1]], device=device, dtype=dtype)
    if size <= 2:
        m = base2
    elif size <= 4:
        m = torch.cat([torch.cat([4*base2+0, 4*base2+2], dim=1),
                       torch.cat([4*base2+3, 4*base2+1], dim=1)], dim=0)
    else:
        m2 = torch.cat([torch.cat([4*base2+0, 4*base2+2], dim=1),
                        torch.cat([4*base2+3, 4*base2+1], dim=1)], dim=0)
        rh = size // m2.shape[0] + (1 if size % m2.shape[0] else 0)
        rw = size // m2.shape[1] + (1 if size % m2.shape[1] else 0)
        m = m2.repeat(rh, rw)[:size, :size]
    m = (m - m.min()) / (m.max() - m.min() + 1e-12)
    return m


def _radial_frequency_grid(h: int, w: int, device, dtype):
    """
    Build a radial spatial-frequency grid in cycles/pixel for 2D FFT shaping.

    Returns tensor of shape (h, w) with radius r = sqrt(fx^2 + fy^2).
    """
    fy = torch.fft.fftfreq(h, d=1.0, device=device)
    fx = torch.fft.fftfreq(w, d=1.0, device=device)
    # Use float dtype for operations, then cast to requested dtype at the end
    Fy = fy.reshape(h, 1).to(torch.float64)
    Fx = fx.reshape(1, w).to(torch.float64)
    r = torch.sqrt(Fx * Fx + Fy * Fy)
    return r.to(dtype=dtype, device=device)


def spectral_noise(h: int, w: int, beta: float, seed: int, device, dtype):
    """
    Generate 2D spatial noise with power spectral density S(f) ∝ f^beta.

    beta = 0   -> white
    beta = -1  -> pink (1/f)
    beta = -2  -> brown/red (1/f^2)
    beta = +1  -> blue (f)
    beta = +2  -> violet/purple (f^2)
    """
    rng = torch.Generator(device=device)
    if seed is not None:
        rng.manual_seed(seed)
    base = torch.randn((h, w), generator=rng, device=device, dtype=dtype)
    F = torch.fft.fft2(base)
    r = _radial_frequency_grid(h, w, device, dtype)
    r = torch.clamp(r, min=1e-6)
    amp = torch.pow(r, beta * 0.5)
    F_shaped = F * amp
    x = torch.fft.ifft2(F_shaped).real
    x = (x - x.min()) / (x.max() - x.min() + 1e-12)
    return x


# -------------------------
# Affine schedule utilities
# -------------------------
def _ease_value(t, curve, back_k=1.70158, bezier=None):
    """Return eased value in [0,1] for t in [0,1] given curve name.
    Supported curves: linear, easeIn/Out/InOut Sine, Quad, Cubic, Quart, Quint, Expo, Circ, Back, cubicBezier.
    """
    t = max(0.0, min(1.0, float(t)))
    c = (curve or "easeInOutSine").lower()
    # Linear
    if c == "linear":
        return t
    # Sine
    if c == "easeinsine":
        return 1 - math.cos((t * math.pi) / 2)
    if c == "easeoutsine":
        return math.sin((t * math.pi) / 2)
    if c == "easeinoutsine":
        return -(math.cos(math.pi * t) - 1) / 2
    # Power families
    pow_map = {"easeinquad": 2, "easeincubic": 3, "easeinquart": 4, "easeinquint": 5}
    if c in pow_map:
        return t ** pow_map[c]
    if c in {"easeoutquad", "easeoutcubic", "easeoutquart", "easeoutquint"}:
        p = {"easeoutquad": 2, "easeoutcubic": 3, "easeoutquart": 4, "easeoutquint": 5}[c]
        return 1 - (1 - t) ** p
    if c in {"easeinoutquad", "easeinoutcubic", "easeinoutquart", "easeinoutquint"}:
        p = {"easeinoutquad": 2, "easeinoutcubic": 3, "easeinoutquart": 4, "easeinoutquint": 5}[c]
        if t < 0.5:
            return (2 * t) ** p / 2
        return 1 - ((-2 * t + 2) ** p) / 2
    # Expo
    if c == "easeinexpo":
        return 0 if t == 0 else 2 ** (10 * t - 10)
    if c == "easeoutexpo":
        return 1 if t == 1 else 1 - 2 ** (-10 * t)
    if c == "easeinoutexpo":
        if t == 0 or t == 1:
            return t
        if t < 0.5:
            return (2 ** (20 * t - 10)) / 2
        return (2 - 2 ** (-20 * t + 10)) / 2
    # Circ
    if c == "easeincirc":
        return 1 - math.sqrt(1 - t * t)
    if c == "easeoutcirc":
        return math.sqrt(1 - (t - 1) ** 2)
    if c == "easeinoutcirc":
        if t < 0.5:
            return (1 - math.sqrt(1 - (2 * t) ** 2)) / 2
        return (math.sqrt(1 - (-2 * t + 2) ** 2) + 1) / 2
    # Back
    if c == "easeinback":
        k = back_k
        return (k + 1) * t ** 3 - k * t ** 2
    if c == "easeoutback":
        k = back_k
        return 1 + (k + 1) * (t - 1) ** 3 + k * (t - 1) ** 2
    if c == "easeinoutback":
        k = back_k * 1.525
        if t < 0.5:
            return ((2 * t) ** 2 * ((k + 1) * 2 * t - k)) / 2
        return (((2 * t - 2) ** 2) * ((k + 1) * (t * 2 - 2) + k) + 2) / 2
    # Cubic Bezier (optional) – approximate on y
    if c == "cubicbezier" and bezier is not None:
        p1x, p1y, p2x, p2y = bezier
        u = 1 - t
        by = (u ** 3) * 0 + 3 * (u ** 2) * t * p1y + 3 * u * (t ** 2) * p2y + (t ** 3) * 1
        return by
    return t


def affine_step_schedule(steps: int, sched: dict):
    """Build per-step multipliers [0..1] using easing and shaping options dict.
    Keys: start, end, bias, exponent, start_offset, end_offset, curve, back_k, bezier.
    """
    start = float(sched.get("start", 0.2))
    end = float(sched.get("end", 0.8))
    bias = float(sched.get("bias", 0.5))
    exponent = float(sched.get("exponent", 1.0))
    start_offset = float(sched.get("start_offset", 0.0))
    end_offset = float(sched.get("end_offset", 0.0))
    curve = str(sched.get("curve", "easeInOutSine"))
    back_k = float(sched.get("back_k", 1.70158))
    bez = sched.get("bezier", None)

    start = min(max(0.0, start), 1.0)
    end = min(max(0.0, end), 1.0)
    if steps <= 0:
        return np.zeros((0,), dtype=np.float32)
    start, end = (start, end) if start <= end else (end, start)
    mid = start + bias * (end - start)
    multipliers = np.zeros(steps, dtype=np.float32)

    start_idx, mid_idx, end_idx = [int(round(x * max(steps - 1, 0))) for x in [start, mid, end]]

    # Rising 0->1 eased
    if mid_idx >= start_idx:
        count = (mid_idx - start_idx) + 1
        ts = np.linspace(0.0, 1.0, count, dtype=np.float32)
        ev = np.array([_ease_value(float(t), curve, back_k=back_k, bezier=bez) for t in ts], dtype=np.float32)
        ev = np.power(ev, max(0.0, exponent))
        if ev.size > 0:
            ev = ev * (1.0 - start_offset) + start_offset
        multipliers[start_idx:mid_idx + 1] = ev

    # Falling 1->0 eased (mirror)
    if end_idx >= mid_idx:
        count = (end_idx - mid_idx) + 1
        ts = np.linspace(1.0, 0.0, count, dtype=np.float32)
        ev = np.array([_ease_value(float(t), curve, back_k=back_k, bezier=bez) for t in ts], dtype=np.float32)
        ev = np.power(ev, max(0.0, exponent))
        if ev.size > 0:
            ev = ev * (1.0 - end_offset) + end_offset
        multipliers[mid_idx:end_idx + 1] = ev

    multipliers[:start_idx] = start_offset
    multipliers[end_idx + 1:] = end_offset
    return multipliers.astype(np.float32)


def _rand_generator(seed: int, device):
    g = torch.Generator(device=device)
    if seed is not None:
        g.manual_seed(seed)
    return g


def _chunked_pairwise_min_dists(px: torch.Tensor, pts: torch.Tensor, k: int = 1, chunk_pts: int = 256, chunk_px: int = 65536):
    """
    Compute k smallest distances from each pixel to points with chunking.
    px: (P,2), pts: (N,2)
    Returns f1 (P,), optional f2 (P,) if k>=2
    """
    P = px.shape[0]
    device = px.device
    dtype = px.dtype
    f = [torch.full((P,), float("inf"), device=device, dtype=dtype) for _ in range(k)]
    for pi in range(0, P, chunk_px):
        pchunk = px[pi:pi+chunk_px]  # (Pc,2)
        # current bests for this slice
        best = [torch.full((pchunk.shape[0],), float("inf"), device=device, dtype=dtype) for _ in range(k)]
        for i in range(0, pts.shape[0], chunk_pts):
            chunk = pts[i:i+chunk_pts]  # (C,2)
            # distances squared
            d2 = (pchunk[:, None, 0] - chunk[None, :, 0])**2 + (pchunk[:, None, 1] - chunk[None, :, 1])**2
            d = torch.sqrt(d2 + 1e-12)
            # merge with current bests
            merged = torch.cat([d] + [b.view(-1, 1) for b in best], dim=1)
            vals, _ = torch.topk(merged, k=k, dim=1, largest=False)
            for j in range(k):
                best[j] = vals[:, j]
        # write back
        for j in range(k):
            f[j][pi:pi+best[j].shape[0]] = best[j]
    return tuple(f)


def poisson_blue_mask_2d(h: int, w: int, radius_px: float, softness: float, seed: int, device, dtype):
    """
    Generate a Poisson-disk point pattern, then map distance-to-nearest into [0,1].
    softness controls how distance maps to mask: m = 1 - exp(-D / softness).
    """
    g = _rand_generator(seed, device)
    # Bridson dart throwing with grid accel
    r = max(1.0, float(radius_px))
    cell = r / math.sqrt(2.0)
    gw, gh = int(math.ceil(w / cell)), int(math.ceil(h / cell))
    grid = -torch.ones((gh, gw), device=device, dtype=torch.int32)
    points = []

    def grid_coords(x, y):
        return int(y // cell), int(x // cell)

    def fits(x, y):
        gi, gj = grid_coords(x, y)
        i0, i1 = max(0, gi - 2), min(gh - 1, gi + 2)
        j0, j1 = max(0, gj - 2), min(gw - 1, gj + 2)
        for ii in range(i0, i1 + 1):
            for jj in range(j0, j1 + 1):
                idx = int(grid[ii, jj].item())
                if idx >= 0:
                    px, py = points[idx]
                    if (px - x) ** 2 + (py - y) ** 2 < r * r:
                        return False
        return True

    # initial point
    x0 = torch.rand((), generator=g, device=device).item() * w
    y0 = torch.rand((), generator=g, device=device).item() * h
    points.append((x0, y0))
    gi, gj = grid_coords(x0, y0)
    grid[gi, gj] = 0
    active = [0]

    k = 30
    max_points = int(4.0 * (h * w) / (r * r + 1e-6))  # heuristic cap
    while active and len(points) < max_points:
        ai = int(active[torch.randint(0, len(active), (1,), generator=g, device=device).item()])
        cx, cy = points[ai]
        placed = False
        for _ in range(k):
            rad = r * (1.0 + torch.rand((), generator=g, device=device).item())
            ang = 2 * math.pi * torch.rand((), generator=g, device=device).item()
            nx = cx + rad * math.cos(ang)
            ny = cy + rad * math.sin(ang)
            if 0 <= nx < w and 0 <= ny < h and fits(nx, ny):
                points.append((nx, ny))
                gi, gj = grid_coords(nx, ny)
                grid[gi, gj] = len(points) - 1
                active.append(len(points) - 1)
                placed = True
                break
        if not placed:
            active.remove(ai)

    if len(points) == 0:
        return torch.zeros((h, w), device=device, dtype=dtype)

    pts = torch.tensor(points, device=device, dtype=dtype)
    yy, xx = torch.meshgrid(torch.arange(h, device=device, dtype=dtype),
                            torch.arange(w, device=device, dtype=dtype), indexing="ij")
    px = torch.stack([xx.reshape(-1), yy.reshape(-1)], dim=1)
    (f1,) = _chunked_pairwise_min_dists(px, pts, k=1)
    D = f1.reshape(h, w)
    m = 1.0 - torch.exp(-D / max(1e-6, float(softness)))
    m = (m - m.min()) / (m.max() - m.min() + 1e-12)
    return m


def worley_edges_2d(h: int, w: int, points_per_kpx: float, metric: str, edge_sharpness: float, seed: int, device, dtype):
    """
    Worley cellular edges via F1/F2 distances. Edge strength ~ 1 - (F2-F1)/norm.
    metric: 'L2' or 'L1'.
    """
    g = _rand_generator(seed, device)
    npts = max(1, int(points_per_kpx * (h * w) / 1000.0))
    pts = torch.stack([
        torch.rand((npts,), generator=g, device=device, dtype=dtype) * w,
        torch.rand((npts,), generator=g, device=device, dtype=dtype) * h,
    ], dim=1)
    yy, xx = torch.meshgrid(torch.arange(h, device=device, dtype=dtype),
                            torch.arange(w, device=device, dtype=dtype), indexing="ij")
    px = torch.stack([xx.reshape(-1), yy.reshape(-1)], dim=1)

    # Distance function
    def dist_fn(a, b):
        if metric.upper() == "L1":
            d = torch.abs(a[:, None, :] - b[None, :, :]).sum(dim=2)
        else:
            d = torch.sqrt(((a[:, None, :] - b[None, :, :]) ** 2).sum(dim=2) + 1e-12)
        return d

    # chunked keep two smallest distances per pixel
    P = px.shape[0]
    f1 = torch.full((P,), float("inf"), device=device, dtype=dtype)
    f2 = torch.full((P,), float("inf"), device=device, dtype=dtype)
    chunk_pts = 256
    chunk_px = 65536
    for pi in range(0, P, chunk_px):
        pchunk = px[pi:pi+chunk_px]
        bf1 = torch.full((pchunk.shape[0],), float("inf"), device=device, dtype=dtype)
        bf2 = torch.full((pchunk.shape[0],), float("inf"), device=device, dtype=dtype)
        for i in range(0, pts.shape[0], chunk_pts):
            chunk = pts[i:i+chunk_pts]
            d = dist_fn(pchunk, chunk)
            merged = torch.cat([d, bf1.view(-1, 1), bf2.view(-1, 1)], dim=1)
            vals, _ = torch.topk(merged, k=2, dim=1, largest=False)
            bf1, bf2 = vals[:, 0], vals[:, 1]
        f1[pi:pi+bf1.shape[0]] = bf1
        f2[pi:pi+bf2.shape[0]] = bf2

    F1 = f1.reshape(h, w)
    F2 = f2.reshape(h, w)
    # Edge measure: small gap between F2 and F1 -> on boundaries
    gap = (F2 - F1)
    gap = gap / (gap.max() + 1e-12)
    edges = 1.0 - gap
    # sharpen
    s = max(0.1, float(edge_sharpness))
    edges = torch.clamp(edges, 0.0, 1.0) ** s
    edges = (edges - edges.min()) / (edges.max() - edges.min() + 1e-12)
    return edges


def tile_oriented_lines_2d(h: int, w: int, tile_size: int, freq_cyc_px: float, jitter: float, seed: int, device, dtype):
    """
    Randomized oriented sinusoid per tile, stitched together; normalized to [0,1].
    """
    s = max(2, int(tile_size))
    g = _rand_generator(seed, device)
    yy, xx = torch.meshgrid(torch.arange(h, device=device, dtype=dtype),
                            torch.arange(w, device=device, dtype=dtype), indexing="ij")
    out = torch.zeros((h, w), device=device, dtype=dtype)
    two_pi = 2.0 * math.pi
    f = float(freq_cyc_px)
    for y0 in range(0, h, s):
        for x0 in range(0, w, s):
            y1 = min(h, y0 + s)
            x1 = min(w, x0 + s)
            th = torch.rand((), generator=g, device=device, dtype=dtype).item() * two_pi
            ph = (torch.rand((), generator=g, device=device, dtype=dtype).item() - 0.5) * two_pi * float(jitter)
            sl = slice(y0, y1)
            sc = slice(x0, x1)
            proj = (xx[sl, sc] - x0) * math.cos(th) + (yy[sl, sc] - y0) * math.sin(th)
            gr = torch.sin(two_pi * f * proj + ph)
            out[sl, sc] = gr
    out = (out - out.min()) / (out.max() - out.min() + 1e-12)
    return out


def dot_screen_jitter_2d(h: int, w: int, cell_size: int, jitter_px: float, fill_ratio: float, seed: int, device, dtype):
    """
    Halftone dot lattice with per-cell jitter of center and radius.
    fill_ratio ~ area fraction per cell (0..1).
    """
    cs = max(2, int(cell_size))
    g = _rand_generator(seed, device)
    yy, xx = torch.meshgrid(torch.arange(h, device=device, dtype=dtype),
                            torch.arange(w, device=device, dtype=dtype), indexing="ij")
    mask = torch.zeros((h, w), device=device, dtype=dtype)
    base_r = 0.5 * math.sqrt(fill_ratio) * cs
    for y0 in range(0, h, cs):
        for x0 in range(0, w, cs):
            y1 = min(h, y0 + cs)
            x1 = min(w, x0 + cs)
            jx = (torch.rand((), generator=g, device=device, dtype=dtype).item() - 0.5) * 2.0 * float(jitter_px)
            jy = (torch.rand((), generator=g, device=device, dtype=dtype).item() - 0.5) * 2.0 * float(jitter_px)
            cx = x0 + cs * 0.5 + jx
            cy = y0 + cs * 0.5 + jy
            rj = base_r * (0.8 + 0.4 * torch.rand((), generator=g, device=device, dtype=dtype).item())
            sl = slice(y0, y1)
            sc = slice(x0, x1)
            d2 = (xx[sl, sc] - cx) ** 2 + (yy[sl, sc] - cy) ** 2
            mask[sl, sc] = torch.where(d2 <= rj * rj, torch.tensor(1.0, device=device, dtype=dtype), torch.tensor(0.0, device=device, dtype=dtype))
    # Slight blur to smooth hard edges
    mask = gaussian_blur(mask.unsqueeze(0).unsqueeze(0), ksize=3, sigma=0.6).squeeze(0).squeeze(0)
    mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-12)
    return mask


def _grid(h: int, w: int, device, dtype):
    yy = torch.arange(h, device=device, dtype=dtype).view(h, 1).expand(h, w)
    xx = torch.arange(w, device=device, dtype=dtype).view(1, w).expand(h, w)
    return xx, yy


def cross_hatch_2d(
    h: int,
    w: int,
    freq_cyc_px: float,
    angles_deg: tuple,
    square: bool,
    phase_jitter: float,
    supersample: int,
    seed: int,
    device,
    dtype,
):
    """
    Generate cross-hatch gratings at given frequency (cycles/pixel) and angles.
    Optionally square-wave, phase jitter, and supersample then downsample.
    Returns normalized [0,1].
    """
    s = max(1, int(supersample))
    H, W = h * s, w * s
    rng = torch.Generator(device=device)
    if seed is not None:
        rng.manual_seed(seed)
    xx, yy = _grid(H, W, device, dtype)

    two_pi = 2.0 * math.pi
    f = float(freq_cyc_px)
    # Allow aliasing if f > 0.5 (user-controlled)
    acc = torch.zeros((H, W), device=device, dtype=dtype)
    ad = angles_deg if isinstance(angles_deg, (list, tuple)) else [0.0, 90.0]
    for a in ad:
        th = math.radians(float(a))
        # Phase jitter in [0, 2π*phase_jitter)
        phi = 0.0
        if phase_jitter > 0.0:
            phi = torch.rand((), generator=rng, device=device, dtype=dtype).item() * two_pi * float(phase_jitter)
        proj = xx * math.cos(th) + yy * math.sin(th)
        gr = torch.sin(two_pi * f * proj + phi)
        if square:
            gr = torch.sign(gr)
        acc = acc + gr
    # Normalize to [0,1]
    if len(ad) > 0:
        acc = acc / float(len(ad))
    acc = (acc - acc.min()) / (acc.max() - acc.min() + 1e-12)

    if s > 1:
        acc = acc.unsqueeze(0).unsqueeze(0)
        # Light blur before downsample to reduce aliasing
        acc = gaussian_blur(acc, ksize=3, sigma=0.6)
        acc = F.interpolate(acc, size=(h, w), mode="area")
        acc = acc.squeeze(0).squeeze(0)
        acc = (acc - acc.min()) / (acc.max() - acc.min() + 1e-12)
    return acc


def highpass_white_2d(h: int, w: int, cutoff_frac: float, order: int, seed: int, device, dtype):
    """
    White noise shaped by Butterworth high-pass filter.
    cutoff_frac: 0..1 fraction of Nyquist radius.
    order: filter order (>=1).
    """
    rng = torch.Generator(device=device)
    if seed is not None:
        rng.manual_seed(seed)
    base = torch.randn((h, w), generator=rng, device=device, dtype=dtype)
    F = torch.fft.fft2(base)
    r = _radial_frequency_grid(h, w, device, dtype)
    r = torch.clamp(r, min=1e-6)
    rmax = r.max().clamp(min=1e-6)
    rc = torch.as_tensor(cutoff_frac, device=device, dtype=dtype).clamp(1e-6, 1.0) * rmax
    n = max(1, int(order))
    # Butterworth high-pass magnitude response
    H = 1.0 / (1.0 + torch.pow(rc / r, 2 * n))
    H = H.to(F.dtype)
    F_hp = F * H
    x = torch.fft.ifft2(F_hp).real
    x = (x - x.min()) / (x.max() - x.min() + 1e-12)
    return x


def ring_noise_2d(h: int, w: int, center_frac: float, bandwidth_frac: float, seed: int, device, dtype):
    """
    Narrow annulus band-pass white noise around center_frac of Nyquist radius.
    """
    rng = torch.Generator(device=device)
    if seed is not None:
        rng.manual_seed(seed)
    base = torch.randn((h, w), generator=rng, device=device, dtype=dtype)
    F = torch.fft.fft2(base)
    r = _radial_frequency_grid(h, w, device, dtype)
    rmax = r.max().clamp(min=1e-6)
    c = torch.as_tensor(center_frac, device=device, dtype=dtype).clamp(0.0, 1.0) * rmax
    bw = torch.as_tensor(bandwidth_frac, device=device, dtype=dtype).clamp(1e-6, 1.0) * rmax
    # Gaussian ring
    sigma = bw
    mask = torch.exp(-0.5 * ((r - c) / sigma) ** 2)
    F_ring = F * mask
    x = torch.fft.ifft2(F_ring).real
    x = (x - x.min()) / (x.max() - x.min() + 1e-12)
    return x


def pink_noise_2d(h: int, w: int, seed: int, device, dtype):
    return spectral_noise(h, w, beta=-1.0, seed=seed, device=device, dtype=dtype)


def brown_noise_2d(h: int, w: int, seed: int, device, dtype):
    # Also known as red noise
    return spectral_noise(h, w, beta=-2.0, seed=seed, device=device, dtype=dtype)


def blue_noise_2d(h: int, w: int, seed: int, device, dtype):
    return spectral_noise(h, w, beta=+1.0, seed=seed, device=device, dtype=dtype)


def violet_noise_2d(h: int, w: int, seed: int, device, dtype):
    # Also called purple noise
    return spectral_noise(h, w, beta=+2.0, seed=seed, device=device, dtype=dtype)


def velvet_noise(h: int, w: int, tap_count: int, seed: int, device, dtype):
    """
    Generate 2D velvet noise: sparse random ±1 impulses, returned in [0,1].

    tap_count: number of impulses to place. Values are clipped to [1, h*w].
    """
    tap_count = int(max(1, min(tap_count, h * w)))
    rng = torch.Generator(device=device)
    if seed is not None:
        rng.manual_seed(seed)

    v = torch.zeros((h, w), device=device, dtype=dtype)
    perm = torch.randperm(h * w, generator=rng, device=device)[:tap_count]
    signs = (torch.randint(0, 2, (tap_count,), generator=rng, device=device) * 2 - 1).to(dtype)
    flat = v.view(-1)
    flat[perm] = signs

    v = (v + 1.0) * 0.5
    return v


def green_noise_2d(h: int, w: int, seed: int, device, dtype, center_frac: float = 0.35, bandwidth_frac: float = 0.15):
    """
    Generate green noise as a mid-frequency band-pass shaped white noise.

    center_frac, bandwidth_frac are relative to max radial frequency (rmax). Typical values:
      center_frac ~ 0.3-0.4, bandwidth_frac ~ 0.1-0.2
    """
    rng = torch.Generator(device=device)
    if seed is not None:
        rng.manual_seed(seed)
    base = torch.randn((h, w), generator=rng, device=device, dtype=dtype)
    F = torch.fft.fft2(base)
    r = _radial_frequency_grid(h, w, device, dtype)
    rmax = r.max().clamp(min=1e-6)
    c = torch.as_tensor(center_frac, device=device, dtype=dtype).clamp(0.0, 1.0) * rmax
    bw = torch.as_tensor(bandwidth_frac, device=device, dtype=dtype).clamp(1e-6, 1.0) * rmax
    sigma = bw
    wmask = torch.exp(-0.5 * ((r - c) / sigma) ** 2)
    F_bp = F * wmask
    x = torch.fft.ifft2(F_bp).real
    x = (x - x.min()) / (x.max() - x.min() + 1e-12)
    return x


def black_noise_2d(h: int, w: int, seed: int, device, dtype, bins: int = 2048):
    """
    Generate black noise as sparse narrowband content: keep a small number of
    frequency bins (in rFFT domain), zero elsewhere, then inverse transform.
    """
    rng = torch.Generator(device=device)
    if seed is not None:
        rng.manual_seed(seed)

    spec_shape = (h, w // 2 + 1)
    total_bins = spec_shape[0] * spec_shape[1]
    k = int(max(1, min(bins, total_bins - 1)))

    spec = torch.zeros(spec_shape, device=device, dtype=torch.complex64)
    idx = torch.randperm(total_bins, generator=rng, device=device)
    idx = idx[idx != 0]
    idx = idx[:k]

    real = torch.randn((k,), generator=rng, device=device)
    imag = torch.randn((k,), generator=rng, device=device)
    vals = torch.complex(real, imag)

    flat = spec.view(-1)
    flat[idx] = vals
    spec = flat.view(spec_shape)

    x = torch.fft.irfft2(spec, s=(h, w))
    x = x.to(dtype)
    x = (x - x.min()) / (x.max() - x.min() + 1e-12)
    return x


def get_cfg_for_step(cfg_values, step_index, total_steps):
    """
    Get the CFG value for a specific step from a list of CFG values.
    If step_index exceeds the list length, repeat the last value.
    
    Args:
        cfg_values: List of CFG float values or single float value
        step_index: Current step index (0-based)
        total_steps: Total number of steps
        
    Returns:
        float: CFG value for the current step
    """
    if isinstance(cfg_values, (int, float)):
        return float(cfg_values)
    
    if isinstance(cfg_values, (list, tuple)) and len(cfg_values) > 0:
        if step_index < len(cfg_values):
            return float(cfg_values[step_index])
        else:
            return float(cfg_values[-1])
    return 4.5