r"""
1D plasma channel vs 0D reactor on a matched NRP discharge (current & energy)
=============================================================================

Run the **same** CH4 NRP discharge as a lumped 0D reactor and as a radially-resolved
1D channel, with **everything matched** -- mechanism, initial conditions, gap, radius,
electric field ``E(t)``, conductivity model, constant volume -- so the only difference
is 0D-lumped vs 1D-radial. The comparison is on the engineering observables: the
**total current** ``I(t)`` and the **energy deposited per pulse**.

- **0D**: :class:`~rizer.plasma.constant_mass_reactor_cpp.ConstantMassPlasmaReactorOdeCpp`
  (native C++ reactor) closed on rizer's NRP circuit, which sets the self-consistent
  field ``E(t) = V_p/gap`` from its own plasma resistance ``R_p = gap/(area*sigma)``.
- **1D**: :class:`~rizer.cantera_ext.plasma_channel.PlasmaChannel` on a radial grid
  (Gaussian ``n_e`` channel, gas heat conduction), driven by that same ``E(t)``. Its
  total current is ``I = E * integral(sigma * 2*pi*r dr)`` and its own plasma
  resistance ``R_p = gap / integral(sigma * 2*pi*r dr)`` are computed for overlay --
  if the two ``R_p(t)`` coincide, the 1D channel would close the same circuit.

Consistency of conductivity: the 1D is given ``nu_m(Te)`` = the collision-sum
``nu_eH`` the C++ 0D model uses at the initial composition, so both compute
``sigma = n_e e^2/(m_e nu_eH)`` from the same model.

The 1D **axis** tracks the 0D electron/gas temperature closely -- the local reduction
holds (over ~100 ns radial heat conduction moves only ~10 um, ``sqrt(alpha t)``). The
**integrated** current and deposited energy, however, come out several times larger in
1D: the prescribed core field also heats and **ionizes the surrounding gas**, so the
conducting channel *broadens* beyond the nominal 0D disk. The ``R_p(t)`` panel makes
this explicit -- the 1D plasma resistance falls below the 0D's as the channel widens,
i.e. a self-consistent 1D circuit would re-load to a lower field. This is the point of
a 1D model: it captures radial channel growth a lumped 0D cannot.

Caveats (this comparison isolates the radial-structure effect, it does not predict the
self-consistent operating point): the field is **prescribed** from the 0D circuit, so
there is no loading feedback onto the broadening 1D channel; and ``nu_m(Te)`` is frozen
at the initial composition. Both amplify the apparent broadening. Closing the circuit
on the 1D's own ``R_p`` (or an operator-split solver with composition-resolved ``sigma``)
is the documented follow-up.

.. tags:: plasma, Cantera, CH4, two-temperature, transient, channel, 0D, 1D, NRP, current
"""  # 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.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_cpp import ConstantMassPlasmaReactorOdeCpp

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 [m]
r_wall = 4.0e-3  # 1D outer (wall) radius [m]

# -- 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]

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

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

# Generator parameters (trapezoidal pulse)
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]

# -- Physical constants -------------------------------------------------#
electron_charge = 1.602176634e-19  # Elementary charge [C]
electron_mass = 9.1093837015e-31  # Electron mass [kg]

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

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

plasma = ct.Solution(mechanism, "plasma", transport_model=None)
names = plasma.species_names
mtcf = get_momentum_transfer_collision_frequencies_list(names, cfm)
x_e = ne_0 / (P0 / (u.k_b * Tg_0))
plasma.Te = Te_0
plasma.TPX = Tg_0, P0, f"CH4:{1 - 2 * x_e:.6e}, e-:{x_e:.6e}, CH4+:{x_e:.6e}"
rho_0 = plasma.density
V0 = gap * np.pi * radius**2
Y0 = plasma.Y.copy()

# %%
# Solve the 0D reactor closed on the NRP circuit.
# ------------------------------------------------

electrical_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 * rho_0,
    gap=gap,
    electric_circuit=electrical_circuit,
    nb_reflections=number_of_reflections,
    polytropic_index=np.inf,
    p_ext=P0,
    species_names=names,
)
solution = scipy.integrate.ode(ode)
solution.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,
)
solution.set_initial_value(np.hstack((Tg_0, Te_0, V0, Y0)), 0.0)
results_0d_lists = {
    k: [v]
    for k, v in dict(
        t=0.0, Tg=Tg_0, Te=Te_0, ne=ne_0, Vp=0.0, Rp=0.0, I=0.0, E=0.0
    ).items()
}
print("running 0D reactor ...")
t_start = time.time()
for ti in np.arange(dt_out, t_end, dt_out):
    solution.integrate(ti)
    if not solution.successful():
        print(f"  0D stopped at {solution.t:.2e}s")
        break
    Tg, Te = solution.y[0], solution.y[1]
    ne = ode._r0d.electron_density(
        Tg, Te, np.ascontiguousarray(solution.y[3:]), ode.mass / solution.y[2]
    )
    Rp, Vp = ode.plasma_resistance, ode.plasma_voltage
    for k, v in dict(
        t=solution.t,
        Tg=Tg,
        Te=Te,
        ne=ne,
        Vp=Vp,
        Rp=Rp,
        I=(Vp / Rp if Rp > 0 else 0.0),
        E=Vp / gap,
    ).items():
        results_0d_lists[k].append(v)
results_0d = {k: np.array(v) for k, v in results_0d_lists.items()}
Edep0 = np.concatenate(
    (
        [0.0],
        np.cumsum(
            0.5
            * (
                results_0d["Vp"][1:] * results_0d["I"][1:]
                + results_0d["Vp"][:-1] * results_0d["I"][:-1]
            )
            * np.diff(results_0d["t"])
        ),
    )
)
print(
    f"  0D done in {time.time() - t_start:.1f}s: "
    f"peak |I|={np.max(np.abs(results_0d['I'])):.2f} A, "
    f"E_dep={Edep0[-1] * 1e3:.3f} mJ"
)

# %%
# Build a consistent collision-sum nu_m(Te) table.
# -------------------------------------------------

Te_grid = np.linspace(800.0, 60000.0, 80)
num_grid = np.array(
    [
        ne_0
        * electron_charge
        * electron_charge
        / (electron_mass * ode._r0d.conductivity(Tg_0, Tk, Y0, rho_0))
        for Tk in Te_grid
    ]
)

# %%
# Solve the 1D channel on the same E(t) and compare.
# ----------------------------------------------------

n_points = 15
u_grid = np.linspace(0, 1, n_points)
r = r_wall * np.sinh(3.5 * u_grid) / np.sinh(3.5)  # clustered near the axis
sigma_gaussian = radius / np.sqrt(2 * np.log(2))  # Gaussian half-width
gaussian_profile = np.exp(-(r**2) / (2 * sigma_gaussian**2))
ne_1d_ic = ne_0 * gaussian_profile + 1e12  # Gaussian n_e channel
# IC matches the 0D at the axis: uniform Tg=Te=Tg_0, only n_e radial.
Y2 = np.zeros((n_points, plasma.n_species))
g = ct.Solution(mechanism, "plasma")
for i in range(n_points):
    xe = ne_1d_ic[i] * u.k_b * Tg_0 / P0
    g.TPX = Tg_0, P0, f"CH4:{1 - 2 * xe:.6e}, e-:{xe:.6e}, CH4+:{xe:.6e}"
    Y2[i] = g.Y
n_steps = int(round(t_end / 1e-9))  # 1 ns internal-step cap (resolves the ramp)
print("running 1D channel ...")
t_start = time.time()
plasma_channel_1D = PlasmaChannel(
    mechanism,
    "plasma",
    r_wall,
    rho_0,
    Y0,
    grid=r,
    gap=gap,
    Te_profile=(r, np.full(n_points, Tg_0)),
    Tg_profile=(r, np.full(n_points, Tg_0)),
    Y0_profile=(r, Y2),
    T_amb=Tg_0,
    nu_m=(Te_grid, num_grid),
    electric_field=(results_0d["t"], results_0d["E"]),
    kappa=0.5,
    reacting=True,
    integrator="bdf",
    dt=1e-9,
    n_steps=n_steps,
    record_every=max(n_steps // 40, 1),
)
times_1D = plasma_channel_1D.t
Tg_1D, Te_1D = plasma_channel_1D.temperature_profiles()
ne_1D = plasma_channel_1D.electron_density()
I1 = plasma_channel_1D.total_current()
num1 = np.interp(Te_1D.ravel(), Te_grid, num_grid).reshape(Te_1D.shape)
sig1 = ne_1D * electron_charge * electron_charge / (electron_mass * num1)
Rp1 = gap / np.trapezoid(sig1 * 2 * np.pi * r[None, :], r, axis=1)
P1 = plasma_channel_1D.joule_power(gap)
Edep1 = np.concatenate(([0.0], np.cumsum(0.5 * (P1[1:] + P1[:-1]) * np.diff(times_1D))))
print(
    f"  1D done in {time.time() - t_start:.1f}s: peak |I|={np.max(np.abs(I1)):.2f} A, "
    f"E_dep={Edep1[-1] * 1e3:.3f} mJ"
)

# %%
# Plot the results and compare the agreement between the two implementations.
# -----------------------------------------------------------------------------
#
# 1D Te/Tg/n_e vary with radius: we plot the on-axis (r=0) value AND a radial
# mean -- n_e-weighted for Te (bulk electron temperature) and volume-averaged for
# Tg and n_e. abs(I), R_p, E_dep are the channel's global (radially-integrated)
# values, not axis quantities.

two_pi_r = 2.0 * np.pi * r
ne_line = np.trapezoid(ne_1D * two_pi_r[None, :], r, axis=1)
Te_neavg = np.trapezoid(ne_1D * Te_1D * two_pi_r[None, :], r, axis=1) / np.maximum(
    ne_line, 1e-300
)
Tg_volavg = np.trapezoid(Tg_1D * two_pi_r[None, :], r, axis=1) / (np.pi * r[-1] ** 2)
ne_volavg = ne_line / (np.pi * r[-1] ** 2)

fig, ax = plt.subplots(2, 3, figsize=(15, 8))
fig.suptitle(
    "CH4 NRP discharge: 0D reactor vs 1D channel "
    "(matched mechanism, IC, gap, field, conductivity, constant volume)"
)


def field_panel(a, y0, yaxis, ymean, ylab, title, mean_lab, log=False):
    """Plot a 0D trace against the 1D on-axis and radial-mean traces."""
    a.plot(results_0d["t"] * 1e9, y0, "-", lw=2, label="0D reactor")
    a.plot(times_1D * 1e9, yaxis, "--", lw=2, label="1D axis (r=0)")
    a.plot(times_1D * 1e9, ymean, ":", lw=2, label=mean_lab)
    a.set_xlabel("Time [ns]")
    a.set_ylabel(ylab)
    a.set_title(title)
    if log:
        a.set_yscale("log")
    a.legend(fontsize="small")


def global_panel(a, y0, y1, ylab, title, lab1="1D channel", log=False):
    """Plot a 0D global trace against the 1D global (radially-integrated) trace."""
    a.plot(results_0d["t"] * 1e9, y0, "-", lw=2, label="0D reactor")
    a.plot(times_1D * 1e9, y1, "--", lw=2, label=lab1)
    a.set_xlabel("Time [ns]")
    a.set_ylabel(ylab)
    a.set_title(title)
    if log:
        a.set_yscale("log")
    a.legend(fontsize="small")


field_panel(
    ax[0, 0],
    results_0d["Te"],
    Te_1D[:, 0],
    Te_neavg,
    "Te [K]",
    "Electron temperature",
    r"1D $\langle T_e\rangle_{n_e}$ (radial)",
)
field_panel(
    ax[0, 1],
    results_0d["Tg"],
    Tg_1D[:, 0],
    Tg_volavg,
    "Tg [K]",
    "Gas temperature",
    r"1D $\langle T_g\rangle$ (radial)",
)
field_panel(
    ax[0, 2],
    results_0d["ne"],
    ne_1D[:, 0],
    ne_volavg,
    "n_e [m^-3]",
    "Electron density",
    r"1D $\langle n_e\rangle$ (radial)",
    log=True,
)
global_panel(
    ax[1, 0],
    np.abs(results_0d["I"]),
    np.abs(I1),
    "|I| [A]",
    "Total current",
    lab1=r"1D channel: $I=E\!\int\!\sigma\,2\pi r\,dr$",
)
global_panel(
    ax[1, 1],
    results_0d["Rp"],
    Rp1,
    "R_p [Ohm]",
    "Plasma resistance",
    lab1=r"1D channel: $R_p=\mathrm{gap}/\!\int\!\sigma\,2\pi r\,dr$",
    log=True,
)
global_panel(
    ax[1, 2], Edep0 * 1e3, Edep1 * 1e3, "E_dep [mJ]", "Cumulative deposited energy"
)
fig.tight_layout(rect=(0, 0, 1, 0.96))
plt.show()

# %%
# Compare peak current, deposited energy, and peak electron temperature.
# -------------------------------------------------------------------------

print(f"\n{'quantity':<24}{'0D':>12}{'1D':>12}{'1D/0D':>9}")
print(
    f"{'peak |I| [A]':<24}{np.max(np.abs(results_0d['I'])):>12.2f}"
    f"{np.max(np.abs(I1)):>12.2f}"
    f"{np.max(np.abs(I1)) / np.max(np.abs(results_0d['I'])):>9.2f}"
)
print(
    f"{'deposited energy [mJ]':<24}{Edep0[-1] * 1e3:>12.3f}"
    f"{Edep1[-1] * 1e3:>12.3f}{Edep1[-1] / max(Edep0[-1], 1e-30):>9.2f}"
)
print(f"{'peak Te [K]':<24}{results_0d['Te'].max():>12.0f}{Te_1D[:, 0].max():>12.0f}")