"""
Finite-difference stencil operators for lattice fields.

This module provides a unified implementation of spatial derivative operators
using finite difference stencils. These work for any dimension (1D, 2D, 3D)
with periodic boundary conditions.

For FFT-based spectral methods, see `spectral.py`.

All operators are JIT-compiled for performance.
"""

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


@jit
def laplacian(field: Array, dx: float) -> Array:
    """
    Compute Laplacian ∇²φ using second-order finite differences.

    Uses the standard 3-point stencil in each dimension:
    ∇²φ = Σ_i [φ(n+i) + φ(n-i) - 2φ(n)] / dx²

    Works for 1D, 2D, or 3D fields with periodic boundary conditions.

    Args:
        field: Field array (1D, 2D, or 3D)
        dx: Lattice spacing (assumed uniform in all directions)

    Returns:
        Laplacian at each lattice site (same shape as input)
    """
    result = jnp.zeros_like(field)
    ndim = field.ndim

    for axis in range(ndim):
        result = result + (
            jnp.roll(field, -1, axis=axis)
            + jnp.roll(field, 1, axis=axis)
            - 2.0 * field
        )

    return result / (dx**2)


@jit
def gradient_squared(field: Array, dx: float) -> Array:
    """
    Compute |∇φ|² using centered finite differences.

    |∇φ|² = Σ_i |∂_iφ|² where ∂_iφ = [φ(n+i) - φ(n-i)] / (2*dx)

    Works for 1D, 2D, or 3D fields. For complex fields, returns |∇φ|².

    Args:
        field: Field array (1D, 2D, or 3D), can be real or complex
        dx: Lattice spacing

    Returns:
        Gradient squared at each site (real-valued, same shape as input)
    """
    grad_sq = jnp.zeros(field.shape, dtype=jnp.float64)
    ndim = field.ndim

    for axis in range(ndim):
        derivative = (
            jnp.roll(field, -1, axis=axis) - jnp.roll(field, 1, axis=axis)
        ) / (2.0 * dx)
        grad_sq = grad_sq + jnp.abs(derivative) ** 2

    return grad_sq


@jit
def gradient(field: Array, dx: float) -> Tuple[Array, ...]:
    """
    Compute gradient ∇φ using centered finite differences.

    Returns a tuple of arrays, one for each spatial direction.

    Args:
        field: Field array (1D, 2D, or 3D)
        dx: Lattice spacing

    Returns:
        Tuple of gradient components (∂_x φ, ∂_y φ, ...) for each dimension
    """
    ndim = field.ndim
    grads = []

    for axis in range(ndim):
        grad_i = (
            jnp.roll(field, -1, axis=axis) - jnp.roll(field, 1, axis=axis)
        ) / (2.0 * dx)
        grads.append(grad_i)

    return tuple(grads)


@jit
def divergence(vector_field: Array, dx: float) -> Array:
    """
    Compute divergence ∇·F for a vector field.

    Args:
        vector_field: Array of shape (ndim, *spatial_shape)
            First axis is the vector component index.
        dx: Lattice spacing

    Returns:
        Divergence at each site (shape = spatial_shape)
    """
    ndim = vector_field.shape[0]
    div = jnp.zeros(vector_field.shape[1:])

    for i in range(ndim):
        # Central difference for ∂_i F_i
        div = div + (
            jnp.roll(vector_field[i], -1, axis=i)
            - jnp.roll(vector_field[i], 1, axis=i)
        ) / (2.0 * dx)

    return div


@jit
def forward_gradient(field: Array, dx: float, axis: int) -> Array:
    """
    Compute forward finite difference ∂_i φ = [φ(n+i) - φ(n)] / dx.

    Useful for gauge-covariant derivatives and link-based calculations.

    Args:
        field: Field array
        dx: Lattice spacing
        axis: Direction of derivative (0, 1, or 2)

    Returns:
        Forward derivative along specified axis
    """
    return (jnp.roll(field, -1, axis=axis) - field) / dx


@jit
def backward_gradient(field: Array, dx: float, axis: int) -> Array:
    """
    Compute backward finite difference ∂_i φ = [φ(n) - φ(n-i)] / dx.

    Args:
        field: Field array
        dx: Lattice spacing
        axis: Direction of derivative (0, 1, or 2)

    Returns:
        Backward derivative along specified axis
    """
    return (field - jnp.roll(field, 1, axis=axis)) / dx
