r"""
Transient two-temperature reacting plasma channel
==================================================

:class:`~rizer.cantera_ext.plasma_channel.PlasmaChannel` is the transient, 1D
radial, two-temperature (:math:`T_e \neq T_g`), finite-rate reacting counterpart
of rizer's 0D :class:`~rizer.plasma.constant_mass_reactor.ConstantMassPlasmaReactorOde`.
It couples plasma chemistry, the gas and electron energy equations, and radial
heat conduction, and exposes the conductive-channel diameter ``D(t) = 2 R(t)``.

This example evolves a hot ionised core into cool ambient air and shows the
radial temperature profiles relaxing while the channel radius grows diffusively
(``R(t) ~ sqrt(alpha t)``), the late-phase mechanism behind rizer's engineering
channel-expansion law.

.. tags:: plasma, Cantera, two-temperature, transient, channel, expansion, 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()

# %%
# Set up a hot core in air and run the transient channel.
# ---------------------------------------------------------

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

# -- Initial conditions --------------------------------------------------- #
outer_radius = 4.0e-3  # Domain outer radius [m]
initial_core_radius = 0.8e-3  # Initial hot-core radius [m]
core_temperature = 6000.0  # Initial core temperature [K]
ambient_temperature = 2500.0  # Ambient (far-field) temperature [K]

gas = ct.Solution(mechanism, "plasma")
gas.TPX = (
    0.5 * (core_temperature + ambient_temperature),
    ct.one_atm,
    "N2:0.79, O2:0.21",
)
gas.equilibrate("TP")

# Piecewise-linear initial radial temperature profile: hot core, transition
# shell, and ambient far field.
profile_radii = np.array(
    [
        0.0,
        0.7 * initial_core_radius,
        initial_core_radius,
        1.3 * initial_core_radius,
        outer_radius,
    ]
)
profile_temperatures = np.array(
    [
        core_temperature,
        core_temperature,
        0.5 * (core_temperature + ambient_temperature),
        ambient_temperature,
        ambient_temperature,
    ]
)

channel = PlasmaChannel(
    mechanism,
    "plasma",
    outer_radius,
    float(gas.density),
    gas.Y.copy(),
    Tg_profile=(profile_radii, profile_temperatures),
    Te_profile=(profile_radii, profile_temperatures),
    T_amb=ambient_temperature,
    kappa=2.0,
    nu_E=1.0e11,
    reacting=True,
    dt=2.0e-7,
    n_steps=30,
    record_every=5,
    n_points=31,
)

# %%
# Extract and plot the results.
# ------------------------------
# Radial gas-temperature profiles over time, and the channel radius R(t).

radii = channel.r
gas_temperature_t, _ = channel.temperature_profiles()
channel_radius = channel.channel_radius("half_max")

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))

for frame, time_s in enumerate(channel.t):
    ax1.plot(radii * 1e3, gas_temperature_t[frame], label=f"{time_s * 1e6:.1f} us")
ax1.set_title("Gas temperature profiles")
ax1.set_xlabel("Radius [mm]")
ax1.set_ylabel("Tg [K]")
ax1.legend(fontsize="small")

ax2.set_title("Channel radius (diffusive growth)")
ax2.set_xlabel("Time [us]")
ax2.set_ylabel("R [mm]")
ax2.plot(channel.t * 1e6, channel_radius * 1e3, "o-")
fig.tight_layout()
plt.show()