r"""
Transient radial conduction vs. the analytical infinite-cylinder solution
===========================================================================

With chemistry, Joule heating, and electron-heavy exchange all off, and both
temperatures starting equal (:math:`T_g=T_e`), the transient two-temperature
:class:`~rizer.cantera_ext.plasma_channel.PlasmaChannel` reduces exactly to
the classical *linear* problem of transient conduction in an infinite
cylinder whose surface temperature is suddenly changed and held fixed
(Carslaw & Jaeger, *Conduction of Heat in Solids*, 1959; e.g. Incropera,
*Fundamentals of Heat and Mass Transfer*, Ch. 5):

.. math::

    \rho c_v \frac{\partial T}{\partial t}
        = \frac{1}{r}\frac{\partial}{\partial r}\!\left(r\,\kappa\,
          \frac{\partial T}{\partial r}\right),
    \qquad T(r,0)=T_i,\quad T(R,t)=T_s .

For constant :math:`\rho`, :math:`c_v`, :math:`\kappa` (diffusivity
:math:`\alpha=\kappa/(\rho c_v)`) this has the closed-form Bessel-series
solution

.. math::

    \frac{T(r,t)-T_s}{T_i-T_s}
        = \sum_{n=1}^{\infty} \frac{2}{\lambda_n J_1(\lambda_n)}\,
          J_0\!\left(\lambda_n \frac{r}{R}\right)\,
          \exp\!\left(-\lambda_n^2\,\frac{\alpha t}{R^2}\right),

where :math:`\lambda_n` are the positive roots of :math:`J_0(\lambda_n)=0`.
This example compares the recorded :math:`T_g(r,t)` field against that series
at **every** recorded radius and time (not just a few snapshots, as in
``tests/plasma/test_plasma_channel_diffusion.py::test_matches_analytical_cylinder_conduction``),
and plots the residual as a 2D map. ``Te`` never moves under these settings
(no source, no transport term for it), which the last panel also confirms.

The residual map has one honest wrinkle worth calling out rather than hiding:
a brief hot streak next to the wall in the first few recorded frames. The
instant the wall's Dirichlet condition switches on, the true solution has a
boundary layer of zero initial thickness that grows as :math:`\sqrt{\alpha t}`
-- any time integrator needs a handful of steps to resolve that, so the very
first frames carry a larger, rapidly-decaying error there. It is a startup
artifact of the time discretisation, not the spatial one, and it is why the
unit test this example extends only starts comparing well past it.

.. tags:: plasma, Cantera, channel, conduction, validation, analytical, 1D
"""  # noqa: D205

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

import cantera as ct
import matplotlib.pyplot as plt
import numpy as np
from scipy.special import j0, j1, jn_zeros

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

# %%
# The analytical solution.
# -------------------------
#
# Same closed form used by the unit test: a truncated sum over the positive
# roots of :math:`J_0`. 1000 terms is ample at the Fourier numbers used here
# (the summed terms decay as :math:`\exp(-\lambda_n^2\alpha t/R^2)` with
# :math:`\lambda_n` growing like :math:`n\pi`).


def cylinder_conduction(r, t, R, alpha, T_i, T_s, n_terms=1000):
    """Transient conduction in an infinite cylinder, T(r,t), on a (t, r) grid."""
    r = np.asarray(r, float)
    t = np.asarray(t, float)
    lam = jn_zeros(0, n_terms)
    Fo = alpha * t[:, None] / R**2  # shape (nt, 1), broadcasts against r
    theta = np.zeros((t.size, r.size))
    for lam_n in lam:
        theta += (
            (2.0 / (lam_n * j1(lam_n)))
            * j0(lam_n * r / R)[None, :]
            * np.exp(-(lam_n**2) * Fo)
        )
    return T_s + (T_i - T_s) * theta


# %%
# Set up a non-reacting, single-initial-temperature conduction problem.
# ------------------------------------------------------------------------
#
# ``kappa`` is injected as a constant table; the mechanism/composition only
# sets ``rho`` and ``cv`` (via a throwaway state at a representative
# temperature -- the real gas's ``cv(T)`` varies mildly over the modest
# temperature range used here).

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

R_max = 5.0e-3  # Outer (wall) radius [m]
T_hot, T_amb, T_rep = 2200.0, 2000.0, 2100.0  # Initial, wall, reference [K]
kappa = 1.0  # Constant gas thermal conductivity [W/m/K]

g = ct.Solution(mechanism, "plasma")
ie = g.species_index("e-")
Y0 = np.zeros(g.n_species)
Y0[g.species_index("N2")] = 0.767
Y0[g.species_index("O2")] = 0.233
Y0[ie] = 1.0e-12

g.TPY = T_rep, ct.one_atm, Y0
rho = g.density
cp_R = g.standard_cp_R
cv_k = (cp_R - 1.0) * ct.gas_constant / g.molecular_weights
cv_heavy = float(np.sum(Y0 * cv_k) - Y0[ie] * cv_k[ie])
alpha = kappa / (rho * cv_heavy)
print(f"rho = {rho:.4f} kg/m^3, cv = {cv_heavy:.1f} J/kg/K, alpha = {alpha:.3e} m^2/s")

# %%
# Run the channel out to Fourier number Fo = alpha*t/R^2 = 0.6, recording
# every step so the (r, t) comparison grid is dense in both directions.
# ---------------------------------------------------------------------------
#
# Uses the ``"bdf"`` backend (CVODE's own graded startup -- a few tiny first
# internal steps -- resolves the sharp early-time transient right after the
# wall's step change much better than a single large fixed Backward-Euler
# step would; see ``ARCHITECTURE.md``'s Model 2 "Time integration").

Fo_end = 0.6
n_steps = 200
t_end = Fo_end * R_max**2 / alpha
dt = t_end / n_steps

rg = np.array([0.0, R_max])
Tg_hot = np.array([T_hot, T_hot])
channel = PlasmaChannel(
    mechanism,
    "plasma",
    R_max,
    float(rho),
    Y0,
    Tg_profile=(rg, Tg_hot),
    Te_profile=(rg, Tg_hot),
    T_amb=T_amb,
    kappa=kappa,
    nu_E=0.0,
    reacting=False,
    electric_field=0.0,
    n_points=201,
    dt=dt,
    n_steps=n_steps,
    record_every=1,
    integrator="bdf",
)

r = channel.r
Tg_num_all, Te_num = channel.temperature_profiles()

# Drop the t=0 frame: the analytical series' r=R boundary term is built from
# J_0(lambda_n), which is *exactly* zero by definition of lambda_n as J_0's
# roots -- so the formula evaluates to T_s at r=R for every t>=0, including
# t=0, while the model's true initial state is uniformly T_hot everywhere
# (including the boundary node). That is a convention mismatch at the single
# point (r=R, t=0), not a numerics discrepancy, so -- exactly as in the unit
# test -- the comparison starts at the first *post-initial* recorded frame.
t = channel.t[1:]
Tg_num = Tg_num_all[1:]
Tg_exact = cylinder_conduction(r, t, R_max, alpha, T_hot, T_amb)

residual = Tg_num - Tg_exact
max_by_frame = np.max(np.abs(residual), axis=1)
print(f"max |Tg_num - Tg_exact| = {max_by_frame.max():.3f} K (whole run)")
print(
    f"                       = {max_by_frame[20:].max():.3f} K (excluding the first 20 frames)"
)
print(f"max |Te_num - T_hot|    = {np.max(np.abs(Te_num - T_hot)):.2e} K (Te is inert)")
print(
    "Note: the largest residual sits at the very first few frames, right at\n"
    "the node next to the wall -- an expected startup artifact (any time\n"
    "integrator needs a few steps to resolve the boundary layer that forms\n"
    "the instant the wall's Dirichlet condition switches on), not a bug. It\n"
    "decays quickly and is visible as a brief hot streak near the wall in\n"
    "the bottom-left residual map below."
)

# %%
# Plot the numerical field, the analytical field, and their residual, plus a
# handful of radial profiles overlaid.
# -------------------------------------------------------------------------------

fig, axes = plt.subplots(2, 2, figsize=(13, 9))
r_mm, t_us = r * 1e3, t * 1e6
vmin, vmax = T_amb, T_hot

ax = axes[0, 0]
im = ax.pcolormesh(r_mm, t_us, Tg_num, shading="auto", vmin=vmin, vmax=vmax)
ax.set_title("Numerical $T_g(r,t)$")
ax.set_xlabel("r [mm]")
ax.set_ylabel(r"t [$\mu$s]")
fig.colorbar(im, ax=ax, label="T [K]")

ax = axes[0, 1]
im = ax.pcolormesh(r_mm, t_us, Tg_exact, shading="auto", vmin=vmin, vmax=vmax)
ax.set_title("Analytical $T(r,t)$")
ax.set_xlabel("r [mm]")
ax.set_ylabel(r"t [$\mu$s]")
fig.colorbar(im, ax=ax, label="T [K]")

ax = axes[1, 0]
res_max = max(np.abs(residual).max(), 1e-6)
im = ax.pcolormesh(
    r_mm, t_us, residual, shading="auto", cmap="RdBu_r", vmin=-res_max, vmax=res_max
)
ax.set_title("Residual (num. $-$ analytical)")
ax.set_xlabel("r [mm]")
ax.set_ylabel(r"t [$\mu$s]")
fig.colorbar(im, ax=ax, label="$\\Delta T$ [K]")

ax = axes[1, 1]
frame_idx = np.linspace(0, len(t) - 1, 4, dtype=int)
colors = plt.colormaps["viridis"](np.linspace(0.15, 0.85, len(frame_idx)))
for f, color in zip(frame_idx, colors, strict=True):
    Fo = alpha * t[f] / R_max**2
    ax.plot(r_mm, Tg_exact[f], "-", color=color, label=f"Fo={Fo:.2f}")
    ax.plot(r_mm[::12], Tg_num[f, ::12], "o", color=color, ms=10)
ax.set_title("Radial profiles")
ax.set_xlabel("r [mm]")
ax.set_ylabel("$T_g$ [K]")
ax.legend(fontsize=8, title="line = analytical\nmarker = numerical", title_fontsize=7)

fig.suptitle("PlasmaChannel1D vs. analytical infinite-cylinder conduction", fontsize=13)
# NB: no fig.tight_layout() here -- the mplstyle's constrained_layout is
# already active, and it conflicts with tight_layout() once a colorbar
# has been added.
plt.show()