r"""
Plot CH₄ average momentum transfer collision frequency.
=======================================================

The average momentum transfer collision frequency of methane (CH₄) with electrons is plotted
as a function of the electron temperature.
Data is taken from the Hayashi and Morgan databases, which are available on the LXCat website.

This collision frequency corresponds to the collision process e- + CH₄ -> e- + CH₄,
which is an elastic collision where the electron is scattered by the methane molecule without any energy loss.

It corresponds to :math:`\bar{\nu_{e h}}^{(1)}`, the average momentum transfer collision frequency, in s^-1,
between electrons and heavy particles, defined in equation (II 6.29) of [Mitchner1973]_.


Notes
-----
* :math:`\nu_{e h}^{(1)}` is the momentum transfer collision frequency, in s^-1, between electrons
  and heavy particles (e.g. neutrals, ions), defined in equation (II 5.6) of [Mitchner1973]_.
* :math:`\bar{\nu_{e h}}^{(1)}` is the average momentum transfer collision frequency, in s^-1, between
  electrons and heavy particles (e.g. neutrals, ions), defined in equation (II 6.29) of [Mitchner1973]_.
* :math:`\nu_{e h}^{E} = \frac{2 m_e}{m_h} \nu_{e h}^{(1)}` is the energy exchange collision frequency,
  in s^-1, between electrons and heavy particles (e.g. neutrals, ions), defined in equation (II 7.6)
  of [Mitchner1973]_.
* :math:`\bar{\nu_{e h}}^{E}` is the average energy exchange collision frequency, in s^-1.


.. tags::

    collision frequency, momentum transfer
"""  # noqa: D205

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

import matplotlib.pyplot as plt
import numpy as np

import rizer.misc.units as u
from rizer.misc.plt_utils import set_mpl_style
from rizer.plasma.collision_frequency import (
    get_momentum_transfer_collision_frequency_model,
)

set_mpl_style()

# %%
# Load the cross section data.
# ----------------------------

# Load cross section data from the Hayashi and Morgan databases.
# They will be used as tabulated cross sections to compute the mean momentum transfer collision frequency.
frequency_model_Hayashi = get_momentum_transfer_collision_frequency_model(
    species="CH4",
    database_to_use="Hayashi.txt",
)
frequency_model_Morgan = get_momentum_transfer_collision_frequency_model(
    species="CH4",
    database_to_use="Morgan.txt",
)

# Compute the mean momentum transfer collision frequency as a function of
# the electron temperature for a given heavy particle density.
electron_temperatures = np.arange(100, 50000, 100, dtype=float)  # K

# Calculate the heavy particle density using the ideal gas law at room temperature and atmospheric pressure.
P = 1.01325e5  # Pa
T = 298.15  # K
n_CH4 = P / (u.k_b * T)  # m⁻³, ideal gas law

average_frequencies_Hayashi = []
average_frequencies_Morgan = []
for T in electron_temperatures:
    average_frequencies_Hayashi.append(
        frequency_model_Hayashi.get_mean_momentum_transfer_collision_frequency(T, n_CH4)
    )
    average_frequencies_Morgan.append(
        frequency_model_Morgan.get_mean_momentum_transfer_collision_frequency(T, n_CH4)
    )

# %%
# Plot the collision frequency as a function of temperature.
# ----------------------------------------------------------

fig, ax1 = plt.subplots()
ax1.set_xlabel(r"$T_\mathrm{e}$ [K]")
ax1.set_ylabel(r"$\bar{\nu}_{\mathrm{e, CH_4}}$ [s⁻¹]")
ax1.plot(
    electron_temperatures,
    average_frequencies_Hayashi,
    label="Hayashi",
)
ax1.plot(
    electron_temperatures,
    average_frequencies_Morgan,
    label="Morgan",
)
ax1.tick_params(axis="x")
ax1.set_xscale("log")
ax1.set_yscale("log")
ax1.set_title(
    "Mean momentum transfer collision frequency of CH₄ with electrons vs. Electron Temperature"
)
ax1.legend(loc="upper left")

plt.show()

# %%