r"""
Extend NASA9 Temperature Limits for Kinetics
============================================

This example shows how to extend NASA9 thermodynamic data to higher temperatures
with a constant-:math:`C_p` continuation and inspect the resulting curves.

Workflow
--------

1. Load an input NASA9 thermo mechanism.
2. Fit an extended NASA9 representation with PMUTT.
3. Write a new YAML file with updated temperature ranges.
4. Plot one species to compare valid, native-extrapolated, and corrected trends.

.. tags:: kinetics, thermo, NASA9, cantera
"""  # noqa: D205, D400

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

from pathlib import Path

from rizer.kin.extend_NASA9_temperature_limits import Nasa9TemperatureLimitsExtender
from rizer.misc.utils import get_path_to_data

# %%
# Configure the extension task.
# -----------------------------

thermo_path = get_path_to_data("mechanisms", "thermo")
data_to_extend = [
    # "burcat_thermo_nasa9.yaml",
    "graphite.yaml",
    "grc_nasa9_Ar.yaml",
    "grc_nasa9_C_H.yaml",
    "grc_nasa9_N_O.yaml",
]
extenders: list[Nasa9TemperatureLimitsExtender] = []

for i, data in enumerate(data_to_extend):
    print(f"Available NASA9 thermo data: {thermo_path / data}")
    input_mechanism = thermo_path / data

    # Write the output near the input file for demonstration purposes.
    output_mechanism = Path(input_mechanism).with_name(
        f"extended_{Path(data).stem}.yaml"
    )

    extender = Nasa9TemperatureLimitsExtender(
        mechanism_to_extend=input_mechanism,
        new_mechanism_to_write=output_mechanism,
        max_temperature_extrapolated=200_000.0,
        number_of_points_fit=500,
    )

    extenders.append(extender)

# %%
# Generate the extended mechanism.
# --------------------------------

for extender in extenders:
    extender.write()
    print(f"Extended mechanism written to: {output_mechanism}")

# %%
# Plot one species thermo continuation.
# -------------------------------------

# For demonstration purposes, we plot the first species of the first mechanism.
extender = extenders[0]
species_to_plot = extender.species_list[0].name
print(f"Plotting thermo data for species: {species_to_plot}")
extender.plot_species_thermo_data(species_to_plot)

# %%