"""
Power spectrum calculations for scalar fields on lattices.

This module provides functions to compute power spectra P(k) from field
configurations in 1D, 2D, and 3D. The power spectrum measures the distribution
of power across different wavenumbers k, which is fundamental for understanding
field fluctuations and correlations.

Physical interpretation:
- P(k) gives the mean square field amplitude at wavenumber k
- For Gaussian random fields: ⟨φ(k) φ*(k')⟩ = (2π)^d δ(k - k') P(k)
- In expanding universes: P(k,t) tracks growth of structure

Usage:
    >>> from jaxlatt.observables import power_spectrum_3d, bin_power_spectrum
    >>> field = ... # 3D field configuration
    >>> k_values, P_k = power_spectrum_3d(field, L)
    >>> k_binned, P_binned, counts = bin_power_spectrum(k_values, P_k, num_bins=50)
"""

import jax.numpy as jnp
from jax import Array
from typing import Tuple
import numpy as np


def power_spectrum_1d(field: Array, L: float) -> Tuple[Array, Array]:
    """
    Compute 1D power spectrum P(k) from a field configuration.

    The power spectrum is defined as:
        P(k) = |φ̃(k)|² / L

    where φ̃(k) is the Fourier transform of the field.

    Args:
        field: 1D array of field values, shape (N,)
        L: Physical length of the lattice

    Returns:
        Tuple of (k_values, P_k) where:
        - k_values: Array of wavenumbers (positive frequencies only)
        - P_k: Power spectrum values at each k

    Example:
        >>> field = jnp.array([...])  # 1D field
        >>> k, P = power_spectrum_1d(field, L=10.0)
        >>> # Plot: plt.loglog(k, P)
    """
    N = field.shape[0]

    # FFT of field
    field_k = jnp.fft.rfft(field)

    # Wavenumbers (positive frequencies only for rfft)
    k_values = jnp.fft.rfftfreq(N, d=L / N) * 2 * jnp.pi

    # Power spectrum: |φ̃(k)|² / L
    # Factor of 2 for k>0 to account for negative frequencies
    P_k = jnp.abs(field_k) ** 2 / L
    P_k = P_k.at[1:].multiply(2)  # Double power for k>0 (symmetry)

    return k_values, P_k


def power_spectrum_2d(field: Array, L: float) -> Tuple[Array, Array, Array]:
    """
    Compute 2D power spectrum P(k_x, k_y) from a field configuration.

    Returns the full 2D spectrum in k-space. For radially averaged spectrum,
    use bin_power_spectrum() on the output.

    Args:
        field: 2D array of field values, shape (N, N)
        L: Physical length of the lattice (assumed square)

    Returns:
        Tuple of (k_x, k_y, P_k) where:
        - k_x: 2D array of x-wavenumbers
        - k_y: 2D array of y-wavenumbers
        - P_k: 2D power spectrum values

    Example:
        >>> field = jnp.array([...])  # 2D field
        >>> kx, ky, P = power_spectrum_2d(field, L=20.0)
        >>> # Radial average:
        >>> k_mag = jnp.sqrt(kx**2 + ky**2)
        >>> k_binned, P_binned, counts = bin_power_spectrum(
        ...     k_mag.flatten(), P.flatten(), num_bins=50
        ... )
    """
    N = field.shape[0]

    # FFT of field
    field_k = jnp.fft.fft2(field)

    # Wavenumbers
    kx_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    ky_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    k_x, k_y = jnp.meshgrid(kx_1d, ky_1d, indexing="ij")

    # Power spectrum: |φ̃(k)|² / L²
    P_k = jnp.abs(field_k) ** 2 / L**2

    return k_x, k_y, P_k


def power_spectrum_3d(field: Array, L: float) -> Tuple[Array, Array, Array, Array]:
    """
    Compute 3D power spectrum P(k_x, k_y, k_z) from a field configuration.

    Returns the full 3D spectrum in k-space. For spherically averaged spectrum,
    use bin_power_spectrum() on the output.

    Args:
        field: 3D array of field values, shape (N, N, N)
        L: Physical length of the lattice (assumed cubic)

    Returns:
        Tuple of (k_x, k_y, k_z, P_k) where:
        - k_x, k_y, k_z: 3D arrays of wavenumber components
        - P_k: 3D power spectrum values

    Example:
        >>> field = jnp.array([...])  # 3D field
        >>> kx, ky, kz, P = power_spectrum_3d(field, L=20.0)
        >>> # Spherical average:
        >>> k_mag = jnp.sqrt(kx**2 + ky**2 + kz**2)
        >>> k_binned, P_binned, counts = bin_power_spectrum(
        ...     k_mag.flatten(), P.flatten(), num_bins=50
        ... )
    """
    N = field.shape[0]

    # FFT of field
    field_k = jnp.fft.fftn(field)

    # Wavenumbers
    kx_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    ky_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    kz_1d = jnp.fft.fftfreq(N, d=L / N) * 2 * jnp.pi
    k_x, k_y, k_z = jnp.meshgrid(kx_1d, ky_1d, kz_1d, indexing="ij")

    # Power spectrum: |φ̃(k)|² / L³
    P_k = jnp.abs(field_k) ** 2 / L**3

    return k_x, k_y, k_z, P_k


def power_spectrum(field: Array, L: float | Tuple[float, ...]) -> Tuple:
    """
    Compute power spectrum automatically based on field dimensionality.

    Dispatches to the appropriate 1D/2D/3D function based on field shape.

    Args:
        field: Field array (1D, 2D, or 3D)
        L: Physical length (scalar or tuple for anisotropic boxes)

    Returns:
        Tuple of (wavenumbers..., P_k) - format depends on dimensionality

    Example:
        >>> field = jnp.array([...])  # Any dimension
        >>> result = power_spectrum(field, L=20.0)
    """
    ndim = len(field.shape)

    # Convert L to scalar if needed
    if isinstance(L, (tuple, list)):
        if len(set(L)) > 1:
            raise ValueError("Anisotropic boxes not yet supported - use equal lengths")
        L_scalar = L[0]
    else:
        L_scalar = L

    if ndim == 1:
        return power_spectrum_1d(field, L_scalar)
    elif ndim == 2:
        return power_spectrum_2d(field, L_scalar)
    elif ndim == 3:
        return power_spectrum_3d(field, L_scalar)
    else:
        raise ValueError(f"Field must be 1D, 2D, or 3D, got {ndim}D")


def bin_power_spectrum(
    k_values: Array,
    P_values: Array,
    num_bins: int = 50,
    k_min: float | None = None,
    k_max: float | None = None,
    log_bins: bool = True,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Bin power spectrum values into radial/spherical shells.

    This function takes the raw k-space power spectrum and averages it into
    bins of |k| to produce a 1D radially-averaged spectrum P(k).

    Args:
        k_values: Flattened array of wavenumber magnitudes |k|
        P_values: Flattened array of power spectrum values
        num_bins: Number of bins for averaging
        k_min: Minimum k value (default: smallest non-zero k)
        k_max: Maximum k value (default: maximum k)
        log_bins: Use logarithmic binning (default: True, recommended for wide k-range)

    Returns:
        Tuple of (k_binned, P_binned, counts) where:
        - k_binned: Center wavenumber of each bin
        - P_binned: Average power in each bin
        - counts: Number of k-modes in each bin

    Example:
        >>> # 3D case
        >>> kx, ky, kz, P = power_spectrum_3d(field, L=20.0)
        >>> k_mag = jnp.sqrt(kx**2 + ky**2 + kz**2)
        >>> k_binned, P_binned, counts = bin_power_spectrum(
        ...     k_mag.flatten(), P.flatten(), num_bins=50
        ... )
        >>> plt.loglog(k_binned, P_binned)

    Notes:
        - Excludes k=0 mode (DC component)
        - Log binning is recommended for fields with power-law spectra
        - Empty bins are removed from output
    """
    # Convert to numpy for binning operations
    k_vals = np.array(k_values)
    P_vals = np.array(P_values)

    # Remove k=0 and any NaN/Inf values
    mask = (k_vals > 0) & np.isfinite(k_vals) & np.isfinite(P_vals)
    k_vals = k_vals[mask]
    P_vals = P_vals[mask]

    if len(k_vals) == 0:
        return np.array([]), np.array([]), np.array([])

    # Set bin edges
    if k_min is None:
        k_min = k_vals.min()
    if k_max is None:
        k_max = k_vals.max()

    if log_bins:
        # Logarithmic bins
        if k_min <= 0:
            k_min = k_vals[k_vals > 0].min()
        bin_edges = np.logspace(np.log10(k_min), np.log10(k_max), num_bins + 1)
    else:
        # Linear bins
        bin_edges = np.linspace(k_min, k_max, num_bins + 1)

    # Bin the data
    k_binned = []
    P_binned = []
    counts = []

    for i in range(num_bins):
        # Find k-modes in this bin
        in_bin = (k_vals >= bin_edges[i]) & (k_vals < bin_edges[i + 1])

        if np.sum(in_bin) > 0:
            # Average k and P in this bin
            k_binned.append(np.mean(k_vals[in_bin]))
            P_binned.append(np.mean(P_vals[in_bin]))
            counts.append(np.sum(in_bin))

    return np.array(k_binned), np.array(P_binned), np.array(counts)


def power_spectrum_with_binning(
    field: Array,
    L: float,
    num_bins: int = 50,
    log_bins: bool = True,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Convenience function: compute and bin power spectrum in one call.

    Automatically handles 1D/2D/3D fields and returns radially/spherically
    averaged power spectrum.

    Args:
        field: Field array (1D, 2D, or 3D)
        L: Physical length of the lattice
        num_bins: Number of bins for averaging
        log_bins: Use logarithmic binning

    Returns:
        Tuple of (k, P, counts) - binned power spectrum

    Example:
        >>> field = jnp.array([...])  # 3D field
        >>> k, P, counts = power_spectrum_with_binning(field, L=20.0, num_bins=50)
        >>> plt.loglog(k, P)
        >>> plt.xlabel('k')
        >>> plt.ylabel('P(k)')
    """
    ndim = len(field.shape)

    if ndim == 1:
        k_values, P_k = power_spectrum_1d(field, L)
        return bin_power_spectrum(k_values, P_k, num_bins, log_bins=log_bins)

    elif ndim == 2:
        k_x, k_y, P_k = power_spectrum_2d(field, L)
        k_mag = jnp.sqrt(k_x**2 + k_y**2)
        return bin_power_spectrum(
            k_mag.flatten(), P_k.flatten(), num_bins, log_bins=log_bins
        )

    elif ndim == 3:
        k_x, k_y, k_z, P_k = power_spectrum_3d(field, L)
        k_mag = jnp.sqrt(k_x**2 + k_y**2 + k_z**2)
        return bin_power_spectrum(
            k_mag.flatten(), P_k.flatten(), num_bins, log_bins=log_bins
        )

    else:
        raise ValueError(f"Field must be 1D, 2D, or 3D, got {ndim}D")
