"""
Symplectic leapfrog integrators for lattice field theory.

Organized by theory type:
- Pure scalar fields (flat space only)
- Pure gauge fields (flat + expanding)
- Coupled scalar-gauge fields (flat + expanding)

All integrators preserve:
- Energy conservation (to numerical precision)
- Symplectic structure
- Gauge constraints (when initially satisfied)

Memory optimization notes:
- donate_argnums allows JAX to reuse input buffers for outputs
- jax.lax.scan is used for evolution loops to avoid Python overhead
- Pre-allocated arrays are used instead of list append where possible
"""

import jax.numpy as jnp
from jax import jit, lax, Array
from functools import partial
from typing import List, Tuple, Callable, NamedTuple, Optional

from jaxlatt.core import (
    CoupledLattice,
    GaugeLattice,
    FRWUniverse,
    friedmann_step_predictor_corrector,
)
from jaxlatt.operators.coupled import scalar_force, covariant_laplacian
from jaxlatt.operators.gauge import gauge_force, evolve_links_half_step
from jaxlatt.core.cosmology import friedmann_acceleration


# =============================================================================
# Pure Gauge Evolution (Flat Space)
# =============================================================================


@partial(jit, donate_argnums=(0, 1))
def _gauge_leapfrog_step_arrays(
    links: jnp.ndarray,
    E: jnp.ndarray,
    dx: float,
    g: float,
    dt: float,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """
    JIT-compiled gauge leapfrog step on raw arrays.

    Note: Uses donate_argnums=(0,1) to allow in-place memory reuse.
    """
    # Half-step for electric field
    force = gauge_force(links, dx, g)
    E_half = E - 0.5 * dt * force

    # Full step for links: U -> U * exp(i dt E)
    links_new = links * jnp.exp(1j * dt * E_half)
    # Ensure unitarity (normalize to unit circle)
    links_new = links_new / jnp.abs(links_new)

    # Second half-step for electric field
    force_new = gauge_force(links_new, dx, g)
    E_new = E_half - 0.5 * dt * force_new

    return links_new, E_new


def gauge_leapfrog_step(lattice: "GaugeLattice", dt: float) -> "GaugeLattice":
    """
    Single leapfrog step for pure gauge U(1) evolution.

    Algorithm:
        1. E(t+dt/2) = E(t) - (dt/2) * F[U(t)]
        2. U(t+dt) = U(t) * exp(i * dt * E(t+dt/2))
        3. E(t+dt) = E(t+dt/2) - (dt/2) * F[U(t+dt)]

    Preserves unitarity, Gauss constraint, and energy.

    Args:
        lattice: Current gauge lattice state
        dt: Time step

    Returns:
        Updated GaugeLattice at t+dt
    """
    links_new, E_new = _gauge_leapfrog_step_arrays(
        lattice.links, lattice.E, lattice.dx, lattice.g, dt
    )
    return lattice.update(links=links_new, E=E_new)


def gauge_evolve(
    lattice: "GaugeLattice",
    dt: float,
    steps: int,
    save_every: int = 1,
    verbose: bool = False,
) -> Tuple[Array, List["GaugeLattice"]]:
    """
    Evolve gauge field using leapfrog integrator.

    Uses jax.lax.scan for the inner loop to avoid Python overhead and
    enable better memory reuse. Snapshots are collected at specified intervals.

    Args:
        lattice: Initial gauge lattice state
        dt: Time step
        steps: Number of steps
        save_every: Save snapshot every N steps
        verbose: Show progress bar

    Returns:
        times: Array of snapshot times
        snapshots: List of GaugeLattice states
    """
    # Calculate number of snapshots (including initial state)
    num_snapshots = steps // save_every + 1

    # Pre-allocate arrays for snapshots
    links_shape = lattice.links.shape
    E_shape = lattice.E.shape

    # Store initial state
    all_links = jnp.zeros((num_snapshots,) + links_shape, dtype=lattice.links.dtype)
    all_E = jnp.zeros((num_snapshots,) + E_shape, dtype=lattice.E.dtype)
    all_links = all_links.at[0].set(lattice.links)
    all_E = all_E.at[0].set(lattice.E)

    # Pre-compute times array
    snapshot_indices = jnp.arange(num_snapshots)
    times = snapshot_indices * save_every * dt

    # Define scan body for evolution between snapshots
    def evolve_segment(carry, _):
        """Evolve for save_every steps and return snapshot."""
        links, E = carry

        # Inner loop: evolve for save_every steps using lax.fori_loop
        def step_fn(_, state):
            links_i, E_i = state
            links_new, E_new = _gauge_leapfrog_step_arrays(
                links_i, E_i, lattice.dx, lattice.g, dt
            )
            return (links_new, E_new)

        links_final, E_final = lax.fori_loop(
            0, save_every, step_fn, (links, E)
        )

        return (links_final, E_final), (links_final, E_final)

    if verbose:
        # For verbose mode, use Python loop with progress bar
        from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn

        snapshots = [lattice]
        current_links = lattice.links
        current_E = lattice.E

        with Progress(
            SpinnerColumn(),
            TextColumn("[bold blue]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeRemainingColumn(),
        ) as progress:
            task = progress.add_task("[cyan]Evolving gauge lattice...", total=steps)

            for step in range(1, steps + 1):
                current_links, current_E = _gauge_leapfrog_step_arrays(
                    current_links, current_E, lattice.dx, lattice.g, dt
                )

                if step % save_every == 0:
                    snap = lattice.update(links=current_links, E=current_E)
                    snapshots.append(snap)

                progress.update(task, advance=1)

        return times, snapshots
    else:
        # Use lax.scan for efficient non-verbose evolution
        num_segments = steps // save_every

        if num_segments > 0:
            (final_links, final_E), (stacked_links, stacked_E) = lax.scan(
                evolve_segment,
                (lattice.links, lattice.E),
                None,
                length=num_segments
            )

            # Combine initial state with evolved snapshots
            all_links = all_links.at[1:num_segments + 1].set(stacked_links)
            all_E = all_E.at[1:num_segments + 1].set(stacked_E)

        # Convert arrays back to GaugeLattice objects
        snapshots = []
        for i in range(num_snapshots):
            snap = lattice.update(links=all_links[i], E=all_E[i])
            snapshots.append(snap)

        return times, snapshots


# =============================================================================
# Coupled Scalar-Gauge Evolution (Flat Space)
# =============================================================================


class CoupledState(NamedTuple):
    """JAX-compatible state for coupled evolution (arrays only)."""

    phi: jnp.ndarray
    pi: jnp.ndarray
    links: jnp.ndarray
    E: jnp.ndarray


@partial(jit, static_argnums=(2,))
def gauge_force_single_direction(
    phi: jnp.ndarray,
    links: jnp.ndarray,
    direction: int,
    g: float,
    dx: float,
) -> jnp.ndarray:
    """
    Compute force on electric field E_i including scalar current.

    F_i = (g dx) Im(U_i S_i*) + (scalar current contribution)

    Args:
        phi: Scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        direction: Direction (0=x, 1=y, 2=z) - static for JIT
        g: Gauge coupling
        dx: Lattice spacing

    Returns:
        F_i: Force on E_i (real-valued)
    """
    from jaxlatt.operators.gauge import staple

    # Get staple for this direction only
    gauge_staple = staple(links, direction)
    U_i = links[direction]

    # Pure gauge force: (g dx) Im(U_i S_i*)
    pure_gauge_force = g * dx * jnp.imag(U_i * jnp.conj(gauge_staple))

    # Scalar current contribution from |D_i φ|² term
    phi_forward = jnp.roll(phi, -1, axis=direction)
    current = jnp.imag(jnp.conj(phi) * U_i * phi_forward) / dx**2

    return pure_gauge_force + current


@jit
def gauge_force_all_directions(
    phi: jnp.ndarray,
    links: jnp.ndarray,
    g: float,
    dx: float,
) -> jnp.ndarray:
    """Compute gauge force for all three directions."""
    F0 = gauge_force_single_direction(phi, links, 0, g, dx)
    F1 = gauge_force_single_direction(phi, links, 1, g, dx)
    F2 = gauge_force_single_direction(phi, links, 2, g, dx)
    return jnp.stack([F0, F1, F2], axis=0)


@partial(jit, donate_argnums=(0, 1, 2, 3))
def coupled_leapfrog_step_arrays(
    phi: jnp.ndarray,
    pi: jnp.ndarray,
    links: jnp.ndarray,
    E: jnp.ndarray,
    m: float,
    lambda_: float,
    g: float,
    dx: float,
    dt: float,
) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """
    Single leapfrog step operating on raw arrays (JIT-compatible).

    Note: Uses donate_argnums=(0,1,2,3) to allow in-place memory reuse
    for phi, pi, links, E arrays. This improves memory efficiency.

    Returns:
        Tuple of (phi_new, pi_new, links_new, E_new)
    """
    # Half-step for scalar field
    phi_half = phi + 0.5 * dt * pi

    # Half-step for links
    links_half = evolve_links_half_step(links, E, dt)

    # Compute forces at half-step
    F_phi = scalar_force(phi_half, links_half, m, lambda_, dx)
    F_gauge = gauge_force_all_directions(phi_half, links_half, g, dx)

    # Full step for momenta
    E_new = E - dt * F_gauge
    pi_new = pi + dt * F_phi

    # Second half-step for fields
    phi_new = phi_half + 0.5 * dt * pi_new
    links_new = evolve_links_half_step(links_half, E_new, dt)

    return phi_new, pi_new, links_new, E_new


def coupled_leapfrog_step(lattice: "CoupledLattice", dt: float) -> "CoupledLattice":
    """
    Single leapfrog step for coupled scalar-gauge evolution.

    Equations of motion:
        dφ/dt = π
        dπ/dt = F_φ
        dA_i/dt = E_i
        dE_i/dt = F_A (with scalar current)

    Args:
        lattice: Current state
        dt: Time step

    Returns:
        CoupledLattice at t+dt
    """
    phi_new, pi_new, links_new, E_new = coupled_leapfrog_step_arrays(
        lattice.phi,
        lattice.pi,
        lattice.links,
        lattice.E,
        lattice.m,
        lattice.lambda_,
        lattice.g,
        lattice.dx,
        dt,
    )

    return CoupledLattice(
        phi=phi_new,
        pi=pi_new,
        links=links_new,
        E=E_new,
        m=lattice.m,
        lambda_=lattice.lambda_,
        g=lattice.g,
        dx=lattice.dx,
        size=lattice.size,
        length=lattice.length,
    )


def coupled_evolve(
    lattice: "CoupledLattice",
    dt: float,
    steps: int,
    save_every: int = 1,
    verbose: bool = False,
) -> Tuple[List[float], List["CoupledLattice"]]:
    """
    Evolve coupled system over multiple time steps.

    Uses jax.lax.scan for the inner loop to avoid Python overhead and
    enable better memory reuse. Snapshots are collected at specified intervals.

    Args:
        lattice: Initial state
        dt: Time step
        steps: Number of steps
        save_every: Save snapshot every N steps
        verbose: Show progress bar

    Returns:
        times: List of saved snapshot times
        snapshots: List of saved lattice states
    """
    # Calculate number of snapshots (including initial state at step 0)
    num_snapshots = steps // save_every + 1

    # Pre-allocate arrays for snapshots
    phi_shape = lattice.phi.shape
    pi_shape = lattice.pi.shape
    links_shape = lattice.links.shape
    E_shape = lattice.E.shape

    # Pre-compute times array
    snapshot_indices = jnp.arange(num_snapshots)
    times = (snapshot_indices * save_every * dt).tolist()

    # Define scan body for evolution between snapshots
    def evolve_segment(carry, _):
        """Evolve for save_every steps and return snapshot."""
        phi, pi, links, E = carry

        # Inner loop: evolve for save_every steps using lax.fori_loop
        def step_fn(_, state):
            phi_i, pi_i, links_i, E_i = state
            phi_new, pi_new, links_new, E_new = coupled_leapfrog_step_arrays(
                phi_i, pi_i, links_i, E_i,
                lattice.m, lattice.lambda_, lattice.g, lattice.dx, dt
            )
            return (phi_new, pi_new, links_new, E_new)

        phi_f, pi_f, links_f, E_f = lax.fori_loop(
            0, save_every, step_fn, (phi, pi, links, E)
        )

        return (phi_f, pi_f, links_f, E_f), (phi_f, pi_f, links_f, E_f)

    if verbose:
        # For verbose mode, use Python loop with progress bar
        from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeRemainingColumn

        snapshots = [lattice]
        current_phi = lattice.phi
        current_pi = lattice.pi
        current_links = lattice.links
        current_E = lattice.E

        with Progress(
            SpinnerColumn(),
            TextColumn("[bold blue]{task.description}"),
            BarColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
            TimeRemainingColumn(),
        ) as progress:
            task = progress.add_task(
                "[cyan]Evolving coupled lattice...", total=steps
            )

            for step in range(1, steps + 1):
                current_phi, current_pi, current_links, current_E = coupled_leapfrog_step_arrays(
                    current_phi, current_pi, current_links, current_E,
                    lattice.m, lattice.lambda_, lattice.g, lattice.dx, dt
                )

                if step % save_every == 0:
                    snap = CoupledLattice(
                        phi=current_phi, pi=current_pi,
                        links=current_links, E=current_E,
                        m=lattice.m, lambda_=lattice.lambda_, g=lattice.g,
                        dx=lattice.dx, size=lattice.size, length=lattice.length,
                    )
                    snapshots.append(snap)

                progress.update(task, advance=1)

        return times, snapshots
    else:
        # Use lax.scan for efficient non-verbose evolution
        num_segments = steps // save_every

        # Store initial state
        all_phi = jnp.zeros((num_snapshots,) + phi_shape, dtype=lattice.phi.dtype)
        all_pi = jnp.zeros((num_snapshots,) + pi_shape, dtype=lattice.pi.dtype)
        all_links = jnp.zeros((num_snapshots,) + links_shape, dtype=lattice.links.dtype)
        all_E = jnp.zeros((num_snapshots,) + E_shape, dtype=lattice.E.dtype)

        all_phi = all_phi.at[0].set(lattice.phi)
        all_pi = all_pi.at[0].set(lattice.pi)
        all_links = all_links.at[0].set(lattice.links)
        all_E = all_E.at[0].set(lattice.E)

        if num_segments > 0:
            _, (stacked_phi, stacked_pi, stacked_links, stacked_E) = lax.scan(
                evolve_segment,
                (lattice.phi, lattice.pi, lattice.links, lattice.E),
                None,
                length=num_segments
            )

            # Combine initial state with evolved snapshots
            all_phi = all_phi.at[1:num_segments + 1].set(stacked_phi)
            all_pi = all_pi.at[1:num_segments + 1].set(stacked_pi)
            all_links = all_links.at[1:num_segments + 1].set(stacked_links)
            all_E = all_E.at[1:num_segments + 1].set(stacked_E)

        # Convert arrays back to CoupledLattice objects
        snapshots = []
        for i in range(num_snapshots):
            snap = CoupledLattice(
                phi=all_phi[i], pi=all_pi[i],
                links=all_links[i], E=all_E[i],
                m=lattice.m, lambda_=lattice.lambda_, g=lattice.g,
                dx=lattice.dx, size=lattice.size, length=lattice.length,
            )
            snapshots.append(snap)

        return times, snapshots


# =============================================================================
# Expanding Universe Evolution
# =============================================================================


def scalar_force_expanding(
    phi: jnp.ndarray,
    links: jnp.ndarray,
    m: float,
    lambda_: float,
    dx: float,
    a: float,
) -> jnp.ndarray:
    """
    Compute force on scalar field in expanding universe.

    Force for π update (excluding friction -2Hπ applied separately):
        F_φ = (∇²φ)/a² - a² dV/dφ*

    Uses autodiff via scalar_potential_force for correct Wirtinger derivatives
    on complex fields.

    Args:
        phi: Scalar field (comoving)
        links: Gauge links
        m: Mass parameter
        lambda_: Self-coupling
        dx: Lattice spacing (comoving)
        a: Scale factor

    Returns:
        F_φ: Force on phi (excluding friction term)
    """
    from jaxlatt.operators.scalar import scalar_potential_force

    # Covariant Laplacian
    laplacian = covariant_laplacian(phi, links, dx)

    # Kinetic term: +∇²φ/a²
    kinetic_force = laplacian / (a**2)

    # Potential force via autodiff: scalar_potential_force returns -dV/dφ*
    # We want: -a² dV/dφ* = a² * scalar_potential_force
    pot_force = scalar_potential_force(phi, m, lambda_)
    potential_force = (a**2) * pot_force

    return kinetic_force + potential_force


@partial(jit, static_argnums=(2,))
def gauge_force_expanding(
    phi: jnp.ndarray,
    links: jnp.ndarray,
    direction: int,
    g: float,
    dx: float,
    a: float,
) -> jnp.ndarray:
    """
    Compute force on electric field in expanding universe.

    Force returned is the RHS (before a² division applied in update).

    Args:
        phi: Scalar field
        links: Gauge links
        direction: Direction (0, 1, 2)
        g: Gauge coupling
        dx: Lattice spacing (comoving)
        a: Scale factor

    Returns:
        F_E: Force on E_i (to be divided by a² in evolution equation)
    """
    from jaxlatt.operators.gauge import staple

    # Get staple for this direction
    gauge_staple = staple(links, direction)
    U_i = links[direction]

    # Pure gauge force: (g dx) Im(U_i S_i*)
    pure_gauge_force = g * dx * jnp.imag(U_i * jnp.conj(gauge_staple))

    # Scalar current contribution
    phi_forward = jnp.roll(phi, -1, axis=direction)
    current = jnp.imag(jnp.conj(phi) * U_i * phi_forward) / dx**2

    return pure_gauge_force + current


@jit
def gauge_force_all_directions_expanding(
    phi: jnp.ndarray,
    links: jnp.ndarray,
    g: float,
    dx: float,
    a: float,
) -> jnp.ndarray:
    """Compute gauge force for all three directions in expanding universe.

    Avoids Python loop by calling each direction explicitly.
    """
    F0 = gauge_force_expanding(phi, links, 0, g, dx, a)
    F1 = gauge_force_expanding(phi, links, 1, g, dx, a)
    F2 = gauge_force_expanding(phi, links, 2, g, dx, a)
    return jnp.stack([F0, F1, F2], axis=0)


def coupled_leapfrog_step_expanding(
    lattice: "CoupledLattice",
    universe: "FRWUniverse",
    rho_func: Callable[[float], float],
    dt: float,
) -> Tuple["CoupledLattice", "FRWUniverse"]:
    """
    Leapfrog step with cosmological expansion and Hubble friction.

    Modified equations in conformal time:
        φ' = π
        π' = F_φ/a² - 2H π       (Hubble friction on π)
        A'_i = E_i
        E'_i = -F_E/a² - 2H E_i  (Hubble friction on E)
        a' = a² H                 (scale factor evolution)

    Args:
        lattice: Current field state (comoving coordinates)
        universe: Current cosmological state
        rho_func: Function rho(a) giving physical energy density
        dt: Time step (conformal time)

    Returns:
        lattice_new: Updated field state
        universe_new: Updated cosmological state
    """
    a = universe.a
    H = universe.H

    # Half-step for fields (no friction yet)
    phi_half = lattice.phi + 0.5 * dt * lattice.pi

    # Update links half-step
    links_half = evolve_links_half_step(lattice.links, lattice.E, dt)

    # Compute forces at half-step
    F_phi = scalar_force_expanding(
        phi_half, links_half, lattice.m, lattice.lambda_, lattice.dx, a
    )

    # Gauge forces for all directions (vectorized, no Python loop)
    F_gauge = gauge_force_all_directions_expanding(
        phi_half, links_half, lattice.g, lattice.dx, a
    )

    # Full step for momenta WITH HUBBLE FRICTION
    pi_new = lattice.pi + dt * (F_phi - 2.0 * H * lattice.pi)
    E_new = lattice.E + dt * ((F_gauge / (a**2)) - 2.0 * H * lattice.E)

    # Second half-step for fields
    phi_new = phi_half + 0.5 * dt * pi_new
    links_new = evolve_links_half_step(links_half, E_new, dt)

    # Safety checks for overflow
    phi_max = jnp.max(jnp.abs(phi_new))
    pi_max = jnp.max(jnp.abs(pi_new))
    E_max = jnp.max(jnp.abs(E_new))

    if jnp.isnan(phi_max) or jnp.isinf(phi_max) or phi_max > 1e6:
        import warnings

        warnings.warn(f"Scalar field overflow: max|φ| = {phi_max:.2e} at a={a:.3f}")
    if jnp.isnan(pi_max) or jnp.isinf(pi_max) or pi_max > 1e6:
        import warnings

        warnings.warn(f"Momentum overflow: max|π| = {pi_max:.2e} at a={a:.3f}")
    if jnp.isnan(E_max) or jnp.isinf(E_max) or E_max > 1e6:
        import warnings

        warnings.warn(f"Electric field overflow: max|E| = {E_max:.2e} at a={a:.3f}")

    # Update lattice
    lattice_new = CoupledLattice(
        phi=phi_new,
        pi=pi_new,
        links=links_new,
        E=E_new,
        m=lattice.m,
        lambda_=lattice.lambda_,
        g=lattice.g,
        dx=lattice.dx,
        size=lattice.size,
        length=lattice.length,
    )

    # Update cosmology
    universe_new = friedmann_step_predictor_corrector(universe, rho_func, dt)

    return lattice_new, universe_new


def coupled_evolve_expanding(
    lattice: "CoupledLattice",
    universe: "FRWUniverse",
    rho_func: Callable[[float], float],
    dt: float,
    steps: int,
    save_every: int = 1,
) -> Tuple[list, list, list]:
    """
    Evolve coupled fields in expanding universe.

    Note: This function cannot use jax.lax.scan because coupled_leapfrog_step_expanding
    contains Python control flow (safety checks with warnings). For fully JIT-compiled
    evolution, use the flat-space coupled_evolve() function instead.

    Args:
        lattice: Initial field configuration
        universe: Initial cosmological state
        rho_func: Function rho(a) for scale factor evolution
        dt: Time step (conformal time)
        steps: Number of steps
        save_every: Save every N steps

    Returns:
        times: List of conformal times
        lattices: List of saved lattice states
        universes: List of saved universe states
    """
    # Calculate number of snapshots for pre-allocation
    num_snapshots = steps // save_every + 1

    # Pre-allocate arrays for field snapshots (more memory efficient than list append)
    phi_shape = lattice.phi.shape
    pi_shape = lattice.pi.shape
    links_shape = lattice.links.shape
    E_shape = lattice.E.shape

    all_phi = jnp.zeros((num_snapshots,) + phi_shape, dtype=lattice.phi.dtype)
    all_pi = jnp.zeros((num_snapshots,) + pi_shape, dtype=lattice.pi.dtype)
    all_links = jnp.zeros((num_snapshots,) + links_shape, dtype=lattice.links.dtype)
    all_E = jnp.zeros((num_snapshots,) + E_shape, dtype=lattice.E.dtype)

    # Pre-allocate for universe parameters
    all_times = jnp.zeros(num_snapshots)
    all_a = jnp.zeros(num_snapshots)
    all_adot = jnp.zeros(num_snapshots)
    all_tau = jnp.zeros(num_snapshots)

    # Current state - no copy needed, step function creates new objects
    current_lattice = lattice
    current_universe = universe
    snapshot_idx = 0

    for step in range(steps + 1):
        if step % save_every == 0:
            # Store snapshot in pre-allocated arrays
            all_phi = all_phi.at[snapshot_idx].set(current_lattice.phi)
            all_pi = all_pi.at[snapshot_idx].set(current_lattice.pi)
            all_links = all_links.at[snapshot_idx].set(current_lattice.links)
            all_E = all_E.at[snapshot_idx].set(current_lattice.E)
            all_times = all_times.at[snapshot_idx].set(current_universe.tau)
            all_a = all_a.at[snapshot_idx].set(current_universe.a)
            all_adot = all_adot.at[snapshot_idx].set(current_universe.adot)
            all_tau = all_tau.at[snapshot_idx].set(current_universe.tau)
            snapshot_idx += 1

        if step < steps:
            current_lattice, current_universe = coupled_leapfrog_step_expanding(
                current_lattice, current_universe, rho_func, dt
            )

    # Convert pre-allocated arrays back to objects for API compatibility
    times = all_times.tolist()
    lattices = []
    universes = []

    for i in range(num_snapshots):
        lat = CoupledLattice(
            phi=all_phi[i], pi=all_pi[i],
            links=all_links[i], E=all_E[i],
            m=lattice.m, lambda_=lattice.lambda_, g=lattice.g,
            dx=lattice.dx, size=lattice.size, length=lattice.length,
        )
        lattices.append(lat)

        univ = FRWUniverse(
            a=float(all_a[i]),
            adot=float(all_adot[i]),
            tau=float(all_tau[i]),
            M_pl=universe.M_pl,
        )
        universes.append(univ)

    return times, lattices, universes


# =============================================================================
# Helper Functions
# =============================================================================


def make_radiation_rho(rho_initial: float, a_initial: float):
    """Create radiation-dominated energy density function ρ ∝ a^{-4}."""

    def rho_func(a: float) -> float:
        return rho_initial * (a_initial / a) ** 4

    return rho_func


def compute_energy_densities_expanding(
    lattice: "CoupledLattice",
    a: float,
) -> dict:
    """
    Compute COMOVING energy densities in expanding universe.

    Phase A: Returns only comoving densities (energy per comoving volume).
    Physical densities with correct a-scaling deferred to Phase B.

    Args:
        lattice: Field state
        a: Scale factor (unused in Phase A, kept for API compatibility)

    Returns:
        dict with:
            "total": total comoving energy density
            "kinetic": scalar kinetic energy density (comoving)
            "gradient": scalar gradient energy density (comoving)
            "potential": scalar potential energy density (comoving)
            "electric": electric energy density (comoving)
            "magnetic": magnetic energy density (comoving)
    """
    from jaxlatt.operators.coupled import coupled_energy

    energy_dict = coupled_energy(
        lattice.phi,
        lattice.pi,
        lattice.links,
        lattice.E,
        lattice.dx,
        lattice.m,
        lattice.lambda_,
        lattice.g,
    )

    volume = (lattice.dx * lattice.size[0]) ** 3

    # Return COMOVING energy densities only (Phase A)
    return {
        "total": energy_dict["total"] / volume,
        "kinetic": energy_dict["scalar_kinetic"] / volume,
        "gradient": energy_dict["scalar_gradient"] / volume,
        "potential": energy_dict["scalar_potential"] / volume,
        "electric": energy_dict["electric"] / volume,
        "magnetic": energy_dict["magnetic"] / volume,
    }


# =============================================================================
# Rescaled Field Evolution (Experimental)
# =============================================================================


def scalar_force_rescaled(
    chi: jnp.ndarray,
    links: jnp.ndarray,
    m: float,
    lambda_: float,
    dx: float,
    a: float,
    addot: float,
) -> jnp.ndarray:
    """
    Force for rescaled field χ = a φ eliminating explicit Hubble friction.

    Derived equation for χ:
        χ'' = (1/a²) ∇²χ + a² m² χ + λ |χ|² χ - (a''/a) χ

    Args:
        chi: Rescaled field χ = a φ
        links: Gauge links
        m, lambda_: Scalar parameters
        dx: Lattice spacing (comoving)
        a: Scale factor
        addot: Conformal second derivative a'' (from Friedmann)

    Returns:
        F_χ array same shape as χ
    """
    # ∇²χ (covariant Laplacian on χ)
    laplacian_chi = covariant_laplacian(chi, links, dx)

    # (1/a²) ∇²χ term
    kinetic_term = laplacian_chi / (a**2)

    # Potential term: a² m² χ + λ |χ|² χ
    chi_sq = jnp.abs(chi) ** 2
    potential_term = (a**2) * (m**2) * chi + lambda_ * chi_sq * chi

    # Curvature term: -(a''/a) χ
    curvature_term = -(addot / a) * chi if a > 0 else 0.0 * chi

    return kinetic_term + potential_term + curvature_term


def coupled_leapfrog_step_rescaled(
    lattice: "CoupledLattice",
    universe: "FRWUniverse",
    rho_func: Callable[[float], float],
    dt: float,
) -> Tuple["CoupledLattice", "FRWUniverse"]:
    """
    Leapfrog step using rescaled field χ=aφ eliminating explicit friction.

    Algorithm:
        χ = a φ
        π_χ = χ'
        F_χ = (1/a²)∇²χ + a² m² χ + λ|χ|²χ - (a''/a)χ
        χ_{n+1/2} = χ_n + (dt/2) π_χ_n
        π_χ_{n+1} = π_χ_n + dt F_χ(χ_{n+1/2})
        χ_{n+1} = χ_{n+1/2} + (dt/2) π_χ_{n+1}
        Update a via Friedmann then back-transform:
            φ = χ / a
            π = (π_χ - a' φ) / a

    Gauge fields still evolved with friction (future: rescale).

    Args:
        lattice: Current field state
        universe: Current cosmological state
        rho_func: Function rho(a) for scale factor evolution
        dt: Time step

    Returns:
        lattice_new: Updated field state
        universe_new: Updated cosmological state
    """
    a = universe.a
    adot = universe.adot

    # Compute a'' for curvature term
    rho_current = rho_func(a)
    addot = float(friedmann_acceleration(rho_current, universe.M_pl, a))

    # Rescaled variables
    chi = a * lattice.phi
    pi_chi = a * lattice.pi + adot * lattice.phi  # χ' = a φ' + a' φ

    # Half-step for χ
    chi_half = chi + 0.5 * dt * pi_chi

    # Update gauge links half-step
    links_half = evolve_links_half_step(lattice.links, lattice.E, dt)

    # Force at half-step
    F_chi = scalar_force_rescaled(
        chi_half, links_half, lattice.m, lattice.lambda_, lattice.dx, a, addot
    )

    # Full update for π_χ
    pi_chi_new = pi_chi + dt * F_chi
    chi_new = chi_half + 0.5 * dt * pi_chi_new

    # Evolve gauge momenta with original friction form (vectorized, no Python loop)
    H = universe.H
    F_gauge = gauge_force_all_directions_expanding(
        lattice.phi, lattice.links, lattice.g, lattice.dx, a
    )
    E_new = lattice.E + dt * ((F_gauge / (a**2)) - 2.0 * H * lattice.E)

    # Second half-step links
    links_new = evolve_links_half_step(links_half, E_new, dt)

    # Update scale factor (predictor-corrector) after field step
    universe_new = friedmann_step_predictor_corrector(universe, rho_func, dt)
    a_new = universe_new.a
    adot_new = universe_new.adot

    # Back-transform to φ, π at new time
    phi_new = chi_new / a_new
    pi_new = (pi_chi_new - adot_new * phi_new) / a_new

    lattice_new = CoupledLattice(
        phi=phi_new,
        pi=pi_new,
        links=links_new,
        E=E_new,
        m=lattice.m,
        lambda_=lattice.lambda_,
        g=lattice.g,
        dx=lattice.dx,
        size=lattice.size,
        length=lattice.length,
    )

    return lattice_new, universe_new
