"""
Utility functions for lattice field initialization and diagnostics.

This module provides helper functions for creating initial field
configurations and analyzing simulation results.
"""

import jax.numpy as jnp
from jax import random, Array
from typing import Tuple, Union, Optional, Any
from jaxlatt.core import Lattice


def gaussian_field_1d(
    key: Any,
    size: int,
    length: float,
    amplitude: float = 1.0,
    center: Optional[float] = None,
    width: Optional[float] = None,
) -> Array:
    """
    Create a 1D Gaussian field profile.

    φ(x) = A * exp(-(x - x₀)² / (2σ²))

    Args:
        key: JAX random key (unused, for API consistency)
        size: Number of lattice points
        length: Physical length of lattice
        amplitude: Amplitude of the Gaussian (default: 1.0)
        center: Center position (default: length/2)
        width: Width σ of the Gaussian (default: length/10)

    Returns:
        1D array with Gaussian profile
    """
    if center is None:
        center = length / 2
    if width is None:
        width = length / 10

    x = jnp.linspace(0, length, size, endpoint=False)
    field = amplitude * jnp.exp(-((x - center) ** 2) / (2 * width**2))

    return field


def gaussian_field_2d(
    key: Any,
    size: Tuple[int, int],
    length: Union[float, Tuple[float, float]],
    amplitude: float = 1.0,
    center: Optional[Tuple[float, float]] = None,
    width: Optional[float] = None,
) -> Array:
    """
    Create a 2D Gaussian field profile.

    φ(x, y) = A * exp(-((x - x₀)² + (y - y₀)²) / (2σ²))

    Args:
        key: JAX random key (unused)
        size: Tuple of (nx, ny) lattice points
        length: Physical length (float for square, tuple for rectangle)
        amplitude: Amplitude of the Gaussian
        center: Center position (x₀, y₀)
        width: Width σ of the Gaussian

    Returns:
        2D array with Gaussian profile
    """
    if isinstance(length, (int, float)):
        lx, ly = length, length
    else:
        lx, ly = length

    if center is None:
        center = (lx / 2, ly / 2)
    if width is None:
        width = lx / 10

    nx, ny = size
    x = jnp.linspace(0, lx, nx, endpoint=False)
    y = jnp.linspace(0, ly, ny, endpoint=False)

    X, Y = jnp.meshgrid(x, y, indexing="ij")

    x0, y0 = center
    r_squared = (X - x0) ** 2 + (Y - y0) ** 2
    field = amplitude * jnp.exp(-r_squared / (2 * width**2))

    return field


def gaussian_field_3d(
    key: Any,
    size: Tuple[int, int, int],
    length: Union[float, Tuple[float, float, float]],
    amplitude: float = 1.0,
    center: Optional[Tuple[float, float, float]] = None,
    width: Optional[float] = None,
) -> Array:
    """
    Create a 3D Gaussian field profile.

    φ(x,y,z) = A * exp(-((x-x0)² + (y-y0)² + (z-z0)²) / (2σ²))

    Args:
        key: JAX random key (unused)
        size: (nx, ny, nz)
        length: Physical length (float for cubic, tuple for rectangular box)
        amplitude: Amplitude A
        center: (x0,y0,z0) center position (default middle)
        width: Gaussian width σ (default: lx/10)

    Returns:
        3D array with Gaussian profile
    """
    if isinstance(length, (int, float)):
        lx = ly = lz = float(length)
    else:
        lx, ly, lz = length
    if center is None:
        center = (lx / 2, ly / 2, lz / 2)
    if width is None:
        width = lx / 10
    nx, ny, nz = size
    x = jnp.linspace(0, lx, nx, endpoint=False)
    y = jnp.linspace(0, ly, ny, endpoint=False)
    z = jnp.linspace(0, lz, nz, endpoint=False)
    X, Y, Z = jnp.meshgrid(x, y, z, indexing="ij")
    x0, y0, z0 = center
    r2 = (X - x0) ** 2 + (Y - y0) ** 2 + (Z - z0) ** 2
    field = amplitude * jnp.exp(-r2 / (2 * width**2))
    return field


def random_field(
    key: Any, size: Union[int, Tuple[int, ...]], amplitude: float = 1.0
) -> Array:
    """
    Create a random field with uniform distribution.

    Args:
        key: JAX random key
        size: Lattice size (int for 1D, tuple for 2D)
        amplitude: Maximum amplitude (field in [-A, A])

    Returns:
        Random field array
    """
    if isinstance(size, int):
        shape = (size,)
    else:
        shape = size

    return amplitude * (2 * random.uniform(key, shape) - 1)


def sine_wave_1d(size: int, length: float, amplitude: float = 1.0, k: int = 1) -> Array:
    """
    Create a 1D sine wave field.

    φ(x) = A * sin(2πkx/L)

    Args:
        size: Number of lattice points
        length: Physical length
        amplitude: Amplitude
        k: Mode number

    Returns:
        1D sine wave array
    """
    x = jnp.linspace(0, length, size, endpoint=False)
    return amplitude * jnp.sin(2 * jnp.pi * k * x / length)


def create_initial_lattice_1d(
    size: int = 128,
    length: float = 10.0,
    field_type: str = "gaussian",
    amplitude: float = 1.0,
    velocity: float = 0.0,
    key: Any = None,
) -> Lattice:
    """
    Create an initial 1D lattice with specified field configuration.

    Args:
        size: Number of lattice points
        length: Physical length
        field_type: Type of initial field ('gaussian', 'sine', 'random', 'zero')
        amplitude: Field amplitude
        velocity: Initial velocity (uniform)
        key: Random key (required for 'random' field)

    Returns:
        Initialized Lattice object
    """
    if field_type == "gaussian":
        if key is None:
            key = random.PRNGKey(0)
        field = gaussian_field_1d(key, size, length, amplitude)
    elif field_type == "sine":
        field = sine_wave_1d(size, length, amplitude, k=1)
    elif field_type == "random":
        if key is None:
            key = random.PRNGKey(0)
        field = random_field(key, size, amplitude)
    elif field_type == "zero":
        field = jnp.zeros(size)
    else:
        raise ValueError(f"Unknown field_type: {field_type}")

    field_dot = jnp.ones(size) * velocity

    return Lattice(size=size, length=length, field=field, field_dot=field_dot)


def create_initial_lattice_2d(
    size: Tuple[int, int] = (64, 64),
    length: Union[float, Tuple[float, float]] = 10.0,
    field_type: str = "gaussian",
    amplitude: float = 1.0,
    velocity: float = 0.0,
    key: Any = None,
) -> Lattice:
    """
    Create an initial 2D lattice with specified field configuration.

    Args:
        size: Tuple of (nx, ny) lattice points
        length: Physical length (float or tuple)
        field_type: Type of initial field
        amplitude: Field amplitude
        velocity: Initial velocity (uniform)
        key: Random key (required for 'random' field)

    Returns:
        Initialized Lattice object
    """
    if field_type == "gaussian":
        if key is None:
            key = random.PRNGKey(0)
        field = gaussian_field_2d(key, size, length, amplitude)
    elif field_type == "random":
        if key is None:
            key = random.PRNGKey(0)
        field = random_field(key, size, amplitude)
    elif field_type == "zero":
        field = jnp.zeros(size)
    else:
        raise ValueError(f"Unknown field_type: {field_type}")

    field_dot = jnp.ones(size) * velocity

    return Lattice(size=size, length=length, field=field, field_dot=field_dot)


def create_initial_lattice_3d(
    size: Tuple[int, int, int] = (32, 32, 32),
    length: Union[float, Tuple[float, float, float]] = 10.0,
    field_type: str = "gaussian",
    amplitude: float = 1.0,
    velocity: float = 0.0,
    key: Any = None,
) -> Lattice:
    """
    Create an initial 3D lattice with specified field configuration.

    Args:
        size: (nx, ny, nz)
        length: Physical length (float cubic or tuple)
        field_type: 'gaussian', 'random', 'zero'
        amplitude: Field amplitude
        velocity: Initial uniform velocity value for φ̇
        key: Random key required for 'random'

    Returns:
        Initialized Lattice object (ndim=3)
    """
    # Normalize length to tuple of 3 floats
    if isinstance(length, (int, float)):
        length = (float(length), float(length), float(length))

    if field_type == "gaussian":
        if key is None:
            key = random.PRNGKey(0)
        field = gaussian_field_3d(key, size, length, amplitude)
    elif field_type == "random":
        if key is None:
            key = random.PRNGKey(0)
        field = random_field(key, size, amplitude)
    elif field_type == "zero":
        field = jnp.zeros(size)
    else:
        raise ValueError(f"Unknown field_type: {field_type}")
    field_dot = jnp.ones(size) * velocity
    return Lattice(size=size, length=length, field=field, field_dot=field_dot)


def compute_statistics(field: Array) -> dict:
    """
    Compute basic statistics of a field configuration.

    Args:
        field: Field array

    Returns:
        Dictionary with mean, std, min, max
    """
    return {
        "mean": float(jnp.mean(field)),
        "std": float(jnp.std(field)),
        "min": float(jnp.min(field)),
        "max": float(jnp.max(field)),
    }


def check_energy_conservation(energies: list, tolerance: float = 0.01) -> bool:
    """
    Check if energy is conserved within a given tolerance.

    Args:
        energies: List of energy values over time
        tolerance: Relative tolerance (default: 1%)

    Returns:
        True if energy variation is within tolerance
    """
    arr = jnp.array(energies)
    mean_energy = float(jnp.mean(arr))
    variation = float(jnp.std(arr) / mean_energy)
    return bool(variation < tolerance)
