r"""
Energy-equation effect of species diffusion (compressibility + enthalpy diffusion)
===================================================================================

Species diffusion in :class:`~rizer.cantera_ext.plasma_channel.PlasmaChannel`
does more than move mass around: it also feeds back into *both* temperature
equations, through two terms (see ``ARCHITECTURE.md``, Model 2, "Governing
equations"):

* **Compressibility** -- from expanding :math:`DP/Dt` for the two-temperature
  ideal-gas equation of state under Fick's law.
* **Enthalpy diffusion** (species heat transport) -- diffusing species carry
  their own sensible enthalpy with them; under Fick's law this reduces to a
  term :math:`\propto \sum_k c_{p,k}\,\nabla Y_k \cdot \nabla T`.

Both vanish when ``D_species == 0`` or the composition is uniform, and
neither is a fabricated source: they *redistribute* energy already carried
by the diffusing species, they do not create or destroy it.

This example isolates the two terms by switching off every other physics
(no conduction, no chemistry, no Joule heating) and diffusing a *localized*
composition perturbation (an N2 <-> O2, plus a tiny O2 <-> e-, mass swap)
that never reaches the wall. With everything else off, **all** of the
resulting :math:`T_g(r,t)`/:math:`T_e(r,t)` evolution shown below comes from
these two terms alone -- and the volume-integrated total (gas + electron)
internal energy stays constant while it happens, exactly as
``tests/plasma/test_plasma_channel_energy_conservation.py`` checks.

.. tags:: plasma, Cantera, channel, diffusion, energy conservation, 1D
"""  # noqa: D205

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

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

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()

# %%
# A localized, wall-isolated composition perturbation.
# -----------------------------------------------------
#
# The wall boundary condition for species is a *Dirichlet ghost* held at the
# constant ambient composition ``Y0`` (see ARCHITECTURE.md's "Dirichlet ghost
# node"), not a zero-flux condition -- so a general profile would leak
# species (and energy) through the wall, contaminating the energy-
# conservation check below. Using a Gaussian bump of half-width ``R_max/12``
# sidesteps that: it has decayed to ``~1e-64`` at the wall, far below
# floating-point resolution, so the wall stays pinned at its ambient value
# for the whole run and the domain is closed in practice.
#
# The bump swaps mass between N2 and O2 (different :math:`c_p`, so the
# enthalpy-diffusion term is non-trivial) *and* a tiny fraction of the
# electron mass (so the electron-specific versions of both terms are
# exercised too) -- while keeping :math:`\sum_k Y_k = 1` exactly everywhere.

mechanism = str(get_path_to_data("mechanisms") / "air_plasma_Laux2000.yaml")

R_max = 5.0e-3  # Outer (wall) radius [m]
T0 = 2000.0  # Uniform initial Tg = Te [K]
D = 2.0e-3  # Species mass diffusivity [m^2/s] (large, so the effect is
# visible over a short, cheap run)

plasma = ct.Solution(mechanism, "plasma")
ie = plasma.species_index("e-")
iN2 = plasma.species_index("N2")
iO2 = plasma.species_index("O2")

Ye0 = 1.0e-12
YN2_bg, YO2_bg = 0.767, 0.233 - Ye0
Y_uniform = np.zeros(plasma.n_species)
Y_uniform[iN2] = YN2_bg
Y_uniform[iO2] = YO2_bg
Y_uniform[ie] = Ye0

n_points = 161
r = np.linspace(0.0, R_max, n_points)
sigma = R_max / 12.0
bump = 0.20 * np.exp(-((r / sigma) ** 2))  # N2 <-> O2 mass swap
bump_e = 0.5 * Ye0 * np.exp(-((r / sigma) ** 2))  # tiny electron-fraction swap

Y_profile = np.zeros((n_points, plasma.n_species))
Y_profile[:, iN2] = YN2_bg + bump
Y_profile[:, iO2] = YO2_bg - bump - bump_e
Y_profile[:, ie] = Ye0 + bump_e
print(f"max |sum(Y) - 1| = {np.abs(Y_profile.sum(axis=1) - 1.0).max():.2e}")

plasma.Te = T0
plasma.TPY = T0, ct.one_atm, Y_uniform
rho = plasma.density
cp_R = plasma.standard_cp_R
W = plasma.molecular_weights
cv_k = (cp_R - 1.0) * ct.gas_constant / W  # [J/kg/K]
cve = 1.5 * ct.gas_constant / W[ie]  # analytic monatomic electron c_v, matching
# the value ReactorRHS::rates() uses internally

# %%
# Run the channel with conduction, chemistry and Joule heating all off.
# ------------------------------------------------------------------------
#
# The only thing that can move energy between ``Tg``/``Te`` here is the two
# species-diffusion energy terms under study.

dt = 5.0e-8
n_steps = 80
channel = PlasmaChannel(
    mechanism,
    "plasma",
    R_max,
    float(rho),
    Y_uniform,
    Y0_profile=(r, Y_profile),
    T_amb=T0,
    D_species=D,
    reacting=False,
    electric_field=0.0,
    nu_E=0.0,
    n_points=n_points,
    dt=dt,
    n_steps=n_steps,
    record_every=8,
    integrator="bdf",
)

t = channel.t
Tg, Te = channel.temperature_profiles()
Y = channel.mass_fractions()  # [nframes, npts, nsp]

print(
    f"max |Tg - T0| = {np.max(np.abs(Tg - T0)):.3f} K  (from compressibility "
    "+ enthalpy diffusion alone -- no conduction/chemistry/Joule)"
)
print(
    f"max |Te - T0| = {np.max(np.abs(Te - T0)):.3f} K  (electron's tiny heat "
    "capacity amplifies the same energy input into a much larger swing)"
)

# %%
# Total (gas + electron) internal energy: redistributed, not created.
# -----------------------------------------------------------------------
#
# Uses the *same* control-volume weights as ``PlasmaChannel1D.cpp``'s
# ``fvDiv()`` (interior faces at the midpoint between nodes, the outer face a
# half-interval beyond the last node), so this is an honest check of the
# scheme's own conserved quantity, not an artifact of a different
# integration rule.


def node_volumes(r):
    """Per-node control-volume area (cylindrical, per unit axial length) [m^2]."""
    faces = np.concatenate(
        (
            [0.0],
            0.5 * (r[:-1] + r[1:]),
            [r[-1] + 0.5 * (r[-1] - r[-2])],
        )
    )
    return 0.5 * (faces[1:] ** 2 - faces[:-1] ** 2)


volumes = node_volumes(r)


def total_energy(frame):
    Yk = Y[frame]
    cv_heavy = Yk @ cv_k - Yk[:, ie] * cv_k[ie]
    e_gas = rho * cv_heavy * Tg[frame]
    e_electron = rho * Yk[:, ie] * cve * Te[frame]
    return np.sum((e_gas + e_electron) * volumes)


energy = np.array([total_energy(f) for f in range(Tg.shape[0])])
rel_drift = (energy - energy[0]) / energy[0]
print(
    f"max relative energy drift = {np.max(np.abs(rel_drift)):.2e} "
    "(discretization-level residual, not a modeling gap)"
)

# %%
# Plot the composition, temperature, and energy-conservation histories.
# --------------------------------------------------------------------------

r_mm = r * 1e3
frame_idx = np.linspace(0, len(t) - 1, 5, dtype=int)
colors = plt.colormaps["viridis"](np.linspace(0.15, 0.85, len(frame_idx)))

fig, axes = plt.subplots(2, 2, figsize=(13, 9))

ax = axes[0, 0]
for f, color in zip(frame_idx, colors, strict=True):
    ax.plot(r_mm, Y[f, :, iN2], color=color, label=f"t={t[f] * 1e6:.2f} $\\mu$s")
ax.set_title("$Y_{N_2}(r)$")
ax.set_xlabel("r [mm]")
ax.set_ylabel("$Y_{N_2}$ [-]")
ax.legend(fontsize=8)

ax = axes[0, 1]
for f, color in zip(frame_idx, colors, strict=True):
    ax.plot(r_mm, Tg[f], color=color, label=f"t={t[f] * 1e6:.2f} $\\mu$s")
ax.axhline(T0, color="k", ls=":", lw=1, label="$T_0$ (no source w/o these terms)")
ax.set_title("$T_g(r)$")
ax.set_xlabel("r [mm]")
ax.set_ylabel("$T_g$ [K]")
ax.legend(fontsize=8)

ax = axes[1, 0]
for f, color in zip(frame_idx, colors, strict=True):
    ax.plot(r_mm, Te[f], color=color, label=f"t={t[f] * 1e6:.2f} $\\mu$s")
ax.axhline(T0, color="k", ls=":", lw=1, label="$T_0$")
ax.set_title("$T_e(r)$")
ax.set_xlabel("r [mm]")
ax.set_ylabel("$T_e$ [K]")
ax.legend(fontsize=8)

ax = axes[1, 1]
ax.plot(t * 1e6, rel_drift, "o-")
ax.set_title("Total energy $E(t)$")
ax.set_xlabel(r"t [$\mu$s]")
ax.set_ylabel(r"$(E(t)-E(0))/E(0)$ [-]")
ax.ticklabel_format(axis="y", style="sci", scilimits=(0, 0))

fig.suptitle(
    "Species-diffusion compressibility + enthalpy-diffusion energy terms", fontsize=12
)
plt.show()