"""Energy diagnostics for field dynamics.

This module provides functions to compute and track various energy components:
- Kinetic energy: (1/2) ⟨φ̇²⟩
- Gradient energy: (1/2) ⟨(∇φ)²⟩
- Potential energy: ⟨V(φ)⟩
- Total energy and conservation checking

Both volume-averaged and volume-integrated functions are provided.

For expanding universe simulations, provides physical energy densities
with proper scale factor scaling for Friedmann equation integration.

Inspired by CosmoLattice's energy measurement framework.
"""

import jax.numpy as jnp
from jax import jit
from functools import partial
from typing import Dict, Callable, Tuple, TYPE_CHECKING
from jaxlatt.core import Lattice
from jaxlatt.operators import gradient_1d, gradient_2d, gradient_3d

if TYPE_CHECKING:
    from jaxlatt.core import CoupledLattice


# =============================================================================
# JIT-compiled Internal Helpers
# =============================================================================


@jit
def _kinetic_energy_averaged_impl(field_dot: jnp.ndarray) -> jnp.ndarray:
    """JIT-compiled kinetic energy computation."""
    return 0.5 * jnp.mean(field_dot**2)


@partial(jit, static_argnums=(1,))
def _gradient_energy_averaged_impl(
    field: jnp.ndarray,
    ndim: int,
    dx_per_dim: Tuple[float, ...],
) -> jnp.ndarray:
    """JIT-compiled gradient energy computation using FFT."""
    field_k = jnp.fft.fftn(field)

    if ndim == 1:
        N = field.shape[0]
        dx = dx_per_dim[0]
        k = 2 * jnp.pi * jnp.fft.fftfreq(N, d=dx)
        grad_squared_k = k**2 * jnp.abs(field_k) ** 2
        total_points = N
    elif ndim == 2:
        Nx, Ny = field.shape
        dx_x, dx_y = dx_per_dim
        kx = 2 * jnp.pi * jnp.fft.fftfreq(Nx, d=dx_x)
        ky = 2 * jnp.pi * jnp.fft.fftfreq(Ny, d=dx_y)
        KX, KY = jnp.meshgrid(kx, ky, indexing="ij")
        k_squared = KX**2 + KY**2
        grad_squared_k = k_squared * jnp.abs(field_k) ** 2
        total_points = Nx * Ny
    else:  # ndim == 3
        Nx, Ny, Nz = field.shape
        dx_x, dx_y, dx_z = dx_per_dim
        kx = 2 * jnp.pi * jnp.fft.fftfreq(Nx, d=dx_x)
        ky = 2 * jnp.pi * jnp.fft.fftfreq(Ny, d=dx_y)
        kz = 2 * jnp.pi * jnp.fft.fftfreq(Nz, d=dx_z)
        KX, KY, KZ = jnp.meshgrid(kx, ky, kz, indexing="ij")
        k_squared = KX**2 + KY**2 + KZ**2
        grad_squared_k = k_squared * jnp.abs(field_k) ** 2
        total_points = Nx * Ny * Nz

    return 0.5 * jnp.mean(grad_squared_k) / total_points


@jit
def _potential_energy_averaged_impl(
    V_field: jnp.ndarray,
) -> jnp.ndarray:
    """JIT-compiled potential energy computation."""
    return jnp.mean(V_field)


# =============================================================================
# Public API - Volume-Averaged Energy Functions
# =============================================================================


def kinetic_energy_averaged(lattice: Lattice) -> float:
    """Compute volume-averaged kinetic energy density.

    Kinetic energy density: ρ_kin = (1/2) φ̇²

    Args:
        lattice: Lattice object with field_dot

    Returns:
        Volume-averaged kinetic energy: ⟨ρ_kin⟩ = (1/2) ⟨φ̇²⟩

    Example:
        >>> E_kin = kinetic_energy_averaged(lattice)
        >>> print(f"Kinetic energy: {E_kin:.6e}")
    """
    return float(_kinetic_energy_averaged_impl(lattice.field_dot))


def gradient_energy_averaged(lattice: Lattice) -> float:
    """Compute volume-averaged gradient energy density.

    Gradient energy density: ρ_grad = (1/2) (∇φ)²

    Uses FFT-based gradient computation for accuracy.

    Args:
        lattice: Lattice object with field

    Returns:
        Volume-averaged gradient energy: ⟨ρ_grad⟩ = (1/2) ⟨(∇φ)²⟩

    Example:
        >>> E_grad = gradient_energy_averaged(lattice)
        >>> print(f"Gradient energy: {E_grad:.6e}")
    """
    ndim = lattice.field.ndim
    if ndim not in (1, 2, 3):
        raise ValueError(f"Unsupported dimensionality: {ndim}")

    dx_per_dim = tuple(L / N for L, N in zip(lattice.length, lattice.size))
    return float(_gradient_energy_averaged_impl(lattice.field, ndim, dx_per_dim))


def potential_energy_averaged(lattice: Lattice, potential_func: Callable) -> float:
    """Compute volume-averaged potential energy density.

    Potential energy density: ρ_pot = V(φ)

    Args:
        lattice: Lattice object with field
        potential_func: Potential function V(φ) -> V
                       Should be a JIT-compiled function

    Returns:
        Volume-averaged potential energy: ⟨ρ_pot⟩ = ⟨V(φ)⟩

    Example:
        >>> from jaxlatt.potentials import quadratic_potential
        >>> V = quadratic_potential(m=1.0)
        >>> E_pot = potential_energy_averaged(lattice, V)
        >>> print(f"Potential energy: {E_pot:.6e}")
    """
    # Evaluate potential at each grid point (potential_func should be JIT-compiled)
    V_field = potential_func(lattice.field)
    return float(_potential_energy_averaged_impl(V_field))


def total_energy_averaged(lattice: Lattice, potential_func: Callable) -> float:
    """Compute total volume-averaged energy density.

    Total energy: E_tot = E_kin + E_grad + E_pot

    Args:
        lattice: Lattice object
        potential_func: Potential function V(φ, φ̇) -> V

    Returns:
        Total volume-averaged energy density

    Example:
        >>> from jaxlatt.potentials import quadratic_potential
        >>> V = quadratic_potential(m=1.0)
        >>> E_tot = total_energy_averaged(lattice, V)
        >>> print(f"Total energy: {E_tot:.6e}")
    """
    E_kin = kinetic_energy_averaged(lattice)
    E_grad = gradient_energy_averaged(lattice)
    E_pot = potential_energy_averaged(lattice, potential_func)

    return E_kin + E_grad + E_pot


def energy_components_averaged(
    lattice: Lattice, potential_func: Callable
) -> Dict[str, float]:
    """Compute all energy components.

    Returns a dictionary with all energy contributions for easy tracking.

    Args:
        lattice: Lattice object
        potential_func: Potential function V(φ, φ̇) -> V

    Returns:
        Dictionary with keys:
            - 'kinetic': Kinetic energy ⟨φ̇²/2⟩
            - 'gradient': Gradient energy ⟨(∇φ)²/2⟩
            - 'potential': Potential energy ⟨V(φ)⟩
            - 'total': Total energy

    Example:
        >>> from jaxlatt.potentials import quadratic_potential
        >>> V = quadratic_potential(m=1.0)
        >>> energies = energy_components_averaged(lattice, V)
        >>> for name, value in energies.items():
        ...     print(f"{name:12s}: {value:.6e}")
    """
    E_kin = kinetic_energy_averaged(lattice)
    E_grad = gradient_energy_averaged(lattice)
    E_pot = potential_energy_averaged(lattice, potential_func)
    E_tot = E_kin + E_grad + E_pot

    return {
        "kinetic": float(E_kin),
        "gradient": float(E_grad),
        "potential": float(E_pot),
        "total": float(E_tot),
    }


def kinetic_energy_integrated(lattice: Lattice) -> float:
    """
    Compute total kinetic energy: E_kin = (1/2) ∫ (∂φ/∂t)² dV
    """
    return kinetic_energy_averaged(lattice) * lattice.volume


def potential_energy_integrated(lattice: Lattice, potential: Callable) -> float:
    """
    Compute total potential energy: E_pot = ∫ V(φ) dV
    """
    return potential_energy_averaged(lattice, potential) * lattice.volume


def energy_components_integrated(
    lattice: Lattice, potential: Callable
) -> Dict[str, float]:
    """
    Compute all energy components separately (volume-integrated).
    """
    # Get volume-averaged components
    components = energy_components_averaged(lattice, potential)

    # Convert to integrated quantities
    volume = lattice.volume
    return {
        "kinetic": components["kinetic"] * volume,
        "gradient": components["gradient"] * volume,
        "potential": components["potential"] * volume,
        "total": components["total"] * volume,
    }


def total_energy_integrated(lattice: Lattice, potential: Callable) -> float:
    """
    Compute total energy of a field configuration (volume-integrated).
    """
    return total_energy_averaged(lattice, potential) * lattice.volume


# TODO: Uncomment when gradient functions are moved to operators module
# def energy_density(
#     lattice: Lattice, potential: Callable
# ) -> jnp.ndarray:
#     """
#     Compute local energy density at each lattice point.
#     """
#     # Kinetic contribution
#     rho_kin = 0.5 * lattice.field_dot**2
#
#     # Potential contribution
#     rho_pot = potential(lattice.field)
#
#     # Gradient contribution
#     if lattice.field.ndim == 1:
#         grad = gradient_1d(lattice.field, lattice.dx)
#         rho_grad = 0.5 * grad**2
#     elif lattice.field.ndim == 2:
#         grad_x, grad_y = gradient_2d(lattice.field, lattice.dx)
#         rho_grad = 0.5 * (grad_x**2 + grad_y**2)
#     elif lattice.field.ndim == 3:
#         gx, gy, gz = gradient_3d(lattice.field, lattice.dx)
#         rho_grad = 0.5 * (gx**2 + gy**2 + gz**2)
#     else:
#         raise ValueError("energy_density only supports 1D, 2D, 3D")
#
#     return rho_kin + rho_grad + rho_pot


# =============================================================================
# Physical Energy Densities for Expanding Universe (FLRW Integration)
# =============================================================================


@jit
def compute_physical_scalar_gradient_energy(
    field: jnp.ndarray,
    scale_factor: float,
    dx: float,
) -> jnp.ndarray:
    """
    Compute physical gradient energy density with scale factor scaling.

    In expanding universe (conformal time):
        ρ_grad^(phys) = (∇φ)²/(2a²)

    where ∇ is the comoving gradient operator.

    Args:
        field: Scalar field configuration
        scale_factor: FLRW scale factor a(τ)
        dx: Lattice spacing (comoving)

    Returns:
        Physical gradient energy density (averaged over lattice)

    Note:
        Uses FFT-based gradient computation for spectral accuracy.
        For expanding universe, this should be called with comoving coordinates.
    """
    ndim = field.ndim

    # FFT-based gradient computation
    field_k = jnp.fft.fftn(field)

    if ndim == 1:
        N = field.shape[0]
        k = 2 * jnp.pi * jnp.fft.fftfreq(N, d=dx)
        grad_squared_k = k**2 * jnp.abs(field_k) ** 2
        total_points = N
    elif ndim == 2:
        Nx, Ny = field.shape
        kx = 2 * jnp.pi * jnp.fft.fftfreq(Nx, d=dx)
        ky = 2 * jnp.pi * jnp.fft.fftfreq(Ny, d=dx)
        KX, KY = jnp.meshgrid(kx, ky, indexing="ij")
        k_squared = KX**2 + KY**2
        grad_squared_k = k_squared * jnp.abs(field_k) ** 2
        total_points = Nx * Ny
    elif ndim == 3:
        Nx, Ny, Nz = field.shape
        kx = 2 * jnp.pi * jnp.fft.fftfreq(Nx, d=dx)
        ky = 2 * jnp.pi * jnp.fft.fftfreq(Ny, d=dx)
        kz = 2 * jnp.pi * jnp.fft.fftfreq(Nz, d=dx)
        KX, KY, KZ = jnp.meshgrid(kx, ky, kz, indexing="ij")
        k_squared = KX**2 + KY**2 + KZ**2
        grad_squared_k = k_squared * jnp.abs(field_k) ** 2
        total_points = Nx * Ny * Nz
    else:
        raise ValueError(f"Unsupported dimensionality: {ndim}")

    # Gradient energy with a^-2 scaling
    a_squared = scale_factor * scale_factor
    grad_energy_comoving = 0.5 * jnp.mean(grad_squared_k) / total_points

    return grad_energy_comoving / a_squared


@jit
def compute_physical_scalar_kinetic_energy(
    field_dot: jnp.ndarray,
    scale_factor: float,
) -> jnp.ndarray:
    """
    Compute physical kinetic energy density with scale factor scaling.

    In conformal time with canonical momentum π = a³ φ̇:
        ρ_kin^(phys) = π²/(2a⁶) = φ̇²/2  (if field_dot = φ̇)

    For proper implementation, we assume field_dot stores the conformal
    time derivative (φ' in conformal time), so:
        ρ_kin^(phys) = (φ')²/(2a²)

    Args:
        field_dot: Time derivative of field (conformal time)
        scale_factor: FLRW scale factor a(τ)

    Returns:
        Physical kinetic energy density (averaged)

    Note:
        Convention depends on how field_dot is stored. Here we assume
        it's the conformal time derivative, giving a^-2 scaling.
    """
    a_squared = scale_factor * scale_factor
    kinetic_comoving = 0.5 * jnp.mean(field_dot**2)

    return kinetic_comoving / a_squared


def compute_physical_potential_energy(
    field: jnp.ndarray,
    potential_func: Callable[[jnp.ndarray], jnp.ndarray],
) -> jnp.ndarray:
    """
    Compute physical potential energy density.

    Potential energy has no scale factor dependence:
        ρ_pot^(phys) = V(φ)

    Args:
        field: Scalar field configuration
        potential_func: Potential function V(φ)

    Returns:
        Physical potential energy density (averaged)
    """
    V_field = potential_func(field)
    return jnp.mean(V_field)


@jit
def compute_physical_gauge_electric_energy(
    E_field: jnp.ndarray,
    scale_factor: float,
) -> jnp.ndarray:
    """
    Compute physical electric field energy density.

    Electric field energy in expanding universe:
        ρ_E^(phys) = E²/(2a⁴)

    where E is the comoving electric field.

    Args:
        E_field: Electric field components (3, N, N, N)
        scale_factor: FLRW scale factor a(τ)

    Returns:
        Physical electric energy density (averaged)
    """
    a_4 = scale_factor**4
    E_squared = jnp.sum(E_field**2, axis=0)  # Sum over field components
    electric_comoving = 0.5 * jnp.mean(E_squared)

    return electric_comoving / a_4


@jit
def compute_physical_gauge_magnetic_energy(
    links: jnp.ndarray,
    scale_factor: float,
    dx: float,
    g: float,
) -> jnp.ndarray:
    """
    Compute physical magnetic field energy density.

    Magnetic field energy in expanding universe:
        ρ_B^(phys) = B²/(2a⁴)

    where B is extracted from plaquettes.

    Args:
        links: Gauge link variables (3, N, N, N)
        scale_factor: FLRW scale factor a(τ)
        dx: Lattice spacing (comoving)
        g: Gauge coupling

    Returns:
        Physical magnetic energy density (averaged)

    Note:
        Imports magnetic_field from operators.gauge to avoid circular dependency.
    """
    from jaxlatt.operators.gauge import magnetic_field

    B = magnetic_field(links, dx, g)
    a_4 = scale_factor**4
    B_squared = jnp.sum(B**2, axis=0)
    magnetic_comoving = 0.5 * jnp.mean(B_squared)

    return magnetic_comoving / a_4


def compute_physical_energy_density(
    lattice: "CoupledLattice",
    scale_factor: float,
    dx: float,
    potential_func: Callable[[jnp.ndarray], jnp.ndarray],
) -> float:
    """
    Compute total physical energy density for Friedmann equation.

    This is the source term ρ_total in the Friedmann equation:
        H² = (8πG/3) ρ_total

    Includes all energy components with correct scale factor scaling:
        - Scalar kinetic: |π|²/(2a⁶) where π is canonical momentum
        - Scalar gradient: (∇φ)²/(2a²)
        - Scalar potential: V(φ)
        - Electric field: E²/(2a⁴)
        - Magnetic field: B²/(2a⁴)

    Args:
        lattice: CoupledLattice with scalar and gauge fields
        scale_factor: Current scale factor a(τ)
        dx: Lattice spacing (comoving)
        potential_func: Scalar potential V(φ)

    Returns:
        Total physical energy density ρ_phys (averaged over lattice)

    Example:
        >>> from jaxlatt.core.potentials import quadratic_potential
        >>> V = quadratic_potential(m=1.0)
        >>> rho = compute_physical_energy_density(lattice, a=1.5, dx=0.5, potential_func=V)
        >>> print(f"Physical energy density: {rho:.6e}")

    Note:
        For scalar-only simulations, pass lattice with dummy gauge fields
        or use scalar-specific functions above.
    """
    # Scalar field contributions
    # Convert canonical momentum π to field_dot: φ̇ = π/a³
    # For complex fields, use |π|²
    pi_squared = (
        jnp.abs(lattice.pi) ** 2 if jnp.iscomplexobj(lattice.pi) else lattice.pi**2
    )
    field_dot = jnp.sqrt(pi_squared) / scale_factor**3

    rho_kinetic = compute_physical_scalar_kinetic_energy(field_dot, scale_factor)

    rho_gradient = compute_physical_scalar_gradient_energy(
        lattice.phi.real if jnp.iscomplexobj(lattice.phi) else lattice.phi,
        scale_factor,
        dx,
    )

    rho_potential = compute_physical_potential_energy(
        lattice.phi.real if jnp.iscomplexobj(lattice.phi) else lattice.phi,
        potential_func,
    )

    # Gauge field contributions
    rho_electric = compute_physical_gauge_electric_energy(lattice.E, scale_factor)

    rho_magnetic = compute_physical_gauge_magnetic_energy(
        lattice.links, scale_factor, dx, lattice.g
    )

    # Total physical energy density
    rho_total = rho_kinetic + rho_gradient + rho_potential + rho_electric + rho_magnetic

    return float(rho_total.real if jnp.iscomplexobj(rho_total) else rho_total)


def compute_physical_energy_components(
    lattice: "CoupledLattice",
    scale_factor: float,
    dx: float,
    potential_func: Callable[[jnp.ndarray], jnp.ndarray],
) -> Dict[str, float]:
    """
    Compute all physical energy density components separately.

    Useful for diagnostics and tracking how energy is distributed
    between different field components during expansion.

    Args:
        lattice: CoupledLattice state
        scale_factor: Current scale factor a(τ)
        dx: Lattice spacing (comoving)
        potential_func: Scalar potential V(φ)

    Returns:
        Dictionary with keys:
            - 'kinetic': Scalar kinetic energy density
            - 'gradient': Scalar gradient energy density
            - 'potential': Scalar potential energy density
            - 'electric': Electric field energy density
            - 'magnetic': Magnetic field energy density
            - 'total': Total physical energy density

    Example:
        >>> components = compute_physical_energy_components(lattice, a=2.0, dx=0.5, V)
        >>> for name, value in components.items():
        ...     print(f"{name:12s}: {value:.6e}")
    """
    # Scalar contributions - handle complex pi
    pi_squared = (
        jnp.abs(lattice.pi) ** 2 if jnp.iscomplexobj(lattice.pi) else lattice.pi**2
    )
    field_dot = jnp.sqrt(pi_squared) / scale_factor**3

    rho_kinetic = compute_physical_scalar_kinetic_energy(field_dot, scale_factor)

    rho_gradient = compute_physical_scalar_gradient_energy(
        lattice.phi.real if jnp.iscomplexobj(lattice.phi) else lattice.phi,
        scale_factor,
        dx,
    )

    rho_potential = compute_physical_potential_energy(
        lattice.phi.real if jnp.iscomplexobj(lattice.phi) else lattice.phi,
        potential_func,
    )

    # Gauge contributions
    rho_electric = compute_physical_gauge_electric_energy(lattice.E, scale_factor)

    rho_magnetic = compute_physical_gauge_magnetic_energy(
        lattice.links, scale_factor, dx, lattice.g
    )

    rho_total = rho_kinetic + rho_gradient + rho_potential + rho_electric + rho_magnetic

    return {
        "kinetic": float(
            rho_kinetic.real if jnp.iscomplexobj(rho_kinetic) else rho_kinetic
        ),
        "gradient": float(
            rho_gradient.real if jnp.iscomplexobj(rho_gradient) else rho_gradient
        ),
        "potential": float(
            rho_potential.real if jnp.iscomplexobj(rho_potential) else rho_potential
        ),
        "electric": float(
            rho_electric.real if jnp.iscomplexobj(rho_electric) else rho_electric
        ),
        "magnetic": float(
            rho_magnetic.real if jnp.iscomplexobj(rho_magnetic) else rho_magnetic
        ),
        "total": float(rho_total.real if jnp.iscomplexobj(rho_total) else rho_total),
    }


def compute_comoving_energy_density(
    lattice: "CoupledLattice",
    dx: float,
    potential_func: Callable[[jnp.ndarray], jnp.ndarray],
) -> float:
    """
    Compute comoving energy density (no scale factor corrections).

    This is what standard (non-expanding) simulations compute.
    Useful for comparison and debugging.

    Args:
        lattice: CoupledLattice state
        dx: Lattice spacing
        potential_func: Scalar potential V(φ)

    Returns:
        Comoving energy density (what you'd get with a=1)
    """
    return compute_physical_energy_density(
        lattice, scale_factor=1.0, dx=dx, potential_func=potential_func
    )


class EnergyTracker:
    """Track energy evolution and conservation over time.

    This class stores energy measurements at each timestep and provides
    methods to analyze energy conservation.

    Attributes:
        times: List of measurement times
        energies: List of energy component dictionaries

    Example:
        >>> from jaxlatt.potentials import quadratic_potential
        >>> V = quadratic_potential(m=1.0)
        >>> tracker = EnergyTracker()
        >>>
        >>> # During evolution
        >>> for t, snapshot in zip(times, snapshots):
        ...     tracker.add_measurement(t, snapshot, V)
        >>>
        >>> # Check conservation
        >>> conservation = tracker.energy_conservation()
        >>> print(f"Energy drift: {conservation['relative_change']:.2e}")
    """

    def __init__(self):
        """Initialize empty energy tracker."""
        self.times = []
        self.energies = []

    def add_measurement(
        self, time: float, lattice: Lattice, potential_func: Callable
    ) -> None:
        """Add energy measurement at given time.

        Args:
            time: Current simulation time
            lattice: Current lattice state
            potential_func: Potential function
        """
        self.times.append(float(time))
        self.energies.append(energy_components_averaged(lattice, potential_func))

    def get_component(self, component: str) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """Get time series of a specific energy component.

        Args:
            component: Energy component name ('kinetic', 'gradient',
                      'potential', or 'total')

        Returns:
            Tuple of (times, energy_values) as numpy arrays

        Example:
            >>> times, E_kin = tracker.get_component('kinetic')
            >>> plt.plot(times, E_kin, label='Kinetic')
        """
        if not self.energies:
            raise ValueError("No measurements recorded yet")

        values = [e[component] for e in self.energies]
        return jnp.array(self.times), jnp.array(values)

    def energy_conservation(self) -> Dict[str, float]:
        """Analyze energy conservation.

        Computes various metrics to assess how well energy is conserved
        during the simulation.

        Returns:
            Dictionary with:
                - 'initial': Initial total energy
                - 'final': Final total energy
                - 'mean': Mean total energy
                - 'std': Standard deviation
                - 'relative_change': (E_final - E_initial) / E_initial
                - 'max_deviation': max|E - E_mean| / E_mean

        Example:
            >>> conservation = tracker.energy_conservation()
            >>> if abs(conservation['relative_change']) > 1e-6:
            ...     print("WARNING: Significant energy drift detected!")
        """
        if len(self.energies) < 2:
            raise ValueError("Need at least 2 measurements for conservation check")

        _, E_tot = self.get_component("total")

        E_initial = E_tot[0]
        E_final = E_tot[-1]
        E_mean = jnp.mean(E_tot)
        E_std = jnp.std(E_tot)

        relative_change = float((E_final - E_initial) / E_initial)
        max_deviation = float(jnp.max(jnp.abs(E_tot - E_mean)) / E_mean)

        return {
            "initial": float(E_initial),
            "final": float(E_final),
            "mean": float(E_mean),
            "std": float(E_std),
            "relative_change": relative_change,
            "max_deviation": max_deviation,
        }

    def summary(self) -> str:
        """Generate human-readable summary of energy tracking.

        Returns:
            Formatted string with energy statistics

        Example:
            >>> print(tracker.summary())
        """
        if not self.energies:
            return "No measurements recorded"

        conservation = self.energy_conservation()

        summary_lines = [
            "=" * 60,
            "ENERGY CONSERVATION SUMMARY",
            "=" * 60,
            f"Time range: {self.times[0]:.2f} to {self.times[-1]:.2f}",
            f"Measurements: {len(self.times)}",
            "",
            "Total Energy:",
            f"  Initial:  {conservation['initial']:.10e}",
            f"  Final:    {conservation['final']:.10e}",
            f"  Mean:     {conservation['mean']:.10e}",
            f"  Std Dev:  {conservation['std']:.10e}",
            "",
            "Conservation Metrics:",
            f"  Relative change: {conservation['relative_change']:+.6e}",
            f"  Max deviation:   {conservation['max_deviation']:+.6e}",
            "",
        ]

        # Add final energy breakdown
        final = self.energies[-1]
        summary_lines.extend(
            [
                "Final Energy Breakdown:",
                f"  Kinetic:   {final['kinetic']:.10e}  ({final['kinetic'] / final['total'] * 100:.2f}%)",
                f"  Gradient:  {final['gradient']:.10e}  ({final['gradient'] / final['total'] * 100:.2f}%)",
                f"  Potential: {final['potential']:.10e}  ({final['potential'] / final['total'] * 100:.2f}%)",
                "=" * 60,
            ]
        )

        return "\n".join(summary_lines)

    def to_dict(self) -> Dict:
        """Export all data as dictionary for saving.

        Returns:
            Dictionary with 'times' and 'energies' arrays

        Example:
            >>> import json
            >>> data = tracker.to_dict()
            >>> with open('energies.json', 'w') as f:
            ...     json.dump(data, f)
        """
        return {"times": self.times, "energies": self.energies}


# =============================================================================
# Volume-Integrated Energy Functions (Backward Compatibility)
# =============================================================================


def kinetic_energy(field_dot: jnp.ndarray, dV: float) -> float:
    """
    Compute total kinetic energy: E_kin = (1/2) ∫ (∂φ/∂t)² dV

    Args:
        field_dot: Time derivative of the field
        dV: Volume element (grid cell volume)

    Returns:
        Total kinetic energy (integrated, not averaged)
    """
    return float(0.5 * jnp.sum(field_dot**2) * dV)


def potential_energy(
    field: jnp.ndarray, potential: Callable[[jnp.ndarray], jnp.ndarray], dV: float
) -> float:
    """
    Compute total potential energy: E_pot = ∫ V(φ) dV

    Args:
        field: Field configuration
        potential: Potential function V(φ)
        dV: Volume element

    Returns:
        Total potential energy (integrated, not averaged)
    """
    V = potential(field)
    return float(jnp.sum(V) * dV)


def compute_energy_components(
    lattice: Lattice, potential: Callable[[jnp.ndarray], jnp.ndarray]
) -> Dict[str, float]:
    """
    Compute all energy components separately (volume-integrated).

    Args:
        lattice: Lattice object containing field and field_dot
        potential: Potential function V(φ)

    Returns:
        Dictionary with keys 'kinetic', 'gradient', 'potential', 'total'
    """
    return energy_components_integrated(lattice, potential)


def compute_energy(
    lattice: Lattice, potential: Callable[[jnp.ndarray], jnp.ndarray]
) -> float:
    """
    Compute total energy of a field configuration (volume-integrated).

    E_total = E_kinetic + E_gradient + E_potential
            = ∫ [(1/2)(∂φ/∂t)² + (1/2)(∇φ)² + V(φ)] dV

    Args:
        lattice: Lattice object containing field and field_dot
        potential: Potential function V(φ)

    Returns:
        Total energy (integrated)
    """
    return total_energy_integrated(lattice, potential)


def energy_density(
    field: jnp.ndarray,
    field_dot: jnp.ndarray,
    potential: Callable[[jnp.ndarray], jnp.ndarray],
    dx: float,
) -> jnp.ndarray:
    """
    Compute local energy density at each lattice point.

    Args:
        field: Field configuration
        field_dot: Time derivative of field
        potential: Potential function
        dx: Lattice spacing

    Returns:
        Energy density at each lattice point
    """
    # Kinetic contribution
    rho_kin = 0.5 * field_dot**2

    # Potential contribution
    rho_pot = potential(field)

    # Gradient contribution
    if field.ndim == 1:
        grad = gradient_1d(field, dx)
        rho_grad = 0.5 * grad**2
    elif field.ndim == 2:
        grad_x, grad_y = gradient_2d(field, dx)
        rho_grad = 0.5 * (grad_x**2 + grad_y**2)
    elif field.ndim == 3:
        gx, gy, gz = gradient_3d(field, dx)
        rho_grad = 0.5 * (gx**2 + gy**2 + gz**2)
    else:
        raise ValueError("energy_density only supports 1D, 2D, 3D")

    return rho_kin + rho_grad + rho_pot
