"""
4th-order Runge-Kutta (RK4) integrator for lattice field evolution.

This provides an alternative to the symplectic leapfrog integrators, offering
higher accuracy (4th order vs 2nd order) at the cost of more force evaluations.

Use cases:
- High-precision evolution when energy drift is acceptable
- Benchmarking and comparison with leapfrog methods
- Short-time dynamics where accuracy matters more than symplectic structure
"""

import jax.numpy as jnp

from jaxlatt.core import CoupledLattice, GaugeLattice
from jaxlatt.operators.coupled import scalar_force
from jaxlatt.operators.gauge import gauge_force


def coupled_rk4_step(
    lattice: CoupledLattice,
    dt: float,
) -> CoupledLattice:
    """
    Single RK4 step for coupled scalar-gauge field evolution (flat space).

    This is a 4th-order accurate integrator that requires 4 force evaluations
    per step, compared to 2 for leapfrog. It does not preserve the symplectic
    structure but offers higher accuracy for smooth dynamics.

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

    Returns:
        Updated CoupledLattice at t+dt

    Note:
        For gauge fields, link updates use exponential map to preserve unitarity.
        Unlike leapfrog, this does NOT guarantee energy conservation.
    """
    phi = lattice.phi
    pi = lattice.pi
    links = lattice.links
    E = lattice.E
    dx = lattice.dx
    m = lattice.m
    lambda_ = lattice.lambda_
    g = lattice.g

    # RK4 for the system:
    # dφ/dt = π
    # dπ/dt = F_scalar(φ, links)
    # dE/dt = -F_gauge(links)
    # dU/dt = i*E*U (link evolution)

    # k1
    k1_phi = pi
    k1_pi = scalar_force(phi, links, m, lambda_, dx, g)
    k1_E = -gauge_force(links, dx, g)

    # k2
    phi_2 = phi + 0.5 * dt * k1_phi
    pi_2 = pi + 0.5 * dt * k1_pi
    links_2 = links * jnp.exp(0.5 * dt * 1j * E)
    links_2 = links_2 / jnp.abs(links_2)  # Ensure unitarity
    E_2 = E + 0.5 * dt * k1_E

    k2_phi = pi_2
    k2_pi = scalar_force(phi_2, links_2, m, lambda_, dx, g)
    k2_E = -gauge_force(links_2, dx, g)

    # k3
    phi_3 = phi + 0.5 * dt * k2_phi
    pi_3 = pi + 0.5 * dt * k2_pi
    links_3 = links * jnp.exp(0.5 * dt * 1j * E_2)
    links_3 = links_3 / jnp.abs(links_3)
    E_3 = E + 0.5 * dt * k2_E

    k3_phi = pi_3
    k3_pi = scalar_force(phi_3, links_3, m, lambda_, dx, g)
    k3_E = -gauge_force(links_3, dx, g)

    # k4
    phi_4 = phi + dt * k3_phi
    pi_4 = pi + dt * k3_pi
    links_4 = links * jnp.exp(dt * 1j * E_3)
    links_4 = links_4 / jnp.abs(links_4)
    E_4 = E + dt * k3_E

    k4_phi = pi_4
    k4_pi = scalar_force(phi_4, links_4, m, lambda_, dx, g)
    k4_E = -gauge_force(links_4, dx, g)

    # Combine with RK4 weights
    phi_new = phi + (dt / 6.0) * (k1_phi + 2 * k2_phi + 2 * k3_phi + k4_phi)
    pi_new = pi + (dt / 6.0) * (k1_pi + 2 * k2_pi + 2 * k3_pi + k4_pi)
    E_new = E + (dt / 6.0) * (k1_E + 2 * k2_E + 2 * k3_E + k4_E)

    # For links, use exponential of average E
    E_avg = (E + 2 * E_2 + 2 * E_3 + E_4) / 6.0
    links_new = links * jnp.exp(1j * dt * E_avg)
    links_new = links_new / jnp.abs(links_new)

    return lattice.replace(phi=phi_new, pi=pi_new, links=links_new, E=E_new)


def gauge_rk4_step(
    lattice: GaugeLattice,
    dt: float,
) -> GaugeLattice:
    """
    Single RK4 step for pure gauge field evolution (flat space).

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

    Returns:
        Updated GaugeLattice at t+dt
    """
    links = lattice.links
    E = lattice.E
    dx = lattice.dx
    g = lattice.g

    # k1
    k1_E = -gauge_force(links, dx, g)

    # k2
    links_2 = links * jnp.exp(0.5 * dt * 1j * E)
    links_2 = links_2 / jnp.abs(links_2)
    E_2 = E + 0.5 * dt * k1_E

    k2_E = -gauge_force(links_2, dx, g)

    # k3
    links_3 = links * jnp.exp(0.5 * dt * 1j * E_2)
    links_3 = links_3 / jnp.abs(links_3)
    E_3 = E + 0.5 * dt * k2_E

    k3_E = -gauge_force(links_3, dx, g)

    # k4
    links_4 = links * jnp.exp(dt * 1j * E_3)
    links_4 = links_4 / jnp.abs(links_4)
    E_4 = E + dt * k3_E

    k4_E = -gauge_force(links_4, dx, g)

    # Combine
    E_new = E + (dt / 6.0) * (k1_E + 2 * k2_E + 2 * k3_E + k4_E)

    # For links, use exponential of average E
    E_avg = (E + 2 * E_2 + 2 * E_3 + E_4) / 6.0
    links_new = links * jnp.exp(1j * dt * E_avg)
    links_new = links_new / jnp.abs(links_new)

    return lattice.replace(links=links_new, E=E_new)
