r"""
1D plasma channel: operator-split vs monolithic CVODE (same physics)
====================================================================

Two transient 1D radial solvers for the CH4 NRP channel, built on the **same** native
C++ per-node physics (:class:`~rizer.plasma.plasma_channel_split.ChannelReactor1D`),
are run on an identical problem and shown to agree -- so the only difference between
them is *numerics*, not physics:

- **Monolithic CVODE** -- the full coupled state ``[Tg, Te, Y]`` over all nodes is
  integrated by a single stiff ``vode``/BDF; radial gas conduction couples neighbours.
  No splitting error, one large stiff system.
- **Operator splitting (Strang)** -- each global step does ½ transport, then an
  **independent** per-node stiff reaction solve (each node is a 0D reactor), then ½
  transport. The OpenSMOKE++/reactPlasFoam pattern.

Both use the composition-resolved C++ conductivity ``sigma = n_e e^2/(m_e nu_eH)`` and
are driven by the field ``E(t)`` from the 0D NRP circuit. The external field forcing is
evaluated continuously inside the reaction substep (it is not a split operator), so the
two methods match to integrator tolerance.

Result: the axis ``Te``, ``n_e``, and the total current ``I(t) = E*integral(sigma 2pi r dr)``
coincide to ~1e-6, while operator splitting runs several times faster (the stiff
electron-energy dynamics are confined to cheap independent per-node solves rather than
one large coupled Jacobian). This validates the split solver and is why it is the
preferred path for the stiff plasma source.

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

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

import time

import cantera as ct
import matplotlib.pyplot as plt
import numpy as np
import scipy.integrate

import rizer.kin.extensible_rate  # noqa: F401 (Register CH4 custom rates)
import rizer.misc.units as u
from rizer.electric_circuit.cable import IdealCable
from rizer.electric_circuit.generator import TrapezoidalGenerator
from rizer.electric_circuit.nrp_circuit import NRPCircuit
from rizer.misc.plt_utils import set_mpl_style
from rizer.misc.simulation.simulation import (
    get_default_collision_frequency_model,
    get_momentum_transfer_collision_frequencies_list,
)
from rizer.misc.utils import get_path_to_data
from rizer.plasma.constant_mass_reactor_cpp import ConstantMassPlasmaReactorOdeCpp
from rizer.plasma.plasma_channel_split import ChannelReactor1D

set_mpl_style()

# %%
# Discharge and circuit setup.
# -----------------------------

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

# -- Geometry parameters ------------------------------------------------ #
gap = 3.8e-3  # Constant gap [m]
radius = 500e-6  # Plasma radius (constant) [m]
r_wall = 4.0e-3  # Radial domain wall [m]

# -- Initial conditions -------------------------------------------------- #
P0 = ct.one_atm  # Initial pressure [Pa]
Tg0 = 1000.0  # Initial gas temperature [K]
ne0 = 1.0e19  # Initial electron density (channel peak) [m^-3]

# -- Generator / pulse parameters ---------------------------------------- #
U_on = 2000.0  # Voltage when on [V]
t_rise = 5e-9  # Rising time [s]
t_on = 25e-9  # On time [s]
t_fall = 15e-9  # Falling time [s]

# -- Time integration parameters ------------------------------------------ #
t_end = 60e-9  # End time [s]
dt = 1e-9  # Output time step [s]

# -- Radial grid parameters -------------------------------------------------- #
n_points = 11  # Number of points in the radial direction [-]

# Load the default per-species collision-frequency model.
cfm = get_default_collision_frequency_model()

# Create the plasma object and get the momentum transfer collision frequencies.
plasma = ct.Solution(mechanism, "plasma", transport_model=None)
names = plasma.species_names
mtcf = get_momentum_transfer_collision_frequencies_list(names, cfm)

# Set the initial conditions for the plasma.
x_e = ne0 / (P0 / (u.k_b * Tg0))
plasma.Te = Tg0
plasma.TPX = Tg0, P0, f"CH4:{1 - 2 * x_e:.6e}, e-:{x_e:.6e}, CH4+:{x_e:.6e}"
rho0 = plasma.density
V0 = gap * np.pi * radius**2

# 0D NRP reactor -> prescribed field E(t).
circuit = NRPCircuit(
    generator=TrapezoidalGenerator(
        R_g=1.0, U_on=U_on, U_off=0.0, t_rise=t_rise, t_on=t_on, t_fall=t_fall
    ),
    cable=IdealCable(L=6.2, Z_c=75.0, c=1.9e8),
)
ode = ConstantMassPlasmaReactorOdeCpp(
    mechanism,
    "plasma",
    mtcf,
    mass=V0 * rho0,
    gap=gap,
    electric_circuit=circuit,
    nb_reflections=1,
    polytropic_index=np.inf,
    p_ext=P0,
    species_names=names,
)
sol = scipy.integrate.ode(ode)
sol.set_integrator(
    "vode",
    method="bdf",
    order=5,
    with_jacobian=True,
    first_step=1e-13,
    max_step=1e-10,
    atol=1e-13,
    rtol=1e-7,
    nsteps=100000,
)
sol.set_initial_value(np.hstack((Tg0, Tg0, V0, plasma.Y.copy())), 0.0)
tt_list, EE_list = [0.0], [0.0]
for ti in np.arange(1e-10, t_end + 1e-12, 1e-10):
    sol.integrate(ti)
    tt_list.append(sol.t)
    EE_list.append(ode.plasma_voltage / gap)
tt, EE = np.array(tt_list), np.array(EE_list)


def E_func(t):
    """Interpolate the reference NRP field E(t) at time t."""
    return float(np.interp(t, tt, EE))


# %%
# 1D initial-condition setup.
# ----------------------------
# Uniform Tg=Te=Tg0, Gaussian n_e channel.


r = r_wall * np.sinh(3.5 * np.linspace(0, 1, n_points)) / np.sinh(3.5)

HWHM = radius  # Half-width at half-maximum of the Gaussian channel [m]
sigma = HWHM / np.sqrt(2 * np.log(2))
gaussian_profile = np.exp(-(r**2) / (2 * sigma**2))

ne1 = ne0 * gaussian_profile + 1e12

plasma = ct.Solution(mechanism, "plasma")
Y0 = np.zeros((n_points, plasma.n_species))
for i in range(n_points):
    xe = ne1[i] * u.k_b * Tg0 / P0
    plasma.TPX = Tg0, P0, f"CH4:{1 - 2 * xe:.6e}, e-:{xe:.6e}, CH4+:{xe:.6e}"
    Y0[i] = plasma.Y

# Plot the initial electron density profile.
plt.figure()
plt.plot(r * 1e3, ne1, "o-", lw=2)
plt.vlines(radius * 1e3, 0, ne0 * 1, colors="r", ls="--", lw=2)
plt.text(radius * 1e3 * 1.1, ne0 * 0.5, "Channel radius", color="r")
plt.xlabel("Radius [mm]")
plt.ylabel("Electron density [m^-3]")
plt.title("Initial electron density profile")
plt.show()


# Uniform initial conditions for gas and electron temperatures.
Tg0_grid = np.full(n_points, Tg0)
Te0_grid = np.full(n_points, Tg0)
# Uniform initial conditions for time integration (output) grid.
t_out = np.arange(0.0, t_end + 1e-12, dt)

channel_reactor = ChannelReactor1D(
    mechanism, "plasma", r, rho0, mtcf, names, reacting=True, kappa=0.5, T_amb=Tg0
)

# %%
# Run the monolithic CVODE solve.
# ---------------------------------

print("monolithic CVODE ...")
t0 = time.time()
history_monolithic = channel_reactor.solve_monolithic(
    t_out, E_func, Tg0_grid, Te0_grid, Y0
)
t_mono = time.time() - t0

# %%
# Run the operator-split solve.
# --------------------------------

print("operator split ...")
t0 = time.time()
history_split = channel_reactor.solve_split(t_out, E_func, Tg0_grid, Te0_grid, Y0)
t_split = time.time() - t0

# %%
# Helper functions to extract axis values and total current.
# --------------------------------------------------------------


def axis(h, key):
    """Extract the on-axis (r=0) value of ``key`` from a solver history."""
    return np.array([v[0] for v in h[key]])


def current(h):
    """Total current I(t) = E * integral(sigma 2 pi r dr) from a solver history."""
    return np.array(
        [
            E_func(h["t"][k])
            * np.trapezoid(
                channel_reactor.conductivity_profile(h["Tg"][k], h["Te"][k], h["Y"][k])
                * 2
                * np.pi
                * r,
                r,
            )
            for k in range(len(h["t"]))
        ]
    )


Im, Is = current(history_monolithic), current(history_split)
tm = np.array(history_monolithic["t"]) * 1e9
ts = np.array(history_split["t"]) * 1e9

# %%
# Plot the comparison panels.
# ------------------------------

fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))
fig.suptitle(
    f"1D CH4 NRP channel: monolithic CVODE ({t_mono:.1f}s) vs "
    f"operator-split ({t_split:.1f}s) -- same C++ physics"
)


def panel(a, ym, ys, ylab, title, log=False):
    """Plot a monolithic-vs-split comparison panel on axis ``a``."""
    a.plot(tm, ym, "-", lw=2.5, label="monolithic CVODE")
    a.plot(ts, ys, "--", lw=2, label="operator split")
    a.set_xlabel("Time [ns]")
    a.set_ylabel(ylab)
    a.set_title(title)
    if log:
        a.set_yscale("log")
    a.legend(fontsize="small")


panel(
    ax[0],
    axis(history_monolithic, "Te"),
    axis(history_split, "Te"),
    "Te (axis) [K]",
    "Electron temperature",
)
panel(
    ax[1],
    axis(history_monolithic, "ne"),
    axis(history_split, "ne"),
    "n_e (axis) [m^-3]",
    "Electron density",
    log=True,
)
panel(ax[2], np.abs(Im), np.abs(Is), "|I| [A]", "Total current")
fig.tight_layout(rect=(0, 0, 1, 0.95))

# %%
# Print the agreement / speedup summary.
# -----------------------------------------

rel = np.max(
    np.abs(axis(history_split, "Te") - axis(history_monolithic, "Te"))
) / np.max(np.abs(axis(history_monolithic, "Te")))
print(f"\naxis Te  max rel(split vs monolithic) = {rel:.2e}")
print(
    f"peak |I|  monolithic={np.max(np.abs(Im)):.2f} A  split={np.max(np.abs(Is)):.2f} A"
)
print(f"speedup (monolithic/split) = {t_mono / max(t_split, 1e-9):.1f}x")
plt.show()

# %%