r"""
Channel diameter expansion: 1D model vs engineering law vs experiment
=====================================================================

Compares three channel-diameter predictions/measurements:

1. **1D transient model** (:class:`rizer.cantera_ext.PlasmaChannel`) -- the
   radial two-temperature reacting plasma channel, run in its diffusion regime
   (radial heat conduction of a hot core). Diameter ``D(t) = 2 R(t)`` from the
   half-maximum radius.
2. **Engineering law** (rizer's diffusion phase,
   :mod:`rizer.hybrid.engineering_model.model.diffusion`):
   ``R_eq(t) = R_0 + sqrt(alpha * t)``.
3. **Experiment**: measured discharge-channel diameters,
   ``data/experiments/T314/run49/section_diameters_mean.npz``.

Caveats (see ``rizer/cantera_ext/ARCHITECTURE.md``):

* The 1D model here is **constant-volume / conduction-driven**, so it reproduces
  the *diffusive* (late) expansion, not the fast hydrodynamic/rarefaction phase
  that dominates the first ~hundreds of ns of the experiment.
* It runs on the native-rate **air** mechanism; the experiment is CH4/N2. A
  quantitative CH4 match needs the Goutier mechanism's custom rates ported to
  C++. The *diffusive scaling* ``R ~ sqrt(alpha t)`` is gas-independent.

.. tags:: plasma, Cantera, channel, expansion, diffusion, experiment, NRP
"""  # 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.io.experimental_data.load_experiment_data import load_experiment_data
from rizer.misc.plt_utils import set_mpl_style
from rizer.misc.utils import get_path_to_data

set_mpl_style()

# %%
# Discharge setup.
# ----------------

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

# -- Scenario parameters ------------------------------------------------ #
R0 = 0.35e-3  # ~ measured initial radius [m]
T_core = 8000.0  # ~ post-pulse core temperature [K]
T_amb = 2000.0  # Ambient temperature [K]
lam = 2.0  # Thermal conductivity used by the 1D model [W/m/K]

# -- Solver parameters ---------------------------------------------------- #
t_end = 1.0e-6  # End time [s]
n_steps = 40  # Number of solver steps [-]


def run_channel(R0, T_core, T_amb, lam, t_end, n_steps):
    """Run the 1D channel (diffusion regime) and return (t, D) [s, m]."""
    g = ct.Solution(mechanism, "plasma")
    g.TPX = 0.5 * (T_core + T_amb), ct.one_atm, "N2:0.79, O2:0.21"
    g.equilibrate("TP")
    R_max = 8.0 * R0
    rg = np.array([0.0, 0.7 * R0, R0, 1.3 * R0, R_max])
    Tg = np.array([T_core, T_core, 0.5 * (T_core + T_amb), T_amb, T_amb])
    ch = PlasmaChannel(
        mechanism,
        "plasma",
        R_max,
        float(g.density),
        g.Y.copy(),
        Tg_profile=(rg, Tg),
        Te_profile=(rg, Tg),
        T_amb=T_amb,
        kappa=lam,
        nu_E=1.0e11,
        reacting=True,
        dt=t_end / n_steps,
        n_steps=n_steps,
        record_every=max(n_steps // 20, 1),
        n_points=41,
    )
    return ch.t, ch.channel_diameter("half_max"), ch


# %%
# Load the experimental data.
# ----------------------------

# Experiment (T314/run49): time [ns], diameter [mm].
exp = load_experiment_data(
    get_path_to_data("experiments", "T314", "run49", "section_diameters_mean.npz")
)
t_exp_ns, D_exp_mm = exp["time_ns"], exp["diameter_mm"]

# %%
# Run the 1D model (diffusion regime), sized to the experiment's scale.
# -----------------------------------------------------------------------

t, D, ch = run_channel(R0, T_core, T_amb, lam, t_end=t_end, n_steps=n_steps)

# %%
# Compute the engineering diffusion law.
# -----------------------------------------

# Engineering diffusion law R_eq = R0 + sqrt(alpha t), alpha = lam/(rho cp).
g = ct.Solution(mechanism, "plasma")
g.TPX = 0.5 * (T_core + T_amb), ct.one_atm, "N2:0.79, O2:0.21"
g.equilibrate("TP")
alpha = lam / (g.density * g.cp_mass)
t_eng = np.linspace(0, t_end, 100)
D_eng_mm = 2.0 * (R0 + np.sqrt(alpha * t_eng)) * 1e3

# %%
# Print the comparison summary.
# --------------------------------

print(f"alpha (engineering) = {alpha:.3e} m^2/s")
print(f"1D model D: {D[0] * 1e3:.3f} -> {D[-1] * 1e3:.3f} mm over {t[-1] * 1e9:.0f} ns")
print(
    f"experiment D: {D_exp_mm[0]:.3f} -> {D_exp_mm[-1]:.3f} mm "
    f"over {t_exp_ns[0]:.0f}-{t_exp_ns[-1]:.0f} ns"
)

# %%
# Plot the results and compare the agreement between the three approaches.
# ----------------------------------------------------------------------------

fig, ax = plt.subplots()
ax.set_title(
    "Discharge-channel diameter expansion\n"
    "(1D conduction model = diffusion regime; air vs CH4 expt.)"
)
ax.set_xlabel("Time [ns]")
ax.set_ylabel("Diameter [mm]")
ax.plot(t * 1e9, D * 1e3, "o-", label="1D Cantera channel (air)")
ax.plot(t_eng * 1e9, D_eng_mm, "--", label=r"engineering $R_0+\sqrt{\alpha t}$")
ax.plot(t_exp_ns, D_exp_mm, "ks", label="experiment T314/run49 (CH4)")
ax.legend()
fig.tight_layout()
plt.show()