"""
Simple scalar field evolution for basic Lattice class.

This module provides time evolution for scalar fields using the Klein-Gordon
equation. Supports both flat spacetime and cosmological evolution with
Hubble friction.

For coupled gauge-scalar systems with full cosmological expansion,
use the leapfrog module instead.
"""

import jax.numpy as jnp
from jax import Array, grad
from typing import Tuple, List, Callable
from rich.progress import (
    Progress,
    SpinnerColumn,
    TextColumn,
    BarColumn,
    TimeRemainingColumn,
)

from jaxlatt.core import Lattice
from jaxlatt.operators.stencil import laplacian as compute_laplacian


def compute_force(
    field: Array,
    potential: Callable[[Array], Array],
    dx: float,
    use_spectral: bool = False,
) -> Array:
    """
    Compute the force term F(φ) = ∇²φ - dV/dφ.

    Args:
        field: Scalar field configuration
        potential: Potential function V(φ)
        dx: Lattice spacing
        use_spectral: If True, use FFT-based spectral Laplacian (exact k² dispersion).
            If False, use finite-difference stencil (default, works with any BC).

    Returns:
        Force term at each site
    """
    # Laplacian term
    if use_spectral:
        from jaxlatt.operators.spectral import laplacian as spectral_laplacian
        lap = spectral_laplacian(field, dx)
    else:
        lap = compute_laplacian(field, dx)

    # Potential derivative using JAX autodiff
    # For array inputs, we compute gradient element-wise
    dV_dphi = grad(lambda phi: jnp.sum(potential(phi)))(field)

    # Total force
    force = lap - dV_dphi

    return force


def evolve_step(
    lattice: Lattice,
    potential: Callable[[Array], Array],
    dt: float,
    H: float = 0.0,
    use_spectral: bool = False,
) -> Lattice:
    """
    Perform a single leapfrog time step for a scalar field.

    Uses the leapfrog (Verlet) scheme with optional Hubble friction:
    1. φ̇_{n+1/2} = φ̇_n + (dt/2) * F(φ_n)
    2. φ_{n+1} = φ_n + dt * φ̇_{n+1/2}
    3. φ̇_{n+1} = φ̇_{n+1/2} + (dt/2) * F(φ_{n+1})

    where F(φ) = ∇²φ - dV/dφ

    For cosmological evolution with Hubble friction (H > 0), the equation becomes:
        φ̈ + 2H φ̇ = ∇²φ - dV/dφ

    The friction term is treated semi-implicitly for numerical stability:
        φ̇_{n+1/2} = [φ̇_n + (dt/2)*F_n] / (1 + H*dt)
        φ̇_{n+1} = [φ̇_{n+1/2} + (dt/2)*F_{n+1}] / (1 + H*dt)

    Args:
        lattice: Current lattice state
        potential: Potential function V(φ)
        dt: Time step
        H: Hubble parameter (conformal). If H=0 (default), no friction is applied.
        use_spectral: If True, use FFT-based spectral Laplacian for exact k² dispersion.
            Recommended for validation and periodic domains. Default: False (finite-difference).

    Returns:
        New lattice state at time t + dt

    Example:
        >>> # Flat spacetime (no friction)
        >>> lattice_new = evolve_step(lattice, potential, dt=0.01)
        >>>
        >>> # With Hubble friction
        >>> lattice_new = evolve_step(lattice, potential, dt=0.01, H=0.1)
        >>>
        >>> # With spectral Laplacian for validation
        >>> lattice_new = evolve_step(lattice, potential, dt=0.01, H=0.1, use_spectral=True)
    """
    # Current state
    field = lattice.field
    field_dot = lattice.field_dot

    # Compute force at current position
    force = compute_force(field, potential, lattice.dx, use_spectral=use_spectral)

    # Half-step velocity update
    field_dot_half = field_dot + 0.5 * dt * force

    # Apply Hubble friction (semi-implicit treatment)
    if H > 0.0:
        field_dot_half = field_dot_half / (1.0 + H * dt)

    # Full-step position update
    field_new = field + dt * field_dot_half

    # Compute force at new position
    force_new = compute_force(field_new, potential, lattice.dx, use_spectral=use_spectral)

    # Half-step velocity update
    field_dot_new = field_dot_half + 0.5 * dt * force_new

    # Apply Hubble friction (semi-implicit treatment)
    if H > 0.0:
        field_dot_new = field_dot_new / (1.0 + H * dt)

    # Return new lattice state
    return Lattice(
        size=lattice.size,
        length=lattice.length,
        field=field_new,
        field_dot=field_dot_new,
    )


def evolve(
    lattice: Lattice,
    potential: Callable[[Array], Array],
    dt: float,
    steps: int,
    snapshot_interval: int = 1,
    verbose: bool = False,
    H: float = 0.0,
    use_spectral: bool = False,
) -> Tuple[Array, List[Lattice]]:
    """
    Evolve a scalar field lattice for multiple time steps.

    Args:
        lattice: Initial lattice state
        potential: Potential function V(φ)
        dt: Time step
        steps: Number of time steps
        snapshot_interval: Save snapshot every N steps (default: 1)
        verbose: Print progress updates (default: False)
        H: Hubble parameter for cosmological friction (default: 0.0, no friction)
        use_spectral: Use FFT-based spectral Laplacian (default: False)

    Returns:
        Tuple of (times, snapshots) where:
        - times: Array of snapshot times
        - snapshots: List of Lattice objects at snapshot times

    Example:
        >>> from jaxlatt.core import ScalarPotential
        >>> from jaxlatt.utils import create_initial_lattice_1d
        >>>
        >>> # Flat spacetime evolution
        >>> lattice = create_initial_lattice_1d(size=128, length=10.0)
        >>> potential = ScalarPotential.quadratic(m=1.0)
        >>> times, snapshots = evolve(lattice, potential, dt=0.01, steps=1000)
        >>>
        >>> # Cosmological evolution with Hubble friction
        >>> times, snapshots = evolve(lattice, potential, dt=0.01, steps=1000, H=0.1)
    """
    snapshots = []
    times = []

    current_lattice = lattice

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

            for step in range(steps):
                if step % snapshot_interval == 0:
                    snapshots.append(current_lattice)
                    times.append(step * dt)

                current_lattice = evolve_step(
                    current_lattice, potential, dt, H=H, use_spectral=use_spectral
                )
                progress.update(task, advance=1)
    else:
        for step in range(steps):
            if step % snapshot_interval == 0:
                snapshots.append(current_lattice)
                times.append(step * dt)

            current_lattice = evolve_step(
                current_lattice, potential, dt, H=H, use_spectral=use_spectral
            )

    # Save final state
    if (steps - 1) % snapshot_interval != 0:
        snapshots.append(current_lattice)
        times.append(steps * dt)

    return jnp.array(times), snapshots
