r"""
0D reactor: native C++ vs Python (and the 1D channel collapsed to a single cell).
=================================================================================

Three implementations of the **same** two-temperature constant-mass 0D plasma
reactor are run on the CH4 Goutier2025 mechanism under an identical NRP discharge,
and shown to agree:

1. **Python reference** -- :class:`~rizer.plasma.constant_mass_reactor.ConstantMassPlasmaReactorOde`
   (stock ``cantera`` + Python ``ExtensibleRate`` rates + ``PlasmaExtension``), the
   validated reference, driven by rizer's NRP electric circuit.
2. **Native C++ 0D** -- :class:`~rizer.plasma.constant_mass_reactor_cpp.ConstantMassPlasmaReactorOdeCpp`:
   the chemistry / two-temperature energy / Joule / collision-sum conductivity are
   evaluated natively in the ``_plasma1d`` extension (Cantera ``MultiRate`` caching),
   while ``scipy`` + the **same NRP circuit** stay in Python. ~50-300x faster per run.
3. **1D channel, 1 cell** -- :class:`~rizer.cantera_ext.plasma_channel.PlasmaChannel`
   with ``n_points=1`` (no transport), i.e. the transient 1D PDE machinery collapsed
   to a single 0D node, driven by the reference field ``E(t)``.

All three now evaluate the exact same composition-resolved collision model
(``mtcf``) live from their own evolving state every step -- `PlasmaChannel1D`
(via the ``mtcf``/``xsec_*``/``ion_Z`` arguments of ``solve_channel_transient``)
gained the same ``CollisionModel`` wiring
:class:`~rizer.plasma.constant_mass_reactor_cpp.ConstantMassPlasmaReactorOdeCpp`
already had (see ``rizer/cantera_ext/ARCHITECTURE.md``'s Model 2 "Governing
equations and closures"). All three reproduce each other to numerical noise
(<3e-3 relative, printed below). This used to require handing the 1D channel a
frozen ``sigma(Te)``/``nu_m(Te)`` table built from a single sampled
composition -- which, however carefully sampled, could only ever approximate
the live model, since the channel always recombines the table with its *own*
independently-evolving electron density. That workaround is gone now that the
channel can evaluate the collision model directly.

.. tags:: plasma, Cantera, CH4, two-temperature, 0D, reactor, C++, validation
"""  # noqa: D205

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

import time

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

import rizer.kin.extensible_rate  # noqa: F401 (Register CH4 custom rates)
import rizer.misc.units as u
from rizer.cantera_ext import PlasmaChannel
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 import ConstantMassPlasmaReactorOde
from rizer.plasma.constant_mass_reactor_cpp import ConstantMassPlasmaReactorOdeCpp
from rizer.plasma.plasma_extension import PlasmaExtension

set_mpl_style()

# %%
# Discharge 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, since gamma = np.inf) [m]
V0 = gap * np.pi * radius**2  # Initial plasma volume [m^3]

# -- Initial conditions ------------------------------------------------ #
P0 = ct.one_atm  # Initial pressure [Pa]
Tg_0 = 1000.0  # Initial gas temperature [K]
Te_0 = 1000.0  # Initial electron temperature [K]
ne_0 = 1.0e19  # Initial electron density [m^-3]

n_tot = P0 / (u.k_b * Tg_0)  # Total number density [m^-3]
x_e = ne_0 / n_tot  # Initial electron mole fraction [-]

# -- Solver parameters ------------------------------------------------- #
t_end = 100e-9  # End time [s]
dt_out = 1e-10  # Output time step [s]

# -- Polytropic process ------------------------------------------------ #
gamma = np.inf  # [-] (infinity for isochoric, 1 for isobaric, etc.)

# -- Electric parameters ----------------------------------------------- #
# Number of reflections to consider
number_of_reflections = 2  # [-]

# Generator parameters (trapezoidal pulse)
U_on = 7e3  # Voltage when on [V]
t_rise = 5e-9  # Rising time [s]
t_on = 6e-9  # On time [s]
t_fall = 6e-9  # Falling time [s]
R_g = 1.0  # Generator resistance [Ohm]

# Cable parameters (ideal lossless)
L = 6.2  # Cable length [m]
Z_c = 75.0  # Cable characteristic impedance [Ohm]
c = 1.9e8  # Cable wave propagation speed [m/s]

electrical_circuit = NRPCircuit(
    generator=TrapezoidalGenerator(
        R_g=R_g, U_on=U_on, t_rise=t_rise, t_on=t_on, t_fall=t_fall
    ),
    cable=IdealCable(L=L, Z_c=Z_c, c=c),
)

# -- Collision frequency model ----------------------------------------- #

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


# %%
# Define two helper functions.
# ----------------------------
#
# The first computes the electron number density from a 0D state vector, and the
# second integrates the 0D reactor ODE with ``scipy.integrate.ode``


def ode_electron_density(
    ode: ConstantMassPlasmaReactorOde | ConstantMassPlasmaReactorOdeCpp, y
):
    """Electron number density from a 0D state vector [Tg, Te, V, Y...]."""
    if isinstance(ode, ConstantMassPlasmaReactorOdeCpp):
        return ode._r0d.electron_density(
            y[0], y[1], np.ascontiguousarray(y[3:]), ode.mass / y[2]
        )
    plasma = ode.plasma
    plasma.Te = y[1]
    plasma.TDY = y[0], ode.mass / y[2], y[3:]
    n_tot = plasma.P / (u.k_b * plasma.mean_temperature)
    return plasma.X[plasma.species_index("e-")] * n_tot


def integrate_0d(
    ode: ConstantMassPlasmaReactorOde | ConstantMassPlasmaReactorOdeCpp,
    label: str,
    tight: bool = True,
) -> tuple[dict[str, np.ndarray], float]:
    """Drive a 0D reactor RHS with scipy vode/BDF over the NRP pulse."""
    solution = scipy.integrate.ode(ode)

    # Set the integrator options. The tight option is used for the Python reference
    # to ensure that it is accurate enough to be a reference for the C++ implementation.
    kw = dict(
        method="bdf", order=5, with_jacobian=True, first_step=1e-15, max_step=1e-10
    )
    kw.update(
        dict(atol=1e-16, rtol=1e-10, nsteps=200000)
        if tight
        else dict(atol=1e-13, rtol=1e-7, nsteps=100000)
    )

    # Set the integrator and initial value.
    solution.set_integrator("vode", **kw)
    solution.set_initial_value(np.hstack((Tg_0, Te_0, V0, Y0)), 0.0)

    # Store the results in a dictionary of lists, which will be converted to a
    # dictionary of numpy arrays at the end.
    results_list = {
        "t": [0.0],
        "Tg": [Tg_0],
        "Te": [Te_0],
        "ne": [ne_0],
        "E": [0.0],
    }

    # Create the time array.
    time_array = np.arange(dt_out, t_end, dt_out)

    # Start the timer.
    t0 = time.time()

    for t in time_array:
        # Integrate the ODE to the next time point.
        solution.integrate(t)

        if not solution.successful():
            print(f"  {label}: stopped at {solution.t:.2e}s")
            break

        # Store the results in the dictionary of lists.
        Tg, Te = solution.y[0], solution.y[1]
        ne = ode_electron_density(ode, solution.y)
        results_list["t"].append(solution.t)
        results_list["Tg"].append(Tg)
        results_list["Te"].append(Te)
        results_list["ne"].append(ne)
        results_list["E"].append(ode.plasma_voltage / gap)

    # Stop the timer and print the wall time and peak electron temperature.
    wall = time.time() - t0
    print(f"  {label}: {wall:.2f}s, peak Te={max(results_list['Te']):.0f} K")

    # Convert the results to a dictionary of numpy arrays.
    results_array = {k: np.array(v) for k, v in results_list.items()}

    return results_array, wall


# %%
# Solve the 0D reactor with both Python and native C++ implementations, and compare.
# ----------------------------------------------------------------------------------

# Create the plasma and gas objects.
plasma = ct.Solution(mechanism, "plasma", transport_model=None)
gas = ct.Solution(mechanism, "gas")

# Set the initial conditions for the plasma.
plasma.Te = Te_0
plasma.TPX = Tg_0, P0, f"CH4:{1 - 2 * x_e:.6e}, e-:{x_e:.6e}, CH4+:{x_e:.6e}"

# Get the initial mass fractions, density, and mass.
Y0 = plasma.Y.copy()  # Initial mass fractions [-]
rho_0 = plasma.density  # Initial density [kg/m^3]
m_0 = rho_0 * V0  # Initial mass [kg]


# Get the momentum transfer collision frequencies for each species.
names = plasma.species_names
mtcf = get_momentum_transfer_collision_frequencies_list(names, cfm)

# (1) Python reference 0D + NRP circuit.
plasma_extension = PlasmaExtension(
    plasma,
    momentum_transfer_collision_frequencies_list=mtcf,
)
ode_py = ConstantMassPlasmaReactorOde(
    plasma,
    plasma_extension,
    mass=m_0,
    initial_radius=radius,
    gap=gap,
    electric_circuit=electrical_circuit,
    nb_reflections=number_of_reflections,
    polytropic_index=gamma,
    p_ext=P0,
)
print("Running Python reference 0D ...")
results_python, wall_py = integrate_0d(ode_py, "python-0D", tight=False)

# (2) Native C++ 0D + the same NRP circuit.
ode_cpp = ConstantMassPlasmaReactorOdeCpp(
    mechanism,
    "plasma",
    mtcf,
    mass=m_0,
    gap=gap,
    electric_circuit=electrical_circuit,
    nb_reflections=number_of_reflections,
    polytropic_index=gamma,
    p_ext=P0,
    species_names=names,
    reacting=True,
    Te_n=600,
    spitzer=True,
)
print("Running native C++ 0D ...")
results_cpp, wall_cpp = integrate_0d(ode_cpp, "cpp-0D", tight=False)

# (3) 1D channel at a single cell, driven by the reference field E(t), using
# the SAME composition-resolved collision model (mtcf) as the 0D reactors --
# sigma and the electron-heavy elastic exchange frequency are now evaluated
# live from the channel's own evolving composition every step, exactly like
# ConstantMassPlasmaReactorOdeCpp already does (see ARCHITECTURE.md's Model 2
# "Governing equations and closures"). No frozen-composition sigma(Te)/nu_m(Te)
# table is needed anymore -- passing `mtcf` here used to require one, sampled
# from some snapshot composition, which could only ever approximate the live
# model this replaces.
number_step = int(round(t_end / dt_out))
print("Running C++ 1D channel @ 1 cell ...")
# Start the timer.
t0 = time.time()
plasma_channel_1D = PlasmaChannel(
    mechanism,
    "plasma",
    radius,
    rho_0,
    Y0,
    n_points=1,
    T_amb=Tg_0,
    mtcf=mtcf,
    Te_n=600,
    spitzer=True,
    electric_field=(results_python["t"], results_python["E"]),
    reacting=True,
    integrator="bdf",
    dt=dt_out,
    n_steps=number_step,
    # Record every output step (dt_out resolution), matching integrate_0d's own
    # cadence for the Python/C++ 0D references: Te spikes sharply over ~1-2 ns
    # here, so a coarser cadence (e.g. every ~80th frame) would UNDER-report the
    # channel's own peak Te by skipping over it, not because the internal BDF
    # dynamics are wrong -- CVODE's adaptive internal stepping already resolves
    # the spike regardless of how often the state is written out.
    record_every=1,
    gap=gap,
)
times_1D = plasma_channel_1D.t
Tg_1D, Te_1D = plasma_channel_1D.temperature_profiles()
ne_1D = plasma_channel_1D.electron_density()
# Stop the timer and print the wall time and peak electron temperature.
wall = time.time() - t0
print(f"  C++ 1D channel: {wall:.2f}s, peak Te={max(Te_1D[:, 0]):.0f} K")

# %%
# Plot the results and compare the agreement between the two implementations.
# ---------------------------------------------------------------------------


def panel(
    ax: Axes,
    y_label: str,
    title: str,
    y_python: np.ndarray,
    y_cpp: np.ndarray,
    y_cpp_1D: np.ndarray | None = None,
    log: bool = False,
) -> None:
    ax.plot(results_python["t"] * 1e9, y_python, "-", lw=2.5, label="Python 0D (ref)")
    ax.plot(results_cpp["t"] * 1e9, y_cpp, "--", lw=2, label="C++ 0D reactor")
    if y_cpp_1D is not None:
        ax.plot(times_1D * 1e9, y_cpp_1D, "-.", lw=2, label="C++ 1D channel")
    ax.set_xlabel("Time [ns]")
    ax.set_ylabel(y_label)
    ax.set_title(title)
    if log:
        ax.set_yscale("log")
    ax.legend(fontsize="small", loc="best")


fig, ax = plt.subplots(1, 3, figsize=(16, 8))
fig.suptitle("CH4 NRP 0D reactor: native C++ vs Python")


panel(
    ax[0],
    r"$T_\text{e}$ [K]",
    "Electron temperature",
    results_python["Te"],
    results_cpp["Te"],
    Te_1D[:, 0],
)
panel(
    ax[1],
    r"$T_\text{g}$ [K]",
    "Gas temperature",
    results_python["Tg"],
    results_cpp["Tg"],
    Tg_1D[:, 0],
)
panel(
    ax[2],
    r"$n_\text{e}$ [cm$^{-3}$]",
    "Electron density",
    results_python["ne"] * 1e-6,
    results_cpp["ne"] * 1e-6,
    ne_1D[:, 0] * 1e-6,
    log=True,
)
plt.show()


# ---------------------------- agreement ---------------------------------
def peak_rel(a, b):
    return abs(np.max(a) - np.max(b)) / max(abs(np.max(a)), 1e-30)


# Print peak Te, Tg, and ne for both implementations, and the relative difference between them.
print("\nComparing 0D Python and native C++ implementations:")
print("-" * 80)
print(
    f"{'Implementation':<28}"
    f"{'peak Te [K]':>14}"
    f"{'peak Tg [K]':>14}"
    f"{'peak ne [cm^-3]':>18}"
)
print(
    f"{'Python 0D (reference)':<28}"
    f"{max(results_python['Te']):>14.0f}"
    f"{max(results_python['Tg']):>14.0f}"
    f"{max(results_python['ne']) * 1e-6:>18.2e}"
)
print(
    f"{'native C++ 0D':<28}"
    f"{max(results_cpp['Te']):>14.0f}"
    f"{max(results_cpp['Tg']):>14.0f}"
    f"{max(results_cpp['ne']) * 1e-6:>18.2e}"
)
print(
    f"{'relative difference':<28}"
    f"{peak_rel(results_python['Te'], results_cpp['Te']):>14.2e}"
    f"{peak_rel(results_python['Tg'], results_cpp['Tg']):>14.2e}"
    f"{peak_rel(results_python['ne'], results_cpp['ne']):>18.2e}"
)
print("-" * 80)

print("\nComparing 0D Python and 1D C++ implementations:")
print("-" * 80)
print(
    f"{'Implementation':<28}"
    f"{'peak Te [K]':>14}"
    f"{'peak Tg [K]':>14}"
    f"{'peak ne [cm^-3]':>18}"
)
print(
    f"{'Python 0D (reference)':<28}"
    f"{max(results_python['Te']):>14.0f}"
    f"{max(results_python['Tg']):>14.0f}"
    f"{max(results_python['ne']) * 1e-6:>18.2e}"
)
print(
    f"{'native C++ 1D channel':<28}"
    f"{max(Te_1D[:, 0]):>14.0f}"
    f"{max(Tg_1D[:, 0]):>14.0f}"
    f"{max(ne_1D[:, 0]) * 1e-6:>18.2e}"
)
print(
    f"{'relative difference':<28}"
    f"{peak_rel(results_python['Te'], Te_1D[:, 0]):>14.2e}"
    f"{peak_rel(results_python['Tg'], Tg_1D[:, 0]):>14.2e}"
    f"{peak_rel(results_python['ne'], ne_1D[:, 0]):>18.2e}"
)
print("-" * 80)


print("\nComparing speed of 0D Python and native C++ implementations:")
print("-" * 80)
print(f"{'Implementation':<28}{'wall time [s]':>14}{'speedup vs Py':>16}")
print(f"{'Python 0D (reference)':<28}{wall_py:>14.2f}{1.0:>15.0f}x")
print(f"{'native C++ 0D':<28}{wall_cpp:>14.2f}{wall_py / max(wall_cpp, 1e-9):>15.0f}x")
print(
    f"{'C++ 1D channel @ 1 cell':<28}{wall:>14.2f}{wall_py / max(wall, 1e-9):>15.0f}x"
)
print("-" * 80)

# %%