Hot spots example: PFR with exponential temperature decay and nucleation analysis.

from pathlib import Path

import cantera as ct
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import yaml

from bloc.chem import get_solid_carbon_mass_fraction_from_states
from bloc.chem import nobili as carbon_utils_nobili
from bloc.paths import get_mechanism_path
from bloc.plotting import SPECIES_COLORS, color_Cs
from bloc.reactors import (
    solve_fixed_time_heat_flux_pfr,
    solve_fixed_time_temperature_profile_pfr,
    switch_mechanism,
)

# NOTE: the previous "enthalpy transfer" strategy used
# `solve_fixed_time_enthalpy_profile_pfr` (still available in `bloc.reactors`).
# It is commented out below in favour of the temperature-difference-driven
# heat-flux coupling; re-add the import to restore it.

SCRIPT_DIR = Path(__file__).resolve().parent
RESULTS_BASE = SCRIPT_DIR / "results"
SKIP_SIMULATION = False


def _get_results_dir(inp):
    """Return results subfolder for this mechanism: results/{mech}/."""
    mech_string = inp["mech"].replace(".yaml", "")
    return RESULTS_BASE / mech_string


def _get_time_ms_str(inp):
    """Return time_exponential in ms as string for filenames, e.g. '10' for 0.01 s."""
    time_ms = int(round(inp["time_exponential"] * 1000))
    return f"{time_ms}ms"


def _save_states_csv(inp, states, results, out_dir, time_ms, suffix):
    """Save states and optional CRECK columns to results/{mech}/states_{time_ms}_{suffix}.csv."""
    fname = f"states_{time_ms}_{suffix}.csv"
    states_csv = out_dir / fname
    states.save(states_csv, overwrite=True)
    states_df = pd.read_csv(states_csv)
    states_df["residence_time"] = states.t
    states_df.to_csv(states_csv)
    if inp["mech"] == "CRECK_Nobili2024.yaml":
        states_df = pd.read_csv(states_csv)
        for cls, rates in results["mass_carbon_rates_by_class"].items():
            states_df[f"mass_carbon_rate_{cls}"] = rates
        for cls, rates in results["reaction_rates_by_class"].items():
            states_df[f"reaction_rate_{cls}"] = rates
        states_df.to_csv(states_csv)


def _load_inputs(yaml_path=None):
    """Load inputs from hot_spots_inputs.yaml and resolve mechanism path."""
    if yaml_path is None:
        yaml_path = SCRIPT_DIR / "hot_spots_inputs.yaml"
    with open(yaml_path, "r") as f:
        inp = yaml.safe_load(f)
    return inp


def run(inp):
    """Run the hot spots simulation: decay first, then increase from 300K with enthalpy from decay."""
    # Unpack: decay uses exponential temperature profile only
    temp_profile_exponential = inp["temp_profile_exponential"]
    equilibrate_before_run = inp["equilibrate_before_run"]
    initial_temperature = inp["initial_temperature"]
    final_temperature = inp["final_temperature"]
    time_exponential = inp["time_exponential"]
    # total_enthalpy_input = inp["total_enthalpy_input"]  # enthalpy-transfer strategy (disabled)
    dt = inp["dt"]
    n_time_exponential = inp["n_time_exponential"]
    out_dir = _get_results_dir(inp)
    out_dir.mkdir(parents=True, exist_ok=True)
    time_ms = _get_time_ms_str(inp)

    # 1) Exponential decay temperature profile only
    if temp_profile_exponential:
        print(
            f"Using exponential decay temperature profile with time_exponential = {time_exponential:.2e} s"
        )

        def temperature_profile_exponential(t, T_init, T_final, tau):
            return (T_init - T_final) * np.exp(-t / tau) + T_final

        temperature_profile_time = np.linspace(
            0, n_time_exponential * time_exponential, 100
        )
        temperature_profile = temperature_profile_exponential(
            temperature_profile_time,
            initial_temperature,
            final_temperature,
            time_exponential,
        )
        temperature_profile = np.vstack((temperature_profile_time, temperature_profile))

    if equilibrate_before_run:
        print(
            f"Equilibrating at initial temperature {initial_temperature} K with Caltech mechanism"
        )
        gas = ct.Solution(get_mechanism_path("Caltech.yaml"))
        gas.TPX = initial_temperature, inp["P_bar"], inp["inlet_composition"]
        print(f" Initial enthalpy: {gas.h * 1e-6:.2e} MJ/kg")
        gas.equilibrate("TP")
        print(f" Equilibrated enthalpy: {gas.h * 1e-6:.2e} MJ/kg")
        top5_idx = np.argsort(gas.X)[-5:][::-1]
        print(
            f"Major 5 species at initial temperature {initial_temperature} K (mole fraction):"
        )
        for idx in top5_idx:
            print(f"  {gas.species_names[idx]}: {gas.X[idx]:.6e}")
        print(
            f"Equilibration completed at temperature {gas.T} K, now switching to {inp['mech']} mechanism..."
        )
        gas = switch_mechanism(gas, inp["mech"], verbose=True)
    else:
        gas = ct.Solution(get_mechanism_path(inp["mech"]))
        gas.TPX = initial_temperature, inp["P_bar"], inp["inlet_composition"]
        print(f"Initial enthalpy: {gas.h * 1e-6:.2e} MJ/kg")

    inlet_composition = gas.X

    # 2) Run decay PFR (exponential decay only)
    print("---------------------- Decay PFR -------------------------")
    print(f"Using time step dt = {dt:.2e} s")
    results_decay = solve_fixed_time_temperature_profile_pfr(
        inlet_composition=inlet_composition,
        P_bar=inp["P_bar"],
        mechanism=inp["mech"],
        temperature_profile_time=temperature_profile,
        dt=dt,
        name="hot_spots_decay",
        verbose=2,
    )
    states_decay = results_decay["states"]
    _save_states_csv(inp, states_decay, results_decay, out_dir, time_ms, suffix="decay")

    # 3) Extract the decay time grid and temperature trajectory (the hot zone).
    times_decay = np.asarray(states_decay.t)
    T_hot_profile = np.asarray(states_decay.T)

    # ------------------------------------------------------------------ #
    # DISABLED STRATEGY A -- "enthalpy transfer".
    # The cold zone was fed the cumulative enthalpy lost by the hot zone,
    # scaled by a mass ratio derived from `total_enthalpy_input`. This
    # borrowed the cold-zone heating *shape* from the hot-zone cooling
    # curve and could over-drive the cold zone past the equilibrium
    # (mixing) temperature -- transferred enthalpy and cold mass were
    # independent knobs with no thermodynamic coupling. Kept for reference.
    #
    # h_profile_decay = states_decay.h
    # dh_cumulative = h_profile_decay[0] - h_profile_decay  # J/kg, >= 0 as decay cools
    # gas_inc = ct.Solution(get_mechanism_path(inp["mech"]))
    # gas_inc.TPX = 300.0, inp["P_bar"] * 1e5, inp["inlet_composition"]
    # h_inc_initial = gas_inc.h
    # delta_h_decay_initial = h_profile_decay[0] - h_inc_initial
    # mass_run_1 = 1.0
    # mass_run_2 = (
    #     delta_h_decay_initial / (total_enthalpy_input * 1e6) * mass_run_1 - mass_run_1
    # )
    # dh_cumulative = dh_cumulative * mass_run_1 / mass_run_2
    # h_profile_increasing = h_inc_initial + dh_cumulative
    # results_increasing = solve_fixed_time_enthalpy_profile_pfr(
    #     inlet_composition=inp["inlet_composition"],
    #     P_bar=inp["P_bar"],
    #     mechanism=inp["mech"],
    #     times=times_decay,
    #     h_profile=h_profile_increasing,
    #     dt=dt,
    #     initial_temperature=300.0,
    #     name="cold_spots_increasing",
    #     verbose=2,
    # )
    # ------------------------------------------------------------------ #

    # 4) STRATEGY B (current) -- temperature-difference-driven heat flux.
    #    The cold zone (from 300 K) is heated by a wall flux toward the hot-zone
    #    temperature with a first-order thermal lag tau:
    #        dT_cold/dt = (T_hot(t) - T_cold) / tau   (excluding chemistry).
    #    Self-limiting: the flux vanishes as T_cold -> T_hot, so the cold zone
    #    never exceeds the hot zone and the delivered enthalpy is an *output*.
    #    tau defaults to the decay time scale; override via the YAML if needed.
    thermal_time_constant = inp.get("cold_zone_time_constant", time_exponential)
    print("---------------------- Increase PFR (heat flux from hot zone) -----")
    results_increasing = solve_fixed_time_heat_flux_pfr(
        inlet_composition=inp["inlet_composition"],
        P_bar=inp["P_bar"],
        mechanism=inp["mech"],
        times=times_decay,
        T_env_profile=T_hot_profile,
        thermal_time_constant=thermal_time_constant,
        dt=dt,
        initial_temperature=300.0,
        name="cold_spots_increasing",
        verbose=2,
    )
    delivered = results_increasing.get("delivered_specific_enthalpy_J_per_kg")
    if delivered is not None:
        print(f"Delivered specific enthalpy to cold zone: {delivered * 1e-6:.2e} MJ/kg")
    states_increasing = results_increasing["states"]
    _save_states_csv(
        inp,
        states_increasing,
        results_increasing,
        out_dir,
        time_ms,
        suffix="increasing",
    )

    return {"decay": results_decay, "increasing": results_increasing}


def load_results(inp, mode="decay"):
    """Load results from results/{mech}/states_{time_ms}_{mode}.csv.

    mode : 'decay' or 'increasing'
    """
    out_dir = _get_results_dir(inp)
    time_ms = _get_time_ms_str(inp)
    fname = f"states_{time_ms}_{mode}.csv"
    states_csv = out_dir / fname
    if not states_csv.exists():
        raise FileNotFoundError(f"States file {states_csv} not found")
    gas = ct.Solution(get_mechanism_path(inp["mech"]))
    states = ct.SolutionArray(gas)
    states.read_csv(states_csv)
    results = {"states": states, "gas": gas}
    if inp["mech"] == "CRECK_Nobili2024.yaml":
        df = pd.read_csv(states_csv)
        for col in df.columns:
            if col.startswith("mass_carbon_rate_") or col.startswith("reaction_rate_"):
                setattr(states, col, np.asarray(df[col].values))
    return results


def _load_precursors_species(inp, results):
    """Load precursors species and names depending on the mechanism."""
    gas = results["gas"]

    # Caltech mechanism: A2, A2R5, A3, A3R5, A4
    if inp["mech"] == "Caltech.yaml":
        precursors_species = ["A2", "A2R5", "A3", "A3R5", "A4", "A4R5"]
    elif inp["mech"] == "CRECK_Nobili2024.yaml":
        species_groups = carbon_utils_nobili.generate_species_groups_nobili(gas)
        precursors_species = species_groups["liquid_particles"]
    else:
        precursors_species = []

    return precursors_species


def plot_results(inp, results, plot_suffix):
    """Plot results and save to results_{plot_suffix}_{time_ms}.png.

    plot_suffix : 'decreasing' for decay plot, 'increasing' for increase-from-300K plot.
    """
    gas = results["gas"]
    states = results["states"]
    precursor_species = _load_precursors_species(inp, results)

    out_dir = _get_results_dir(inp)
    out_dir.mkdir(parents=True, exist_ok=True)
    time_ms = _get_time_ms_str(inp)

    # Find nucleation peak if CRECK and rate data present
    nucleation_peak_time = None
    if inp["mech"] == "CRECK_Nobili2024.yaml":
        rate_attr = getattr(states, "mass_carbon_rate_Nucleation", None)
        if rate_attr is not None:
            nucleation_peak_time = states.t[np.argmax(rate_attr)]
    if inp["mech"] == "CRECK_Nobili2024.yaml":
        fig, axs = plt.subplots(2, 2, figsize=(12, 10), sharex=False)
        ax1, ax3, ax4, ax5 = axs.ravel()
    else:
        fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharex=False)
        ax1, ax2, ax3 = axs.ravel()

    if plot_suffix == "decreasing":
        fig.suptitle(
            f"Hot Spot temperature relaxation (decay)\n"
            f"Exponential decay (tau={time_ms}) simulated with {inp['mech'].replace('.yaml', '')}"
        )
    else:
        fig.suptitle(
            f"Cold Spots\n"
            f"From 300 K, heat flux from hot zone (tau={time_ms}), {inp['mech'].replace('.yaml', '')}"
        )

    # PLOT 1: Plot temperature profile
    ax1.plot(states.t * 1000, states.T, label="Temperature", c="tab:red")
    ax1.set_xlim(0, np.max(states.t) * 1000)
    ax1.set_ylabel("Temperature (K)")
    ax1.set_xlabel("Time (ms)")
    ax1.set_title("Temperature profile")
    if inp["mech"] == "CRECK_Nobili2024.yaml" and nucleation_peak_time is not None:
        ax1.axvline(
            nucleation_peak_time * 1000,
            color="tab:blue",
            linestyle="--",
            label="Nucleation peak",
        )
    ax1.legend(fontsize="small", loc="best", ncol=2)

    # PLOT 2: Plot enthalpy profile + mass heat capacity on right axis (Cantera: cp_mass)
    if inp["mech"] != "CRECK_Nobili2024.yaml":
        (line_h,) = ax2.plot(
            states.t * 1000, states.h * 1e-6, label="Enthalpy", c="tab:purple"
        )
        ax2.set_xlim(0, np.max(states.t) * 1000)
        ax2.set_ylabel("Enthalpy (MJ/kg)")
        ax2.set_xlabel("Time (ms)")
        ax2.set_title("Enthalpy profile")
        ax2_cp = ax2.twinx()
        (line_cp,) = ax2_cp.plot(
            states.t * 1000,
            states.cp_mass,
            label=r"$c_p$",
            c="tab:green",
        )
        ax2_cp.set_ylabel(r"Heat capacity $c_p$ (J/(kg·K))")
        ax2_cp.set_ylim(2000, 6300)
        ax2_cp.tick_params(axis="y", labelcolor="tab:green")
        if inp["mech"] == "CRECK_Nobili2024.yaml" and nucleation_peak_time is not None:
            ax2.axvline(
                nucleation_peak_time * 1000,
                color="tab:blue",
                linestyle="--",
                label="Nucleation peak",
            )
        ax2.legend(
            [line_h, line_cp],
            ["Enthalpy", r"$c_p$"],
            fontsize="small",
            loc="best",
            ncol=2,
        )

    # PLOT 3: Fixed species list with colors from bloc/plot.py
    _SPECIES_TO_PLOT = ["CH4", "H2", "C2H2", "C6H6"]
    species_names = gas.species_names
    for sp in _SPECIES_TO_PLOT:
        if sp not in species_names:
            continue
        ax3.plot(
            states.t * 1000,
            states(sp).Y,
            label=sp,
            color=SPECIES_COLORS[sp],
        )
    mass_fraction_precursors = np.zeros(len(states.t))
    for sp in precursor_species:
        if sp in species_names:
            mass_fraction_precursors += states(sp).Y[:, 0]
    ax3.plot(
        states.t * 1000,
        mass_fraction_precursors,
        label="Liquid particles"
        if inp["mech"] == "CRECK_Nobili2024.yaml"
        else "Precursors",
        linestyle="--",
    )
    Y_solid = get_solid_carbon_mass_fraction_from_states(states)
    ax3.plot(
        states.t * 1000,
        Y_solid,
        label="Solid carbon",
        color=color_Cs,
        linewidth=2,
    )
    ax3.set_ylabel("Mass fraction")
    ax3.set_yscale("log")
    ax3.set_ylim(1e-6, 1)
    ax3.set_xlim(0, np.max(states.t) * 1000)
    ax3.set_xlabel("Time (ms)")
    ax3.set_title("Species concentrations")
    ax3.legend(fontsize="small", loc="best", ncol=2)

    # PLOT 4: Plot mass carbon rates by class
    if inp["mech"] == "CRECK_Nobili2024.yaml":
        rate_attr = getattr(states, "mass_carbon_rate_Nucleation", None)
        if rate_attr is not None:
            ax4.plot(
                states.t * 1000, states.mass_carbon_rate_Nucleation, label="Nucleation"
            )
            ax4.plot(
                states.t * 1000,
                states.mass_carbon_rate_SurfaceGrowthPAH,
                label="SurfaceGrowthPAH",
            )
            ax4.plot(
                states.t * 1000,
                states.mass_carbon_rate_SurfaceGrowthRadicals,
                label="SurfaceGrowthRadicals",
            )
            ax4.plot(states.t * 1000, states.mass_carbon_rate_HACA, label="HACA")
            ax4.plot(
                states.t * 1000,
                states.mass_carbon_rate_Coalescence,
                label="Coalescence",
            )
            if nucleation_peak_time is not None:
                ax4.axvline(
                    nucleation_peak_time * 1000,
                    color="black",
                    linestyle="--",
                    label="Nucleation peak",
                )
        ax4.set_xlabel("Time (ms)")
        ax4.set_ylabel("Mass carbon rate (kg/m^3/s)")
        ax4.set_yscale("log")
        ax4.set_ylim(1e-6, 1)
        ax4.set_xlim(0, np.max(states.t) * 1000)
        ax4.set_title("Mass carbon rates by class")
        ax4.legend(fontsize="small", loc="best", ncol=2)

    # PLOT 5: Histogram of nucleated mass vs temperature (mass per time step binned by T)
    if inp["mech"] == "CRECK_Nobili2024.yaml":
        rate_attr = getattr(states, "mass_carbon_rate_Nucleation", None)
        if rate_attr is not None:
            dt_arr = np.diff(np.concatenate(([0], states.t)))
            mass_per_step = np.array(states.mass_carbon_rate_Nucleation) * dt_arr
            T_history = np.asarray(states.T)
            T_min, T_max = T_history.min(), T_history.max()
            n_bins = max(30, int((T_max - T_min) / 50))
            bin_edges = np.linspace(T_min, T_max, n_bins + 1)
            mass_hist, _ = np.histogram(
                T_history, bins=bin_edges, weights=mass_per_step
            )
            bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2
            bin_width = np.diff(bin_edges)
            ax5.bar(
                bin_centers,
                mass_hist,
                width=bin_width * 0.9,
                align="center",
                color="tab:blue",
                edgecolor="gray",
            )
            ax5.set_xlabel("Temperature (K)")
            ax5.set_ylabel("Nucleated mass (kg/m^3)")
            ax5.set_title("Nucleation mass distribution vs temperature")
            ax5.set_xlim(T_min, T_max)

    fig.tight_layout()
    fig.savefig(out_dir / f"results_{plot_suffix}_{time_ms}.png", bbox_inches="tight")
    print(f"Saved figure to {out_dir / f'results_{plot_suffix}_{time_ms}.png'}")


if __name__ == "__main__":
    print("---------------- Running hot spots simulation ----------------")
    inp = _load_inputs()
    if SKIP_SIMULATION:
        time_ms = _get_time_ms_str(inp)
        print(
            f"Skipping simulation and loading results from results/.../states_{time_ms}_*.csv"
        )
        results_decay = load_results(inp, mode="decay")
        results_increasing = load_results(inp, mode="increasing")
    else:
        both = run(inp)
        results_decay = both["decay"]
        results_increasing = both["increasing"]
    plot_results(inp, results_decay, plot_suffix="decreasing")
    plot_results(inp, results_increasing, plot_suffix="increasing")