r"""
CH4 plasma channel: field pulse on a pre-ionized channel, then relaxation (n_e video)
=====================================================================================

Transient 1D radial run of :class:`~rizer.cantera_ext.plasma_channel.PlasmaChannel`
on the **CH4 Goutier2025** mechanism: an initially ionized methane channel
(`n_e = 1e13 cm^-3` on axis, Gaussian) is driven by a **trapezoidal electric-field
pulse** (ramp up / hold / ramp down, mirroring rizer's 0D NRP discharge), then left
to relax. Renders a video of the electron number density `n_e(r, t)`.

This is the CH4 counterpart of ``plot_channel_field_ionization.py`` (which uses air).
It exercises the same pipeline: an explicit non-uniform grid, an arbitrary Python
``E(t)`` schedule, `sigma = n_e e^2/(m_e nu_m)`, the stiff **BDF** backend
(``integrator="bdf"``), and the recorded `n_e(r, t)`.

Physics note -- why CH4 *relaxes* rather than ionizes up
--------------------------------------------------------
The Goutier electron-impact rates are numerically fine (finite to >=3 eV), and CH4
runs cleanly here. But methane net-ionizes only above ~1 eV; below that it
net-recombines. At accessible fields the electron energy is consumed by CH4's large
electron-impact **dissociation** cross-sections, which clamp the electron temperature
around ~0.3 eV (~3500 K) -- well below the ionization threshold -- until the methane
is depleted. So the pre-ionized channel **relaxes** (recombination + dissociation +
diffusion) under the pulse; it does not avalanche. (Air, by contrast, ionizes
strongly at a few thousand K -- see ``plot_channel_field_ionization.py``.) The field
is kept modest and ramped: a strong instantaneous field would drive the stiff
electron-energy equation faster than the integrator can follow.

.. tags:: plasma, Cantera, CH4, two-temperature, transient, channel, relaxation, 1D, CVODE
"""  # noqa: D205

# %%
# Import the required libraries.
# -------------------------------

import cantera as ct
import numpy as np

import rizer.kin.extensible_rate  # noqa: F401 (Register CH4 custom rates)
from rizer.cantera_ext import PlasmaChannel
from rizer.misc.plt_utils import set_mpl_style
from rizer.misc.utils import get_path_to_data

set_mpl_style()

# %%
# Scenario parameters.
# ---------------------

mechanism = str(get_path_to_data("mechanisms") / "Goutier2025" / "CH4_to_C2H2.yaml")
phase = "plasma"

# -- Geometry parameters ------------------------------------------------ #
wall_radius = 4.0e-3  # Wall radius [m]
channel_width = 500.0e-6  # Channel width (FWHM) [m]
n_points = 21  # Number of radial grid points [-]

# -- Initial conditions -------------------------------------------------- #
wall_temperature = 1000.0  # Wall/ambient temperature [K]
core_temperature = 1500.0  # Initial gas/electron temp on axis [K] (gentle: Te=Tg)
peak_electron_density = 1.0e19  # Initial electron density on axis [m^-3] = 1e13 cm^-3
ion_species = "CH4+"  # Quasi-neutral partner for the initial electrons

sigma_gaussian = channel_width / (2.0 * np.sqrt(2.0 * np.log(2.0)))  # Gaussian sigma

# -- Electric field pulse (trapezoidal) ---------------------------------- #
peak_field = 5.0e5  # Peak applied field [V/m] (modest; CH4 dissociation-clamps Te)
t_rise = 30.0e-9  # Field rise time [s]
t_hold = 50.0e-9  # Field hold time [s]
t_fall = 30.0e-9  # Field fall time [s]
t_end = 150.0e-9  # Total simulated time [s]

# -- Collision frequency model --------------------------------------------#
nu_m = 1.0e12  # e-neutral momentum-transfer frequency [1/s]

# -- Solver parameters ----------------------------------------------------#
dt = 0.5e-9  # Output cadence / BDF max step [s]


# %%
# Define helper functions.
# --------------------------
#
# ``stretched_grid`` builds the non-uniform radial grid, ``build_initial_state``
# constructs the Gaussian pre-ionized channel, ``field_pulse`` returns the
# trapezoidal ``E(t)`` schedule, and ``make_video`` renders the ``n_e(r, t)``
# animation.


def stretched_grid(R, n, beta=3.5):
    """Non-uniform radial grid, refined near the axis via a sinh stretch."""
    x = np.linspace(0.0, 1.0, n)
    return R * np.sinh(beta * x) / np.sinh(beta)


def build_initial_state(r):
    """Gaussian warm, weakly-ionized CH4 channel; Te=Tg (field heats electrons)."""
    g = np.exp(-(r**2) / (2.0 * sigma_gaussian**2))
    Tg = wall_temperature + (core_temperature - wall_temperature) * g
    ne = peak_electron_density * g + 1.0e12
    gas = ct.Solution(mechanism, phase)
    P = ct.one_atm
    gas.TPX = wall_temperature, P, "CH4:1"  # ambient/wall = cold methane
    Y0 = gas.Y.copy()
    gas.TPX = 0.5 * (core_temperature + wall_temperature), P, "CH4:1"
    rho = float(gas.density)
    Y_2d = np.zeros((len(r), gas.n_species))
    for i in range(len(r)):
        x_e = ne[i] * ct.boltzmann * Tg[i] / P  # consistent X_e for target n_e
        gas.TPX = (
            Tg[i],
            P,
            f"CH4:{1 - 2 * x_e:.6e}, e-:{x_e:.6e}, {ion_species}:{x_e:.6e}",
        )
        Y_2d[i] = gas.Y
    return rho, Y0, Tg, Y_2d


def field_pulse():
    """Trapezoidal E(t): 0 ->(rise) peak_field ->(hold) peak_field ->(fall) 0 -> 0."""
    t1, t2, t3 = t_rise, t_rise + t_hold, t_rise + t_hold + t_fall
    return (
        np.array([0.0, t1, t2, t3, t_end]),
        np.array([0.0, peak_field, peak_field, 0.0, 0.0]),
    )


def make_video(ch, path="ch4_channel_ne.mp4"):
    """Render an n_e(r) animation over time, falling back across writers."""
    import matplotlib.pyplot as plt
    from matplotlib.animation import FFMpegWriter, FuncAnimation, PillowWriter

    r_mm = ch.r * 1e3
    ne = ch.electron_density() * 1e-6  # cm^-3
    t = ch.t
    t_off = t_rise + t_hold + t_fall
    fig, ax = plt.subplots(figsize=(7, 4.5))
    (line,) = ax.semilogy([], [], "o-", color="teal", ms=3)
    ax.set_xlim(0, r_mm.max())
    ax.set_ylim(ne.min() * 0.5, ne.max() * 2)
    ax.set_xlabel("Radius [mm]")
    ax.set_ylabel(r"$n_e$ [cm$^{-3}$]")
    title = ax.set_title("")

    def update(f):
        line.set_data(r_mm, ne[f])
        phase = "field ON" if t[f] <= t_off else "field OFF (relaxing)"
        title.set_text(f"CH4 channel  t = {t[f] * 1e9:5.0f} ns   ({phase})")
        return line, title

    anim = FuncAnimation(fig, update, frames=len(t), interval=120, blit=False)
    saved = None
    for writer, p in (
        (FFMpegWriter(fps=8), path),
        (PillowWriter(fps=8), path.replace(".mp4", ".gif")),
    ):
        try:
            anim.save(p, writer=writer)
            saved = p
            break
        except Exception as exc:  # noqa: BLE001
            print(f"  ({type(writer).__name__} unavailable: {exc})")
    plt.close(fig)
    return saved


# %%
# Run the channel solve.
# ------------------------

print(
    f"CH4 channel: E={peak_field / 1e3:.0f} kV/m pulse "
    f"({t_rise * 1e9:.0f}/{t_hold * 1e9:.0f}/{t_fall * 1e9:.0f} ns), total {t_end * 1e9:.0f} ns"
)

r = stretched_grid(wall_radius, n_points)
rho, Y0, Tg, Y_2d = build_initial_state(r)
nu_E = (2.0 * 9.1093837015e-31 / (16.0e-3 / ct.avogadro)) * nu_m
number_step = int(round(t_end / dt))
ch = PlasmaChannel(
    mechanism,
    phase,
    wall_radius,
    rho,
    Y0,
    grid=r,
    Tg_profile=(r, Tg),
    Te_profile=(r, Tg),
    Y0_profile=(r, Y_2d),
    T_amb=wall_temperature,
    nu_m=nu_m,
    nu_E=nu_E,
    electric_field=field_pulse(),  # arbitrary E(t) ramp from Python
    kappa=0.5,
    kappa_e=0.05,
    reacting=True,
    integrator="bdf",
    dt=dt,
    n_steps=number_step,
    record_every=max(number_step // 30, 1),
)

# %%
# Print status lines.
# ---------------------

ne = ch.electron_density()
Tg_result, Te_result = ch.temperature_profiles()
print(
    f"axis n_e: {ne[0, 0]:.2e} -> {ne[-1, 0]:.2e} m^-3  (relaxes; CH4 dissociation-clamped)"
)
print(
    f"axis Te : {Te_result[0, 0]:.0f} -> peak {Te_result[:, 0].max():.0f} -> "
    f"{Te_result[-1, 0]:.0f} K"
)

# %%
# Render the video.
# --------------------

saved = make_video(ch)
if saved:
    print(f"saved n_e animation -> {saved}")

# %%
# Plot the static summary figure.
# ----------------------------------

import matplotlib.pyplot as plt

r_mm = ch.r * 1e3
idx = np.unique(np.linspace(0, len(ch.t) - 1, 6).astype(int))
fig, (axn, axt) = plt.subplots(1, 2, figsize=(11, 4.2))
for f in idx:
    axn.semilogy(r_mm, ch.electron_density()[f] * 1e-6, label=f"{ch.t[f] * 1e9:.0f} ns")
axn.set_title("CH4 electron density n_e(r)")
axn.set_xlabel("Radius [mm]")
axn.set_ylabel(r"$n_e$ [cm$^{-3}$]")
axn.legend(fontsize="small")
axt.plot(ch.t * 1e9, Te_result[:, 0] * 8.617e-5, "o-", color="teal")
axt.axvspan(
    0,
    (t_rise + t_hold + t_fall) * 1e9,
    color="gold",
    alpha=0.25,
    label="field on",
)
axt.set_title("Axis electron temperature")
axt.set_xlabel("Time [ns]")
axt.set_ylabel(r"$T_e$ [eV]")
axt.legend(fontsize="small")
fig.tight_layout()
plt.show()