"""
Force operators for lattice field evolution using autodiff.

This module provides a unified interface for computing forces in lattice
field theory simulations. All potential-derived forces use JAX autodiff,
guaranteeing correctness and energy conservation in symplectic integrators.

Design principles:
1. Forces are computed as F = -∂H/∂q via autodiff
2. Spatial operators (Laplacians) use finite differences (appropriate for lattices)
3. Factory functions return JIT-compiled, cached force functions
4. Support for both flat spacetime and expanding universe (FRW)

Example:
    >>> from jaxlatt.core.potentials import ScalarPotential
    >>> from jaxlatt.operators.forces import make_coupled_scalar_force
    >>>
    >>> potential = ScalarPotential.quartic(m=1.0, lambda_=0.1)
    >>> force_fn = make_coupled_scalar_force(potential, dx=0.5)
    >>> F = force_fn(phi, links)
"""

from functools import lru_cache
from typing import Callable, Tuple, Union

import jax.numpy as jnp
from jax import Array, grad, jit

from jaxlatt.core.potentials import ScalarPotential, PotentialFunction


# =============================================================================
# Simple Scalar Field Forces (no gauge coupling)
# =============================================================================


def make_scalar_force(
    potential: Union[ScalarPotential, PotentialFunction],
    dx: float,
) -> Callable[[Array], Array]:
    """
    Create force function for free scalar field (no gauge coupling).

    Computes F(φ) = ∇²φ - dV/dφ where:
    - ∇² is the discrete Laplacian (finite differences)
    - dV/dφ is computed via autodiff

    Args:
        potential: ScalarPotential instance or legacy potential function
        dx: Lattice spacing

    Returns:
        JIT-compiled function: field -> force

    Example:
        >>> from jaxlatt.core.potentials import ScalarPotential
        >>> V = ScalarPotential.quadratic(m=1.0)
        >>> force_fn = make_scalar_force(V, dx=0.1)
        >>> F = force_fn(field)
    """
    from jaxlatt.evolution.scalar import compute_laplacian

    # Handle both ScalarPotential and legacy functions
    if isinstance(potential, ScalarPotential):
        potential_force = potential.force
    else:
        # Legacy: create force from function
        @jit
        def potential_force(field: Array) -> Array:
            return -grad(lambda f: jnp.sum(potential(f)).real, holomorphic=False)(field)

    @jit
    def force(field: Array) -> Array:
        lap = compute_laplacian(field, dx)
        return lap + potential_force(field)

    return force


# =============================================================================
# Coupled Scalar-Gauge Forces (flat spacetime)
# =============================================================================


def make_coupled_scalar_force(
    potential: Union[ScalarPotential, PotentialFunction],
    dx: float,
) -> Callable[[Array, Array], Array]:
    """
    Create force for scalar field coupled to gauge field (flat spacetime).

    Computes F_φ = D²φ - dV/dφ where:
    - D² is the covariant Laplacian (includes gauge links)
    - dV/dφ is computed via autodiff

    Args:
        potential: ScalarPotential instance or legacy potential function
        dx: Lattice spacing

    Returns:
        JIT-compiled function: (phi, links) -> force

    Example:
        >>> from jaxlatt.core.potentials import ScalarPotential
        >>> V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        >>> force_fn = make_coupled_scalar_force(V, dx=0.5)
        >>> F = force_fn(phi, links)
    """
    from jaxlatt.operators.coupled import covariant_laplacian

    # Handle both ScalarPotential and legacy functions
    if isinstance(potential, ScalarPotential):
        potential_force = potential.force
    else:

        @jit
        def potential_force(field: Array) -> Array:
            return -grad(lambda f: jnp.sum(potential(f)).real, holomorphic=False)(field)

    @jit
    def force(phi: Array, links: Array) -> Array:
        lap = covariant_laplacian(phi, links, dx)
        return lap + potential_force(phi)

    return force


# =============================================================================
# Expanding Universe Forces (FRW spacetime)
# =============================================================================


def make_coupled_scalar_force_expanding(
    potential: Union[ScalarPotential, PotentialFunction],
    dx: float,
) -> Callable[[Array, Array, float], Array]:
    """
    Create force for scalar field in expanding FRW universe.

    Computes F_φ = (1/a²)D²φ - a² dV/dφ* where:
    - D² is the covariant Laplacian
    - dV/dφ is computed via autodiff
    - a is the scale factor

    Note: The Hubble friction term -2Hπ is NOT included here;
    it should be applied separately in the evolution equation.

    Args:
        potential: ScalarPotential instance or legacy potential function
        dx: Comoving lattice spacing

    Returns:
        JIT-compiled function: (phi, links, a) -> force

    Example:
        >>> V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        >>> force_fn = make_coupled_scalar_force_expanding(V, dx=0.5)
        >>> F = force_fn(phi, links, a=2.0)
    """
    from jaxlatt.operators.coupled import covariant_laplacian

    # Handle both ScalarPotential and legacy functions
    if isinstance(potential, ScalarPotential):
        potential_force = potential.force
    else:

        @jit
        def potential_force(field: Array) -> Array:
            return -grad(lambda f: jnp.sum(potential(f)).real, holomorphic=False)(field)

    @jit
    def force(phi: Array, links: Array, a: float) -> Array:
        # Covariant Laplacian
        lap = covariant_laplacian(phi, links, dx)

        # Kinetic term: +D²φ/a²
        kinetic = lap / (a**2)

        # Potential term: -a² dV/dφ*
        # potential_force returns F = -dV/dφ
        # We want: -a² dV/dφ = a² * F
        pot = (a**2) * potential_force(phi)

        return kinetic + pot

    return force


def make_rescaled_scalar_force(
    potential: Union[ScalarPotential, PotentialFunction],
    dx: float,
) -> Callable[[Array, Array, float, float], Array]:
    """
    Create force for rescaled field χ = aφ (eliminates Hubble friction).

    The rescaled field evolution equation is:
        χ'' = (1/a²)∇²χ - a² dV_eff/dχ - (a''/a)χ

    where V_eff is the effective potential for χ.

    Args:
        potential: ScalarPotential for physical field φ (not χ)
        dx: Comoving lattice spacing

    Returns:
        JIT-compiled function: (chi, links, a, addot) -> force

    Example:
        >>> V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        >>> force_fn = make_rescaled_scalar_force(V, dx=0.5)
        >>> F = force_fn(chi, links, a=2.0, addot=0.1)
    """
    from jaxlatt.operators.coupled import covariant_laplacian

    # Get the potential function for computing derivatives
    if isinstance(potential, ScalarPotential):
        V_fn = lambda field: potential(field)
    else:
        V_fn = potential

    @jit
    def force(chi: Array, links: Array, a: float, addot: float) -> Array:
        # Covariant Laplacian of χ
        lap_chi = covariant_laplacian(chi, links, dx)

        # Kinetic term: (1/a²)∇²χ
        kinetic = lap_chi / (a**2)

        # Potential term via autodiff
        # For χ = aφ, we need -dV/dχ where V is expressed in terms of χ
        # V(φ) = V(χ/a), and the EOM has a⁴V(χ/a) structure
        # The force is: -a² dV/d(χ/a) * (1/a) = -a dV/dφ|_{φ=χ/a}
        def V_effective(chi_field: Array) -> Array:
            phi = chi_field / a
            return jnp.sum(V_fn(phi)).real * (a**4)

        dV_dchi = grad(V_effective, holomorphic=False)(chi)

        # Curvature term: -(a''/a)χ
        curvature = -(addot / a) * chi

        return kinetic - dV_dchi + curvature

    return force


# =============================================================================
# Gauge Field Forces
# =============================================================================


def make_gauge_force(
    g: float,
    dx: float,
) -> Callable[[Array, Array], Array]:
    """
    Create force on electric field from pure gauge + scalar coupling.

    Computes F_E for all three directions. This is geometric and
    doesn't particularly benefit from autodiff, but is provided here
    for API consistency.

    Args:
        g: Gauge coupling constant
        dx: Lattice spacing

    Returns:
        JIT-compiled function: (phi, links) -> force array (3, N, N, N)

    Example:
        >>> force_fn = make_gauge_force(g=0.5, dx=0.1)
        >>> F_E = force_fn(phi, links)
    """
    from jaxlatt.evolution.leapfrog import gauge_force_all_directions

    @jit
    def force(phi: Array, links: Array) -> Array:
        return gauge_force_all_directions(phi, links, g, dx)

    return force


def make_gauge_force_expanding(
    g: float,
    dx: float,
) -> Callable[[Array, Array, float], Array]:
    """
    Create force on electric field in expanding universe.

    Args:
        g: Gauge coupling constant
        dx: Comoving lattice spacing

    Returns:
        JIT-compiled function: (phi, links, a) -> force array (3, N, N, N)
    """
    from jaxlatt.evolution.leapfrog import gauge_force_expanding

    def force(phi: Array, links: Array, a: float) -> Array:
        # Compute force for each direction
        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)

    return jit(force)


# =============================================================================
# Full Hamiltonian-Derived Forces (for validation)
# =============================================================================


def make_forces_from_hamiltonian(
    potential: Union[ScalarPotential, PotentialFunction],
    g: float,
    dx: float,
) -> Tuple[Callable, Callable]:
    """
    Create scalar and gauge forces from full Hamiltonian via autodiff.

    This guarantees F = -∂H/∂q exactly, useful for:
    1. Validating other force implementations
    2. Energy conservation checks
    3. Novel potentials without analytical derivatives

    Note: Slightly slower than specialized implementations due to
    full Hamiltonian evaluation. Use make_coupled_scalar_force for production.

    Args:
        potential: ScalarPotential or legacy potential function
        g: Gauge coupling constant
        dx: Lattice spacing

    Returns:
        Tuple of (scalar_force_fn, gauge_force_fn)

    Example:
        >>> V = ScalarPotential.quartic(m=1.0, lambda_=0.1)
        >>> scalar_F, gauge_F = make_forces_from_hamiltonian(V, g=0.5, dx=0.1)
        >>> F_phi = scalar_F(phi, pi, links, E)
    """
    from jaxlatt.operators.coupled import covariant_gradient_squared
    from jaxlatt.operators.gauge import magnetic_energy_density, electric_energy_density

    # Get potential function
    if isinstance(potential, ScalarPotential):
        V_fn = lambda field: potential(field)
    else:
        V_fn = potential

    def total_hamiltonian(phi: Array, pi: Array, links: Array, E: Array) -> Array:
        """Compute total Hamiltonian H = T + V."""
        # Scalar kinetic energy
        T_scalar = 0.5 * jnp.sum(jnp.abs(pi) ** 2)

        # Scalar gradient energy (covariant)
        grad_sq = covariant_gradient_squared(phi, links, dx)
        T_grad = 0.5 * jnp.sum(grad_sq)

        # Scalar potential energy
        V_scalar = jnp.sum(V_fn(phi)).real

        # Electric energy
        T_E = jnp.sum(electric_energy_density(E))

        # Magnetic energy
        T_B = jnp.sum(magnetic_energy_density(links, dx, g))

        return T_scalar + T_grad + V_scalar + T_E + T_B

    @jit
    def scalar_force(phi: Array, pi: Array, links: Array, E: Array) -> Array:
        """F_φ = -∂H/∂φ*"""
        return -grad(lambda p: total_hamiltonian(p, pi, links, E).real, holomorphic=False)(
            phi
        )

    @jit
    def gauge_force(phi: Array, pi: Array, links: Array, E: Array) -> Array:
        """F_U = -∂H/∂U* (approximation for U(1))"""
        # For U(1) gauge theory, this is more subtle
        # We compute gradient w.r.t. link phases
        return -grad(
            lambda L: total_hamiltonian(phi, pi, L, E).real, holomorphic=False
        )(links)

    return scalar_force, gauge_force


# =============================================================================
# Validation Utilities
# =============================================================================


def validate_force_against_manual(
    autodiff_force: Array,
    manual_force: Array,
    rtol: float = 1e-5,
    atol: float = 1e-8,
) -> dict:
    """
    Compare autodiff force with manual implementation.

    Args:
        autodiff_force: Force computed via autodiff
        manual_force: Force computed manually
        rtol: Relative tolerance
        atol: Absolute tolerance

    Returns:
        Dictionary with comparison metrics
    """
    diff = autodiff_force - manual_force
    max_abs_error = float(jnp.max(jnp.abs(diff)))
    rms_error = float(jnp.sqrt(jnp.mean(jnp.abs(diff) ** 2)))
    max_rel_error = float(
        jnp.max(jnp.abs(diff) / (jnp.abs(manual_force) + atol))
    )

    matches = jnp.allclose(autodiff_force, manual_force, rtol=rtol, atol=atol)

    return {
        "matches": bool(matches),
        "max_abs_error": max_abs_error,
        "rms_error": rms_error,
        "max_rel_error": max_rel_error,
    }
