r"""
Steady LTE arc column: Cantera 1D solver vs analytical Elenbaas-Heller
======================================================================

:class:`~rizer.cantera_ext.thermal_plasma_column.ThermalPlasmaColumn` solves the
steady radial Elenbaas-Heller equation

.. math::

    \frac{1}{r}\frac{d}{dr}\!\left(r\,\kappa(T)\frac{dT}{dr}\right)
        + \sigma(T) E^2 = 0

numerically with a custom Cantera ``Domain1D`` (Newton + adaptive grid), and is
validated against the analytical
:class:`~rizer.thermal_plasma.elenbaas_heller.ElenbaasHeller` solution. Fed the
same tabulated LTE properties, the two agree to within the analytical model's
piecewise-linear linearization error.

.. tags:: plasma, Cantera, Elenbaas-Heller, arc, LTE, 1D
"""  # noqa: D205

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

import matplotlib.pyplot as plt
import numpy as np

from rizer.cantera_ext import ThermalPlasmaColumn
from rizer.misc.plt_utils import set_mpl_style
from rizer.thermal_plasma.elenbaas_heller import ElenbaasHeller
from rizer.thermal_plasma.fit_LTE_data import FitLTEData

set_mpl_style()

# %%
# Load H2 LTE data and build both models for the same field.
# ------------------------------------------------------------

column_radius = 10e-3  # Column radius [m]
electric_field = 5000.0  # Electric field [V/m]

lte_data = FitLTEData(
    gas_name_transport="H2",
    gas_name_radiation="H2",
    pressure_atm=1,
    source_transport="Boulos2023",
    source_radiation="Gueye2017",
    emission_radius_mm=0,
    max_temperature_fit=12000.0,
)

column = ThermalPlasmaColumn(
    R=column_radius, electric_field=electric_field, gas_data=lte_data, n_points=61
)
elenbaas_heller = ElenbaasHeller(
    R=column_radius, electric_field=electric_field, current=None, gas_data=lte_data
)

# %%
# Compare the radial temperature profiles, then plot the results.
# -------------------------------------------------------------------

radius, temperature_numerical = column.temperature_profile()
temperature_analytical = np.array(
    [elenbaas_heller.get_temperature_vs_radius(ri) for ri in radius]
)

print(
    f"T_center: numerical {temperature_numerical[0]:.0f} K vs analytical "
    f"{elenbaas_heller.get_temperature_vs_radius(0.0):.0f} K"
)
print(
    f"current : numerical {column.current():.1f} A vs analytical "
    f"{elenbaas_heller.analytical_current():.1f} A"
)

fig, ax = plt.subplots()
ax.set_title(
    f"H2 arc column, E = {electric_field:.0f} V/m, R = {column_radius * 1e3:.0f} mm"
)
ax.set_xlabel("Radius [mm]")
ax.set_ylabel("Temperature [K]")
ax.plot(radius * 1e3, temperature_numerical, label="Numerical (Cantera Domain1D)")
ax.plot(radius * 1e3, temperature_analytical, "--", label="Analytical Elenbaas-Heller")
ax.legend()
fig.tight_layout()
plt.show()