"""Perlin noise-based synthetic anomaly generator.

This module provides functionality to generate synthetic anomalies using Perlin noise
patterns. The generator can create both noise-based and image-based anomalies with
configurable parameters.

Example:
    >>> from anomalib.data.utils.generators.perlin import generate_perlin_noise
    >>> import torch
    >>> # Generate 256x256 noise with default random scale
    >>> noise = generate_perlin_noise(256, 256)
    >>> print(noise.shape)
    torch.Size([256, 256])

    >>> # Generate 512x512 noise with fixed scale
    >>> noise = generate_perlin_noise(512, 512, scale=(8, 8))
    >>> print(noise.shape)
    torch.Size([512, 512])

    >>> # Generate noise on GPU if available
    >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    >>> noise = generate_perlin_noise(128, 128, device=device)
"""

# Copyright (C) 2022-2025 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from pathlib import Path

import torch
from torchvision import io
from torchvision.datasets.folder import IMG_EXTENSIONS
from torchvision.transforms import v2


def generate_perlin_noise(
    height: int,
    width: int,
    scale: tuple[int, int] = None,
    device: torch.device = None,
) -> torch.Tensor:
    """Generate a Perlin noise pattern.

    This function generates a Perlin noise pattern using a grid-based gradient noise
    approach. The noise is generated by interpolating between randomly generated
    gradient vectors at grid vertices. The interpolation uses a quintic curve for
    smooth transitions.

    Args:
        height: Desired height of the noise pattern.
        width: Desired width of the noise pattern.
        scale: Tuple of ``(scale_x, scale_y)`` for noise granularity. If ``None``,
            random scales will be used. Larger scales produce coarser noise patterns,
            while smaller scales produce finer patterns.
        device: Device to generate the noise on. If ``None``, uses current default
            device.

    Returns:
        torch.Tensor: Tensor of shape ``[height, width]`` containing the noise
            pattern, with values roughly in ``[-1, 1]`` range.

    Example:
        >>> # Generate 256x256 noise with default random scale
        >>> noise = generate_perlin_noise(256, 256)
        >>> print(noise.shape)
        torch.Size([256, 256])

        >>> # Generate 512x512 noise with fixed scale
        >>> noise = generate_perlin_noise(512, 512, scale=(8, 8))
        >>> print(noise.shape)
        torch.Size([512, 512])

        >>> # Generate noise on GPU if available
        >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        >>> noise = generate_perlin_noise(128, 128, device=device)
    """
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Handle scale parameter
    if scale is None:
        min_scale, max_scale = 0, 6
        scalex = 2 ** torch.randint(min_scale, max_scale, (1,), device=device).item()
        scaley = 2 ** torch.randint(min_scale, max_scale, (1,), device=device).item()
    else:
        scalex, scaley = scale

    # Ensure dimensions are powers of 2 for proper noise generation
    def nextpow2(value: int) -> int:
        return int(2 ** torch.ceil(torch.log2(torch.tensor(value))).int().item())

    pad_h = nextpow2(height)
    pad_w = nextpow2(width)

    # Generate base grid
    delta = (scalex / pad_h, scaley / pad_w)
    d = (pad_h // scalex, pad_w // scaley)

    grid = (
        torch.stack(
            torch.meshgrid(
                torch.arange(0, scalex, delta[0], device=device),
                torch.arange(0, scaley, delta[1], device=device),
                indexing="ij",
            ),
            dim=-1,
        )
        % 1
    )

    # Generate random gradients
    angles = 2 * torch.pi * torch.rand(int(scalex) + 1, int(scaley) + 1, device=device)
    gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim=-1)

    def tile_grads(slice1: list[int], slice2: list[int]) -> torch.Tensor:
        return (
            gradients[slice1[0] : slice1[1], slice2[0] : slice2[1]]
            .repeat_interleave(int(d[0]), 0)
            .repeat_interleave(int(d[1]), 1)
        )

    def dot(grad: torch.Tensor, shift: list[float]) -> torch.Tensor:
        return (
            torch.stack(
                (grid[:pad_h, :pad_w, 0] + shift[0], grid[:pad_h, :pad_w, 1] + shift[1]),
                dim=-1,
            )
            * grad[:pad_h, :pad_w]
        ).sum(dim=-1)

    # Calculate noise values at grid points
    n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0])
    n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0])
    n01 = dot(tile_grads([0, -1], [1, None]), [0, -1])
    n11 = dot(tile_grads([1, None], [1, None]), [-1, -1])

    # Interpolate between grid points using quintic curve
    def fade(t: torch.Tensor) -> torch.Tensor:
        return 6 * t**5 - 15 * t**4 + 10 * t**3

    t = fade(grid[:pad_h, :pad_w])
    noise = torch.sqrt(torch.tensor(2.0, device=device)) * torch.lerp(
        torch.lerp(n00, n10, t[..., 0]),
        torch.lerp(n01, n11, t[..., 0]),
        t[..., 1],
    )

    # Crop to desired dimensions
    return noise[:height, :width]
