r"""
Example: Calculating the properties of an :math:`CH_4` plasma.
==============================================================

Compute the properties of an :math:`CH_4` plasma using minplascalc.
Data are compared with reference data from the literature.

Properties calculated include:

- Density
- Enthalpy
- Heat capacity
- Viscosity
- Thermal conductivity
- Electrical conductivity


Species data is loaded from `./data/minplascalc/species` folder.
These data have been generated using the `generate_minplascalc_data.py` script.

Reference data are [Wu2016]_ and [Niu2016]_.
"""  # noqa: D205

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

import matplotlib.pyplot as plt
import minplascalc as mpc
import numpy as np

from rizer.io.thermo_transport_data_reader import ThermoTransportDataReader
from rizer.misc.utils import get_path_to_data

# %%
# Create mixture object for the species we're interested in.
# ----------------------------------------------------------
#
# Next, we create a minplascalc LTE mixture object. Here we use a helper function
# in minplascalc which creates the object directly from a list of the species names.

species_and_concentrations_low_temperature = {
    "CH4": 1,
    "CH3": 0,
    "CH2": 0,
    "CH": 0,
    "C": 0,
    "C+": 0,
    "C2H6": 0,
    "C2H4": 0,
    "C2H2": 0,
    "C2H": 0,
    "C2": 0,
    "H2": 0,
    "H": 0,
    "H+": 0,
}

species_and_concentrations_mid_temperature = {
    "C": 1,
    "C+": 0,
    "C2": 0,
    "H": 4,
    "H+": 0,
    "H2": 0,
    "H2+": 0,
}

species_and_concentrations_high_temperature = {
    "C": 1,
    "C+": 0,
    "C++": 0,
    "C+++": 0,
    "C++++": 0,
    "C+++++": 0,
    "H": 4,
    "H+": 0,
}


SPECIES_PATH = get_path_to_data("minplascalc", "species")

species_and_concentrations_list = [
    species_and_concentrations_low_temperature,
    species_and_concentrations_mid_temperature,
    species_and_concentrations_high_temperature,
]

species_mpc_list = [
    [
        mpc.species.from_file(SPECIES_PATH / f"{species_name}.json")
        for species_name in species_and_concentrations
    ]
    for species_and_concentrations in species_and_concentrations_list
]

methane_mixture_list = [
    mpc.mixture.LTE(
        species=species_mpc,
        x0=species_and_concentrations.values(),
        T=300,
        P=101325,
        gfe_initial_particles=1e23,
        gfe_rtol=1e-10,
        gfe_max_iter=1000,
    )
    for species_mpc, species_and_concentrations in zip(
        species_mpc_list, species_and_concentrations_list
    )
]


# %%
# Set a range of temperatures to calculate the equilibrium compositions at.
# -------------------------------------------------------------------------

# Set a range of temperatures to calculate the equilibrium compositions at.
low_temperatures = np.arange(400, 6_000 + 100, 100, dtype=float)
mid_temperatures = np.arange(6_000, 12_000 + 250, 250, dtype=float)
high_temperatures = np.arange(12_000, 90_000 + 250, 250, dtype=float)
temperatures_list = [low_temperatures, mid_temperatures, high_temperatures]

# Initialise lists to store the property values at each temperature.
ni_list: list[list[np.ndarray]] = [[] for _ in range(len(temperatures_list))]
density_list: list[np.ndarray] = [
    np.zeros_like(temperatures, dtype=float) for temperatures in temperatures_list
]
h_list: list[np.ndarray] = [
    np.zeros_like(temperatures, dtype=float) for temperatures in temperatures_list
]
cp_list: list[np.ndarray] = [
    np.zeros_like(temperatures, dtype=float) for temperatures in temperatures_list
]
viscosity_list: list[np.ndarray] = [
    np.zeros_like(temperatures, dtype=float) for temperatures in temperatures_list
]
thermal_conductivity_list: list[np.ndarray] = [
    np.zeros_like(temperatures, dtype=float) for temperatures in temperatures_list
]
electrical_conductivity_list: list[np.ndarray] = [
    np.zeros_like(temperatures, dtype=float) for temperatures in temperatures_list
]

# %%
# Perform the composition calculations.
# --------------------------------------
#
# Now we can perform the property calculations.
#
# Note that execution of this calculation is fairly compute intensive and the following code
# snippet may take several minute to complete.

for i, (methane_mixture, temperatures) in enumerate(
    zip(methane_mixture_list, temperatures_list)
):
    for idx_T, T in enumerate(temperatures):
        print(f"Calculating properties at T = {T} K...")
        methane_mixture.T = T
        P = 1e5
        k_B = 1.38e-23
        methane_mixture.gfe_initial_particles = P / (k_B * T) / 10.0

        ni_list[i].append(methane_mixture.calculate_composition())

        density_list[i][idx_T] = methane_mixture.calculate_density()
        h_list[i][idx_T] = methane_mixture.calculate_enthalpy()
        cp_list[i][idx_T] = methane_mixture.calculate_heat_capacity()
        viscosity_list[i][idx_T] = methane_mixture.calculate_viscosity()
        thermal_conductivity_list[i][idx_T] = (
            methane_mixture.calculate_thermal_conductivity()
        )
        electrical_conductivity_list[i][idx_T] = (
            methane_mixture.calculate_electrical_conductivity()
        )


print("Done!")


# %%
# Save the data in a CSV file.
# ----------------------------


# For enthalpy, compute the difference between the first temperature of the second temperature range,
# and the last temperature of the first temperature range.
# Then, subtract this difference from the enthalpy of the second temperature range.
delta_h = h_list[1][0] - h_list[0][-1]
h_list[1] -= delta_h
# Do the same for the third temperature range.
delta_h = h_list[2][0] - h_list[1][-1]
h_list[2] -= delta_h

# Remove the last point of the first zone, and the first point of the second zone, for all properties.
temperatures_list[0] = temperatures_list[0][:-1]
ni_list[0] = ni_list[0][:-1]
density_list[0] = density_list[0][:-1]
h_list[0] = h_list[0][:-1]
cp_list[0] = cp_list[0][:-1]
viscosity_list[0] = viscosity_list[0][:-1]
thermal_conductivity_list[0] = thermal_conductivity_list[0][:-1]
electrical_conductivity_list[0] = electrical_conductivity_list[0][:-1]
temperatures_list[1] = temperatures_list[1][1:]
ni_list[1] = ni_list[1][1:]
density_list[1] = density_list[1][1:]
h_list[1] = h_list[1][1:]
cp_list[1] = cp_list[1][1:]
viscosity_list[1] = viscosity_list[1][1:]
thermal_conductivity_list[1] = thermal_conductivity_list[1][1:]
electrical_conductivity_list[1] = electrical_conductivity_list[1][1:]

# Concatenate the data into a single array.
data = np.array(
    [
        np.concatenate(temperatures_list),
        np.concatenate(density_list),
        np.concatenate(h_list),
        np.concatenate(cp_list),
        np.concatenate(viscosity_list),
        np.concatenate(thermal_conductivity_list),
        np.concatenate(electrical_conductivity_list),
    ]
).transpose()

# Save the data to a CSV file.
np.savetxt(
    get_path_to_data("minplascalc", "CH4_1atm_minplascalc.csv", force_return=True),
    data,
    delimiter=",",
    header="T [K], density [kg/m^3], enthalpy [J/kg], Cp [J/(kg.K)], "
    "viscosity [Pa.s], thermal conductivity [W/(m.K)], electrical conductivity [S/m]",
    comments="",
)


# %%.
# Load reference data.
# --------------------

# Load reference data from Wu2016 and Niu2016 for comparison.
data_ref = ThermoTransportDataReader(
    gas_name="CH4", pressure_atm=1, source="Wu2016_Niu2016", ignore_missing_values=True
)
temperatures_ref = data_ref.temperature
density_ref = data_ref.density
enthalpy_ref = data_ref.enthalpy
viscosity_ref = data_ref.dynamic_viscosity
thermal_conductivity_ref = data_ref.thermal_conductivity
electrical_conductivity_ref = data_ref.electrical_conductivity

# %%
# Plot the results.
# -----------------
#
# Now we can visualise the properties by plotting them against temperature, to see how they vary.

fig, axs = plt.subplots(2, 3, figsize=(12, 8))

ax = axs[0, 0]
ax.set_title(r"$\mathregular{CH_4}$ plasma density at 1 atm")
ax.set_xlabel("T [K]")
ax.set_ylabel("$\\mathregular{\\rho [kg/m^3]}$")
ax.semilogy(data[:, 0], data[:, 1], "k", label="minplascalc")
# ax.semilogy(temperatures_ref, density_ref, "k--", label="Wu2016")

ax.legend()

ax = axs[0, 1]
ax.set_title(r"$\mathregular{CH_4}$ plasma heat capacity at 1 atm")
ax.set_xlabel("T [K]")
ax.set_ylabel(r"$\mathregular{C_P [J/(kg.K)]}$")
ax.plot(data[:, 0], data[:, 3], "k", label="minplascalc")
# ax.plot(temperatures_ref, cp_ref, "k--", label="Wu2016")
ax.legend()

ax = axs[0, 2]
ax.set_title(r"$\mathregular{CH_4}$ plasma enthalpy at 1 atm")
ax.set_xlabel("T [K]")
ax.set_ylabel(r"$\mathregular{H [J/kg]}$")
ax.plot(data[:, 0], data[:, 2], "k", label="minplascalc")
ax.scatter(temperatures_ref, enthalpy_ref, marker="x", color="red", label="Wu2016")
ax.legend()

ax = axs[1, 0]
ax.set_title(r"$\mathregular{CH_4}$ plasma viscosity")
ax.set_xlabel("T [K]")
ax.set_ylabel("$\\mathregular{\\mu [Pa.s]}$")
ax.plot(data[:, 0], data[:, 4], "k", label="minplascalc")
# ax.plot(temperatures_ref, viscosity_ref, "k--", label="Niu2016")
ax.legend()

ax = axs[1, 1]
ax.set_title(r"$\mathregular{CH_4}$ plasma thermal conductivity")
ax.set_xlabel("T [K]")
ax.set_ylabel("$\\mathregular{\\kappa [W/(m.K)]}$")
ax.plot(data[:, 0], data[:, 5], "k", label="minplascalc")
ax.scatter(
    temperatures_ref, thermal_conductivity_ref, marker="x", color="red", label="Niu2016"
)
ax.legend()

ax = axs[1, 2]
ax.set_title(r"$\mathregular{CH_4}$ plasma electrical conductivity")
ax.set_xlabel("T [K]")
ax.set_ylabel("$\\mathregular{\\sigma [S/m]}$")
ax.semilogy(data[:, 0], data[:, 6], "k", label="minplascalc")
ax.semilogy(
    temperatures_ref,
    electrical_conductivity_ref,
    marker="x",
    color="red",
    label="Niu2016",
)
ax.legend()


plt.tight_layout()

# %%