r"""
Vlasov's collisionless kinetic equation: phase mixing and reversibility
==========================================================================

Vlasov (1938) argued that in a plasma, particles feel the smooth field
of all the others far more than occasional close collisions, and dropped
the collision term:

.. math::

    \partial_t f + v\,\partial_x f + \frac{q}{m}E\,\partial_v f = 0,
    \qquad \partial_x E = \rho .

Without collisions nothing is dissipated. The distribution function is
carried unchanged along particle orbits, so phase-space volume, and any
functional such as :math:`\int f\ln f`, is conserved exactly. A wave can
still *appear* to decay, because the perturbation is sheared into ever
finer filaments in velocity. But the information stays in those
filaments, and reversing every velocity brings the wave back.

This example solves the Vlasov-Poisson system by its characteristics,
one particle per characteristic, with
:func:`~physicskit.plasma.kinetic.pic_step`. It shows the filamentation
directly in phase space, then runs the system backward in time and
watches the decayed wave reappear.
"""

# %%
import matplotlib.pyplot as plt
import numpy as np

from physicskit.plasma.kinetic import deposit_number_density, interpolate_field, landau_damping_ic, pic_step, solve_poisson_1d

k, v_th, alpha = 0.5, 1.0, 0.2
L, ng, dt = 2 * np.pi / k, 64, 0.1
x, v = landau_damping_ic(200_000, L, k, alpha, v_th, seed=1938)


def field_at(x):
    return interpolate_field(x, solve_poisson_1d(1.0 - deposit_number_density(x, L, ng), L), L)


# Stagger the velocities half a step back, as the leapfrog scheme requires.
v = v + 0.5 * field_at(x) * dt  # qm = -1

# %%
# Forward: the wave decays by phase mixing
# --------------------------------------------
# Snapshots of the perturbation :math:`\delta f = f - f_0` in phase space.
# Particles at different speeds drift apart, and the single density ripple
# is sheared into stripes that grow ever narrower in :math:`v`. By
# :math:`t=30` they are finer than the plotting bins: coarse-grained, the
# perturbation looks gone, though every particle still remembers it.
n_steps = 300
snap_steps = (0, 100, 300)
snapshots, energy_fwd = {}, []
x0_copy = x.copy()
for s in range(n_steps + 1):
    if s in snap_steps:
        snapshots[s] = (x.copy(), v.copy())
    if s < n_steps:
        x, v, fe = pic_step(x, v, L, ng, dt)
        energy_fwd.append(fe)

xb, vb = np.linspace(0, L, 49), np.linspace(-3.5, 3.5, 71)
H0 = np.histogram2d(x0_copy, snapshots[0][1], bins=(xb, vb))[0].mean(axis=0, keepdims=True)
fig1, axes = plt.subplots(1, 3, figsize=(14, 3.8), sharey=True)
for ax, s in zip(axes, snap_steps):
    H = np.histogram2d(*snapshots[s], bins=(xb, vb))[0]
    dH = (H - H0) / H0.max()
    ax.pcolormesh(xb, vb, dH.T, cmap="RdBu_r", vmin=-0.2, vmax=0.2, shading="auto")
    ax.set_title(f"t = {s * dt:.0f} / $\\omega_p$")
    ax.set_xlabel("x")
axes[0].set_ylabel("v")
fig1.suptitle(r"$\delta f(x, v)$: the perturbation filaments in velocity")
fig1.tight_layout()

# %%
# Backward: reverse every velocity
# ------------------------------------
# The leapfrog step is exactly time-reversible. Kick once more to get
# the velocity half a step *ahead*, flip its sign, and step forward
# again: each particle retraces its orbit, and the filaments fold back
# into the original ripple. The field energy climbs back to where it
# started.
v_ahead = v - field_at(x) * dt
x_b, v_b = x.copy(), -v_ahead
energy_back = []
for _ in range(n_steps):
    x_b, v_b, fe = pic_step(x_b, v_b, L, ng, dt)
    energy_back.append(fe)

err = np.max(np.abs((x_b - x0_copy + L / 2) % L - L / 2))
print(f"after {n_steps} steps forward and {n_steps} back, max particle displacement from start: {err:.1e}")
print(f"field energy: start {energy_fwd[0]:.2e}, after decay {energy_fwd[-1]:.2e}, after reversal {energy_back[-1]:.2e}")

t = np.arange(n_steps) * dt
fig2, ax = plt.subplots(figsize=(7, 3.8))
ax.semilogy(t, energy_fwd, color="steelblue", label="forward in time")
ax.semilogy(t[-1] + t, energy_back, color="firebrick", label="velocities reversed")
ax.axvline(t[-1], color="0.5", ls=":")
ax.set_xlabel(r"$\omega_p t$")
ax.set_ylabel("electrostatic field energy")
ax.set_title("Collisionless 'damping' undone by time reversal")
ax.legend(fontsize=8)
fig2.tight_layout()

plt.show()

# %%
# Check
# -----
# Collisionless damping is reversible: reversing every velocity retraces the
# particle orbits to round-off and the field energy, which had fallen over
# 1000-fold, comes back.
assert err < 1e-9
assert energy_fwd[-1] < 1e-3 * energy_fwd[0] and energy_back[-1] > 0.9 * energy_fwd[0]