"""
FFT-based spatial derivative operators for lattice fields.

This module implements spatial derivatives using spectral methods (FFT),
which are efficient and accurate for periodic boundary conditions.

Supported dimensions: 1D, 2D, and 3D.
"""

import jax.numpy as jnp
from jax import jit, Array
from typing import Tuple


@jit
def laplacian_1d(field: Array, dx: float) -> Array:
    """
    Compute the Laplacian of a 1D field using FFT.

    For periodic boundary conditions, the Laplacian in Fourier space is:
    ∇²φ(x) = IFFT[-k² * FFT(φ(x))]

    Args:
        field: 1D array of field values
        dx: Lattice spacing

    Returns:
        Laplacian of the field
    """
    n = field.shape[0]
    # Compute k-space coordinates
    k = jnp.fft.fftfreq(n, d=dx) * 2 * jnp.pi

    # Transform to Fourier space
    field_k = jnp.fft.fft(field)

    # Apply Laplacian operator: -k²
    laplacian_k = -(k**2) * field_k

    # Transform back to real space
    laplacian = jnp.fft.ifft(laplacian_k).real

    return laplacian


@jit
def laplacian_2d(field: Array, dx: float) -> Array:
    """
    Compute the Laplacian of a 2D field using FFT.

    For periodic boundary conditions:
    ∇²φ(x,y) = IFFT[-(kx² + ky²) * FFT(φ(x,y))]

    Args:
        field: 2D array of field values
        dx: Lattice spacing (assumed equal in both directions)

    Returns:
        Laplacian of the field
    """
    nx, ny = field.shape

    # Compute k-space coordinates
    kx = jnp.fft.fftfreq(nx, d=dx) * 2 * jnp.pi
    ky = jnp.fft.fftfreq(ny, d=dx) * 2 * jnp.pi

    # Create 2D k-space grid
    KX, KY = jnp.meshgrid(kx, ky, indexing="ij")
    k_squared = KX**2 + KY**2

    # Transform to Fourier space
    field_k = jnp.fft.fft2(field)

    # Apply Laplacian operator: -k²
    laplacian_k = -k_squared * field_k

    # Transform back to real space
    laplacian = jnp.fft.ifft2(laplacian_k).real

    return laplacian


@jit
def laplacian_3d(field: Array, dx: float) -> Array:
    """
    Compute the Laplacian of a 3D field using FFT.

    ∇²φ(x,y,z) = IFFTN[-(kx² + ky² + kz²) * FFTN(φ)]

    Args:
        field: 3D array of field values
        dx: Lattice spacing (assumed equal in all directions)

    Returns:
        Laplacian of the field
    """
    nx, ny, nz = field.shape
    kx = jnp.fft.fftfreq(nx, d=dx) * 2 * jnp.pi
    ky = jnp.fft.fftfreq(ny, d=dx) * 2 * jnp.pi
    kz = jnp.fft.fftfreq(nz, d=dx) * 2 * jnp.pi
    KX, KY, KZ = jnp.meshgrid(kx, ky, kz, indexing="ij")
    k_squared = KX**2 + KY**2 + KZ**2
    field_k = jnp.fft.fftn(field)
    laplacian_k = -k_squared * field_k
    laplacian = jnp.fft.ifftn(laplacian_k).real
    return laplacian


def laplacian(field: Array, dx: float) -> Array:
    """
    Compute the Laplacian of a field (supports 1D, 2D, 3D).

    Args:
        field: Array of field values
        dx: Lattice spacing

    Returns:
        Laplacian of the field
    """
    if field.ndim == 1:
        return laplacian_1d(field, dx)
    elif field.ndim == 2:
        return laplacian_2d(field, dx)
    elif field.ndim == 3:
        return laplacian_3d(field, dx)
    else:
        raise ValueError(
            f"laplacian only supports 1D, 2D, 3D fields, got {field.ndim}D"
        )


@jit
def gradient_1d(field: Array, dx: float) -> Array:
    """
    Compute the gradient of a 1D field using FFT.

    Args:
        field: 1D array of field values
        dx: Lattice spacing

    Returns:
        Gradient of the field
    """
    n = field.shape[0]
    k = jnp.fft.fftfreq(n, d=dx) * 2 * jnp.pi

    field_k = jnp.fft.fft(field)
    gradient_k = 1j * k * field_k
    gradient = jnp.fft.ifft(gradient_k).real

    return gradient


@jit
def gradient_2d(field: Array, dx: float) -> Tuple[Array, Array]:
    """
    Compute the gradient of a 2D field using FFT.

    Args:
        field: 2D array of field values
        dx: Lattice spacing

    Returns:
        Tuple of (gradient_x, gradient_y)
    """
    nx, ny = field.shape

    kx = jnp.fft.fftfreq(nx, d=dx) * 2 * jnp.pi
    ky = jnp.fft.fftfreq(ny, d=dx) * 2 * jnp.pi

    KX, KY = jnp.meshgrid(kx, ky, indexing="ij")

    field_k = jnp.fft.fft2(field)

    gradient_x_k = 1j * KX * field_k
    gradient_y_k = 1j * KY * field_k

    gradient_x = jnp.fft.ifft2(gradient_x_k).real
    gradient_y = jnp.fft.ifft2(gradient_y_k).real

    return gradient_x, gradient_y


@jit
def gradient_energy_1d(field: Array, dx: float) -> Array:
    """
    Compute the gradient energy ½∫(∇φ)² dV for a 1D field.

    Args:
        field: 1D array of field values
        dx: Lattice spacing

    Returns:
        Total gradient energy
    """
    grad = gradient_1d(field, dx)
    return 0.5 * jnp.sum(grad**2) * dx


@jit
def gradient_energy_2d(field: Array, dx: float) -> Array:
    """
    Compute the gradient energy ½∫(∇φ)² dV for a 2D field.

    Args:
        field: 2D array of field values
        dx: Lattice spacing

    Returns:
        Total gradient energy
    """
    grad_x, grad_y = gradient_2d(field, dx)
    return 0.5 * jnp.sum(grad_x**2 + grad_y**2) * dx**2


@jit
def gradient_3d(field: Array, dx: float) -> Tuple[Array, Array, Array]:
    """
    Compute the gradient of a 3D field using FFT.

    Args:
        field: 3D array of field values
        dx: Lattice spacing

    Returns:
        Tuple (∂φ/∂x, ∂φ/∂y, ∂φ/∂z)
    """
    nx, ny, nz = field.shape
    kx = jnp.fft.fftfreq(nx, d=dx) * 2 * jnp.pi
    ky = jnp.fft.fftfreq(ny, d=dx) * 2 * jnp.pi
    kz = jnp.fft.fftfreq(nz, d=dx) * 2 * jnp.pi
    KX, KY, KZ = jnp.meshgrid(kx, ky, kz, indexing="ij")
    field_k = jnp.fft.fftn(field)
    grad_x = jnp.fft.ifftn(1j * KX * field_k).real
    grad_y = jnp.fft.ifftn(1j * KY * field_k).real
    grad_z = jnp.fft.ifftn(1j * KZ * field_k).real
    return grad_x, grad_y, grad_z


@jit
def gradient_energy_3d(field: Array, dx: float) -> Array:
    """
    Compute gradient energy ½∫(∇φ)² dV for 3D field.
    """
    gx, gy, gz = gradient_3d(field, dx)
    return 0.5 * jnp.sum(gx**2 + gy**2 + gz**2) * dx**3


def gradient_energy(field: Array, dx: float) -> Array:
    """
    Compute the gradient energy ½∫(∇φ)² dV (supports 1D, 2D, 3D).

    Args:
        field: Array of field values
        dx: Lattice spacing

    Returns:
        Total gradient energy including 1/2 factor
    """
    if field.ndim == 1:
        return gradient_energy_1d(field, dx)
    elif field.ndim == 2:
        return gradient_energy_2d(field, dx)
    elif field.ndim == 3:
        return gradient_energy_3d(field, dx)
    else:
        raise ValueError("gradient_energy only supports 1D, 2D, 3D fields")
