#include "gridfire/engine/views/engine_multiscale.h"
#include "gridfire/exceptions/error_engine.h"
#include "gridfire/engine/procedures/priming.h"
#include "gridfire/utils/sundials.h"
#include "gridfire/utils/logging.h"

#include "gridfire/engine/scratchpads/blob.h"
#include "gridfire/engine/scratchpads/utils.h"
#include "gridfire/engine/scratchpads/engine_multiscale_scratchpad.h"

#include <stdexcept>
#include <vector>
#include <ranges>
#include <unordered_map>
#include <unordered_set>
#include <fstream>

#include <queue>

#include <algorithm>
#include <numeric>

#include "fourdst/atomic/species.h"
#include "quill/LogMacros.h"
#include "quill/Logger.h"

#include "kinsol/kinsol.h"
#include "sundials/sundials_context.h"
#include "sunmatrix/sunmatrix_dense.h"
#include "sunlinsol/sunlinsol_dense.h"

#include  "xxhash64.h"
#include "fourdst/composition/exceptions/exceptions_composition.h"
#include "fourdst/composition/utils/composition_hash.h"

namespace {
    using namespace fourdst::atomic;


    std::vector<std::vector<Species>> findConnectedComponentsBFS(
        const std::unordered_map<Species, std::vector<Species>>& graph,
        const std::vector<Species>& nodes
    ) {
        std::vector<std::vector<Species>> components;
        std::unordered_set<Species> visited;

        for (const Species& start_node : nodes) {
            if (!visited.contains(start_node)) {
                std::vector<Species> current_component;
                std::queue<Species> q;

                q.push(start_node);
                visited.insert(start_node);

                while (!q.empty()) {
                    Species u = q.front();
                    q.pop();
                    current_component.push_back(u);

                    if (graph.contains(u)) {
                        for (const auto& v : graph.at(u)) {
                            if (!visited.contains(v)) {
                                visited.insert(v);
                                q.push(v);
                            }
                        }
                    }
                }
                components.push_back(current_component);
            }
        }
        return components;
    }

    std::vector<std::vector<Species>> findConnectedComponentsBFS(
        const std::unordered_map<Species, std::set<Species>>& graph,
        const std::vector<Species>& nodes
    ) {
        std::unordered_map<Species, std::vector<Species>> adjList;
        for (const auto& [u, neighbors] : graph) {
            adjList[u] = std::vector<Species>(neighbors.begin(), neighbors.end());
        }
        return findConnectedComponentsBFS(adjList, nodes);
    }

    struct SpeciesSetIntersection {
        const Species species;
        std::size_t count;
    };

    std::expected<SpeciesSetIntersection, std::string> get_intersection_info (
        const std::unordered_set<Species>& setA,
        const std::unordered_set<Species>& setB
    ) {
        // Iterate over the smaller of the two
        auto* outerSet = &setA;
        auto* innerSet = &setB;
        if (setA.size() > setB.size()) {
            outerSet = &setB;
            innerSet = &setA;
        }

        std::size_t matchCount = 0;
        const Species* firstMatch = nullptr;

        for (const Species& sp : *outerSet) {
            if (innerSet->contains(sp)) {
                if (matchCount == 0) {
                    firstMatch = &sp;
                }
                ++matchCount;
                if (matchCount > 1) {
                    break;
                }
            }
        }
        if (!firstMatch) {
            // No matches found
            return std::unexpected{"Intersection is empty"};
        }
        if (matchCount == 0) {
            // No matches found
            return std::unexpected{"No intersection found"};
        }

        // Return the first match and the count of matches
        return SpeciesSetIntersection{*firstMatch, matchCount};

    }

    bool has_distinct_reactant_and_product_species (
        const std::unordered_set<Species>& poolSpecies,
        const std::unordered_set<Species>& reactants,
        const std::unordered_set<Species>& products
    ) {
        const auto reactant_result = get_intersection_info(poolSpecies, reactants);
        if (!reactant_result) {
            return false; // No reactants found
        }
        const auto [reactantSample, reactantCount] = reactant_result.value();

        const auto product_result = get_intersection_info(poolSpecies, products);
        if (!product_result) {
            return false; // No products found
        }

        const auto [productSample, productCount] = product_result.value();

        // If either side has ≥2 distinct matches, we can always pick
        // one from each that differ.
        if (reactantCount > 1 || productCount > 1) {
            return true;
        }

        // Exactly one match on each side → they must differ
        return reactantSample != productSample;
    }

    struct DisjointSet {
        std::vector<size_t> parent;
        explicit DisjointSet(const size_t n) {
            parent.resize(n);
            std::iota(parent.begin(), parent.end(), 0);
        }

        size_t find(const size_t i) { // NOLINT(*-no-recursion)
            if (parent.at(i) == i) return i;
            return parent.at(i) = find(parent.at(i)); // Path compression
        }

        void unite(const size_t i, const size_t j) {
            const size_t root_i = find(i);
            const size_t root_j = find(j);
            if (root_i != root_j) {
                parent.at(root_i) = root_j;
            }

        }
    };
}

namespace gridfire::engine {
    using fourdst::atomic::Species;

    MultiscalePartitioningEngineView::MultiscalePartitioningEngineView(
        DynamicEngine& baseEngine
    ) : m_baseEngine(baseEngine) {
    }

    const std::vector<Species> & MultiscalePartitioningEngineView::getNetworkSpecies(
        scratch::StateBlob& ctx
    ) const {
        return m_baseEngine.getNetworkSpecies(ctx);
    }

    std::expected<StepDerivatives<double>, EngineStatus> MultiscalePartitioningEngineView::calculateRHSAndEnergy(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho,
        bool trust
    ) const {
        LOG_TRACE_L2(m_logger, "Calculating RHS and Energy in MultiscalePartitioningEngineView at T9 = {}, rho = {}.", T9, rho);
        LOG_TRACE_L2(m_logger, "Input composition is {}", [&comp]() -> std::string {
            std::stringstream ss;
            size_t i = 0;
            for (const auto& [species, abundance] : comp) {
                ss << species.name() << ": " << abundance;
                if (i < comp.size() - 1) {
                    ss << ", ";
                }
                i++;
            }
            return ss.str();
        }());
        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, trust);
        LOG_TRACE_L2(m_logger, "Equilibrated composition prior to calling base engine is {}", [&qseComposition, &comp]() -> std::string {
            std::stringstream ss;
            size_t i = 0;
            for (const auto& [species, abundance] : qseComposition) {
                ss << species.name() << ": " << abundance;
                if (comp.contains(species)) {
                    ss << " (input: " << comp.getMolarAbundance(species) << ")";
                }
                if (i < qseComposition.size() - 1) {
                    ss << ", ";
                }
                i++;
            }
            return ss.str();
        }());

        const auto result = m_baseEngine.calculateRHSAndEnergy(ctx, qseComposition, T9, rho, false);
        LOG_TRACE_L2(m_logger, "Base engine calculation of RHS and Energy complete.");

        if (!result) {
            LOG_TRACE_L2(m_logger, "Base engine returned stale error during RHS and Energy calculation.");
            return std::unexpected{result.error()};
        }

        auto deriv = result.value();
        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);

        LOG_TRACE_L2(m_logger, "Zeroing out algebraic species derivatives.");
        for (const auto& species : state->algebraic_species) {
            deriv.dydt[species] = 0.0; // Fix the algebraic species to the equilibrium abundances we calculate.
        }
        LOG_TRACE_L2(m_logger, "Done Zeroing out algebraic species derivatives.");
        return deriv;
    }

    EnergyDerivatives MultiscalePartitioningEngineView::calculateEpsDerivatives(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho
    ) const {
        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);
        return m_baseEngine.calculateEpsDerivatives(ctx, qseComposition, T9, rho);
    }

    NetworkJacobian MultiscalePartitioningEngineView::generateJacobianMatrix(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho
    ) const {
        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);
        return m_baseEngine.generateJacobianMatrix(ctx, qseComposition, T9, rho, state->dynamic_species);
    }

    NetworkJacobian MultiscalePartitioningEngineView::generateJacobianMatrix(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho,
        const std::vector<Species> &activeSpecies
    ) const {
        bool activeSpeciesIsSubset = true;
        for (const auto& species : activeSpecies) {
            if (!involvesSpecies(ctx, species)) activeSpeciesIsSubset = false;
        }
        if (!activeSpeciesIsSubset) {
            std::string msg = std::format(
                "Active species set contains species ({}) not present in network partition. Cannot generate jacobian matrix due to this.",
                [&]() -> std::string {
                    std::stringstream ss;
                    for (const auto& species : activeSpecies) {
                        if (!involvesSpecies(ctx, species)) {
                            ss << species << " ";
                        }
                    }
                    return ss.str();
                }()
            );
            LOG_CRITICAL(m_logger, "{}", msg);
            throw std::runtime_error(msg);
        }

        std::vector<Species> dynamicActiveSpeciesIntersection;
        for (const auto& species : activeSpecies) {
            if (involvesSpeciesInDynamic(ctx, species)) {
                dynamicActiveSpeciesIntersection.push_back(species);
            }
        }

        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);

        return m_baseEngine.generateJacobianMatrix(ctx, qseComposition, T9, rho, dynamicActiveSpeciesIntersection);
    }

    NetworkJacobian MultiscalePartitioningEngineView::generateJacobianMatrix(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho,
        const SparsityPattern &sparsityPattern
    ) const {
        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);
        return m_baseEngine.generateJacobianMatrix(ctx, qseComposition, T9, rho, sparsityPattern);
    }

    double MultiscalePartitioningEngineView::calculateMolarReactionFlow(
        scratch::StateBlob& ctx,
        const reaction::Reaction &reaction,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho
    ) const {
        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);

        return m_baseEngine.calculateMolarReactionFlow(ctx, reaction, qseComposition, T9, rho);
    }

    const reaction::ReactionSet & MultiscalePartitioningEngineView::getNetworkReactions(
        scratch::StateBlob& ctx
    ) const {
        return m_baseEngine.getNetworkReactions(ctx);
    }

    std::expected<std::unordered_map<Species, double>, EngineStatus> MultiscalePartitioningEngineView::getSpeciesTimescales(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho
    ) const {
        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);
        const auto result  = m_baseEngine.getSpeciesTimescales(ctx, qseComposition, T9, rho);
        if (!result) {
            return std::unexpected{result.error()};
        }

        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        std::unordered_map<Species, double> speciesTimescales = result.value();
        for (const auto& algebraicSpecies : state->algebraic_species) {
            speciesTimescales[algebraicSpecies] = std::numeric_limits<double>::infinity(); // Algebraic species have infinite timescales.
        }
        return speciesTimescales;
    }

    std::expected<std::unordered_map<Species, double>, EngineStatus> MultiscalePartitioningEngineView::getSpeciesDestructionTimescales(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho
    ) const {
        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);
        const auto result = m_baseEngine.getSpeciesDestructionTimescales(ctx, qseComposition, T9, rho);
        if (!result) {
            return std::unexpected{result.error()};
        }
        std::unordered_map<Species, double> speciesDestructionTimescales = result.value();
        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        for (const auto& algebraicSpecies : state->algebraic_species) {
            speciesDestructionTimescales[algebraicSpecies] = std::numeric_limits<double>::infinity(); // Algebraic species have infinite destruction timescales.
        }
        return speciesDestructionTimescales;
    }

    fourdst::composition::Composition MultiscalePartitioningEngineView::project(
        scratch::StateBlob& ctx,
        const NetIn &netIn
    ) const {
        const fourdst::composition::Composition baseUpdatedComposition = m_baseEngine.project(ctx, netIn);

        NetIn baseUpdatedNetIn = netIn;
        baseUpdatedNetIn.composition = baseUpdatedComposition;
        fourdst::composition::Composition equilibratedComposition = partitionNetwork(ctx, baseUpdatedNetIn);

        auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        state->composition_cache.clear();

        return equilibratedComposition;
    }

    screening::ScreeningType MultiscalePartitioningEngineView::getScreeningModel(
        scratch::StateBlob& ctx
    ) const {
        return m_baseEngine.getScreeningModel(ctx);
    }

    const DynamicEngine & MultiscalePartitioningEngineView::getBaseEngine() const {
        return m_baseEngine;
    }

    std::vector<std::vector<Species>> MultiscalePartitioningEngineView::analyzeTimescalePoolConnectivity(
        scratch::StateBlob& ctx,
        const std::vector<std::vector<Species>> &timescale_pools,
        const fourdst::composition::Composition &comp,
        double T9,
        double rho
    ) const {
        std::vector<std::vector<Species>> final_connected_pools;

        for (const auto& pool : timescale_pools) {
            if (pool.empty()) {
                continue; // Skip empty pools
            }

            // For each timescale pool, we need to analyze connectivity.
            auto connectivity_graph = buildConnectivityGraph(ctx, pool, comp, T9, rho);
            auto components = findConnectedComponentsBFS(connectivity_graph, pool);
            final_connected_pools.insert(final_connected_pools.end(), components.begin(), components.end());
        }
        return final_connected_pools;

    }

    std::vector<MultiscalePartitioningEngineView::QSEGroup> MultiscalePartitioningEngineView::pruneValidatedGroups(
        scratch::StateBlob& ctx,
        const std::vector<QSEGroup> &groups,
        const std::vector<reaction::ReactionSet> &groupReactions,
        const fourdst::composition::Composition &comp,
        const double T9,
        const double rho
    ) const {
        const auto result = m_baseEngine.getSpeciesTimescales(ctx, comp, T9, rho);
        if (!result) {
            throw std::runtime_error("Base engine returned stale error during pruneValidatedGroups timescale retrieval.");
        }
        const std::unordered_map<Species, double>& speciesTimescales = result.value();
        std::vector<QSEGroup> newGroups;
        for (const auto &[group, reactions] : std::views::zip(groups, groupReactions)) {
            if (reactions.empty()) { // If a QSE group has gotten here it should have reactions associated with it. If it doesn't that is a serious error.
                LOG_CRITICAL(m_logger, "No reactions specified for QSE group {} during pruning analysis.", group.toString(false));
                throw std::runtime_error("No reactions specified for QSE group " + group.toString(false) + " during pruneValidatedGroups flux analysis.");
            }
            LOG_TRACE_L2(m_logger, "Attempting pruning of QSE Group {:40} (reactions: {}).", group.toString(false), utils::iterable_to_delimited_string(reactions, ", ", [](const auto& r) { return r->id(); }));
            std::unordered_map<size_t, double> reactionFluxes;
            std::unordered_map<size_t, const reaction::Reaction&> reactionLookup;
            reactionFluxes.reserve(reactions.size());

            double mean_molar_abundance = 0;
            for (const auto& species : group.algebraic_species) {
                mean_molar_abundance += comp.getMolarAbundance(species);
            }
            mean_molar_abundance /= static_cast<double>(group.algebraic_species.size());
            { // Safety Valve to ensure valid log scaling
                if (mean_molar_abundance <= 0) {
                    LOG_CRITICAL(m_logger, "Non-positive mean molar abundance {} calculated for QSE group during pruning analysis.", mean_molar_abundance);
                    throw std::runtime_error("Non-positive mean molar abundance calculated during pruneValidatedGroups flux analysis.");
                }
            }

            for (const auto& reaction : reactions) {
                const double flux = m_baseEngine.calculateMolarReactionFlow(ctx, *reaction, comp, T9, rho);
                size_t hash = reaction->hash(0);
                if (reactionFluxes.contains(hash)) {
                    throw std::runtime_error("Duplicate reaction hash found during pruneValidatedGroups flux analysis.");
                }

                { // Safety Valve to ensure valid log scaling
                    if (flux <= 0) {
                        LOG_CRITICAL(m_logger, "Non-positive flux {} calculated for reaction {} during pruning analysis.", flux, reaction->id());
                        throw std::runtime_error("Non-positive flux calculated during pruneValidatedGroups flux analysis.");
                    }
                }
                double lAbundanceNormalizedFlux = std::log(flux/mean_molar_abundance);
                reactionFluxes.emplace(hash, lAbundanceNormalizedFlux);

                assert(!std::isnan(lAbundanceNormalizedFlux) && !std::isinf(lAbundanceNormalizedFlux) && "Invalid log abundance normalized flux calculated during pruneValidatedGroups flux analysis.");

                reactionLookup.emplace(hash, *reaction);
            }

            std::vector<size_t> sorted_reactions_based_on_flow;
            for (const auto &hash: reactionFluxes | std::views::keys) {
                sorted_reactions_based_on_flow.push_back(hash);
            }

            std::ranges::sort(sorted_reactions_based_on_flow, [&reactionFluxes](const size_t a, const size_t b) {
                return std::abs(reactionFluxes.at(a)) < std::abs(reactionFluxes.at(b));
            });

            std::unordered_map<size_t, double> pruned_reaction_fluxes;
            for (const auto& [hash, normalizedFlux] : reactionFluxes) {
                if (normalizedFlux > -30) { // TODO: replace -30 with some more physically motivated value
                    pruned_reaction_fluxes.emplace(hash, normalizedFlux);
                    LOG_TRACE_L2(m_logger, "Retaining reaction {:15} with log(mean abundance normalized flux) {:10.4E} during pruning.", reactionLookup.at(hash).id(), normalizedFlux);
                } else {
                    LOG_TRACE_L2(m_logger, "Pruning   reaction {:15} with log(mean abundance normalized flux) {:10.4E} during pruning.", reactionLookup.at(hash).id(), normalizedFlux);
                }
            }

            std::set<Species> reachableAlgebraicSpecies;
            std::set<Species> reachableSeedSpecies;

            std::unordered_map<Species, std::set<Species>> connectivity_graph;

            for (const auto& reactionHash : pruned_reaction_fluxes | std::views::keys) {
                const auto& reaction = reactionLookup.at(reactionHash);

                for (const auto& reactant : reaction.reactants()) {
                    if (group.algebraic_species.contains(reactant)) {
                        reachableAlgebraicSpecies.insert(reactant);
                    } else if (group.seed_species.contains(reactant)) {
                        reachableSeedSpecies.insert(reactant);
                    }

                    if (!connectivity_graph.contains(reactant)) {
                        connectivity_graph.emplace(reactant, std::set<Species>{});
                    }
                    for (const auto& product : reaction.products()) {
                        connectivity_graph.at(reactant).insert(product);
                    }
                }

                for (const auto& product : reaction.products()) {
                    if (group.algebraic_species.contains(product)) {
                        reachableAlgebraicSpecies.insert(product);
                    } else if (group.seed_species.contains(product)) {
                        reachableSeedSpecies.insert(product);
                    }
                }

            }

            LOG_TRACE_L2(
                m_logger,
                "{}",
                [&group, &reachableAlgebraicSpecies, &reachableSeedSpecies]() -> std::string {
                    std::stringstream ss;
                    ss << "Pruned QSE Group. Group Started with Algebraic Species: {";
                    int i = 0;
                    for (const auto& species : group.algebraic_species) {
                        ss << species.name();
                        if (i < group.algebraic_species.size() - 1) {
                            ss << ", ";
                        }
                        i++;
                    }
                    ss << "} and Seed Species: {";
                    i = 0;
                    for (const auto& species : group.seed_species) {
                        ss << species.name();
                        if (i < group.seed_species.size() - 1) {
                            ss << ", ";
                        }
                        i++;
                    }
                    ss << "}. After pruning, reachable Algebraic Species: {";
                    i = 0;
                    for (const auto& species : reachableAlgebraicSpecies) {
                        ss << species.name();
                        if (i < reachableAlgebraicSpecies.size() - 1) {
                            ss << ", ";
                        }
                        i++;
                    }
                    ss << "} and reachable Seed Species: {";
                    i = 0;
                    for (const auto& species : reachableSeedSpecies) {
                        ss << species.name();
                        if (i < reachableSeedSpecies.size() - 1) {
                            ss << ", ";
                        }
                        i++;
                    }
                    ss << "}.";
                    return ss.str();
                }()
            );

            std::vector<std::vector<Species>> connected_components = findConnectedComponentsBFS(
                connectivity_graph,
                std::vector<Species>(
                    reachableAlgebraicSpecies.begin(),
                    reachableAlgebraicSpecies.end()
                )
            );

            for (const auto& subgraph: connected_components) {
                QSEGroup g;
                for (const auto& species: subgraph) {
                    if (reachableAlgebraicSpecies.contains(species)) {
                        g.algebraic_species.insert(species);
                    } else if (reachableSeedSpecies.contains(species)) {
                        g.seed_species.insert(species);
                    }
                }
                if (!g.seed_species.empty() && !g.algebraic_species.empty()) {
                    double meanTimescale = 0;
                    for (const auto &species : g.algebraic_species) {
                        meanTimescale += speciesTimescales.at(species);
                    }
                    meanTimescale /= static_cast<double>(g.algebraic_species.size());
                    g.mean_timescale = meanTimescale;
                    newGroups.push_back(g);
                }
            }
        }
        return newGroups;
    }

    std::vector<MultiscalePartitioningEngineView::QSEGroup> MultiscalePartitioningEngineView::merge_coupled_groups(
        scratch::StateBlob& ctx,
        const std::vector<QSEGroup> &groups,
        const std::vector<reaction::ReactionSet> &groupReactions
    ) {
        if (groups.empty()) { return {}; }
        DisjointSet dsu(groups.size());

        std::unordered_map<std::string_view, size_t> reaction_owner;

        for (size_t group_idx = 0; group_idx < groups.size(); ++group_idx) {
            for (const auto& reaction : groupReactions[group_idx]) {
                const std::string_view rxn_id = reaction->id();
                if (reaction_owner.contains(rxn_id)) {
                    dsu.unite(group_idx, reaction_owner[rxn_id]);
                } else {
                    reaction_owner[rxn_id] = group_idx;
                }
            }
        }

        std::unordered_map<size_t, std::pair<QSEGroup, size_t>> merged_groups_map;
        for (size_t i = 0; i < groups.size(); ++i) {
            size_t root = dsu.find(i);

            if (!merged_groups_map.contains(root)) {
                merged_groups_map[root].first = QSEGroup();
                merged_groups_map.at(root).first.mean_timescale = 0.0;
                merged_groups_map.at(root).first.is_in_equilibrium = true;
            }

            QSEGroup& target_group = merged_groups_map.at(root).first;
            const QSEGroup& source_group = groups[i];

            target_group.seed_species.insert(source_group.seed_species.begin(), source_group.seed_species.end());
            target_group.algebraic_species.insert(source_group.algebraic_species.begin(), source_group.algebraic_species.end());

            target_group.mean_timescale += source_group.mean_timescale;
            merged_groups_map.at(root).second += 1;
        }

        std::vector<QSEGroup> final_merged_groups;
        for (auto &groupInfo: merged_groups_map | std::views::values) {
            QSEGroup& group = groupInfo.first;
            for (const auto& alg_species : group.algebraic_species) {
                group.seed_species.erase(alg_species); // A species cannot be both seed and algebraic. Algebraic takes precedence.
            }
            group.mean_timescale /= static_cast<double>(groupInfo.second);

            final_merged_groups.push_back(std::move(group));
        }

        return final_merged_groups;
    }

    fourdst::composition::Composition MultiscalePartitioningEngineView::partitionNetwork(
        scratch::StateBlob& ctx,
        const NetIn &netIn
    ) const {
        auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        // --- Step 0. Prime the network ---
        const PrimingReport primingReport = m_baseEngine.primeEngine(ctx, netIn);
        const fourdst::composition::Composition& comp = primingReport.primedComposition;
        const double T9 = netIn.temperature / 1e9;
        const double rho = netIn.density;

        // --- Step 0.5 Clear previous state ---
        LOG_TRACE_L1(m_logger, "Partitioning network...");
        LOG_TRACE_L1(m_logger, "Clearing previous state...");
        state->qse_groups.clear();
        state->qse_solvers.clear();
        state->dynamic_species.clear();
        state->algebraic_species.clear();
        state->composition_cache.clear(); // We need to clear the cache now cause the same comp, temp, and density may result in a different value

        // --- Step 1. Identify distinct timescale regions ---
        LOG_TRACE_L1(m_logger, "Identifying fast reactions...");
        const std::vector<std::vector<Species>> timescale_pools = partitionByTimescale(ctx, comp, T9, rho);
        LOG_TRACE_L1(m_logger, "Found {} timescale pools.", timescale_pools.size());

        // --- Step 2. Select the mean slowest pool as the base dynamical group ---
        LOG_TRACE_L1(m_logger, "Identifying mean slowest pool...");
        const size_t mean_slowest_pool_index = identifyMeanSlowestPool(ctx, timescale_pools, comp, T9, rho);
        LOG_TRACE_L1(m_logger, "Mean slowest pool index: {}", mean_slowest_pool_index);

        // --- Step 3. Push the slowest pool into the dynamic species list ---
        for (const auto& slowSpecies : timescale_pools[mean_slowest_pool_index]) {
            state->dynamic_species.push_back(slowSpecies);
        }

        // --- Step 4. Pack Candidate QSE Groups ---
        std::vector<std::vector<Species>> candidate_pools;
        for (size_t i = 0; i < timescale_pools.size(); ++i) {
            if (i == mean_slowest_pool_index) continue; // Skip the slowest pool
            LOG_TRACE_L1(m_logger, "Group {} with {} species identified for potential QSE.", i, timescale_pools[i].size());
            candidate_pools.push_back(timescale_pools[i]);
        }

        LOG_TRACE_L1(m_logger, "Preforming connectivity analysis on timescale pools...");
        const std::vector<std::vector<Species>> connected_pools = analyzeTimescalePoolConnectivity(ctx, candidate_pools, comp, T9, rho);
        LOG_TRACE_L1(m_logger, "Found {} connected pools (compared to {} timescale pools) for QSE analysis.", connected_pools.size(), timescale_pools.size());


        // --- Step 5. Identify potential seed species for each candidate pool ---
        LOG_TRACE_L1(m_logger, "Identifying potential seed species for candidate pools...");
        const std::vector<QSEGroup> candidate_groups = constructCandidateGroups(ctx, connected_pools, comp, T9, rho);
        LOG_TRACE_L1(m_logger, "Found {} candidate QSE groups for further analysis ({})", candidate_groups.size(), utils::iterable_to_delimited_string(candidate_groups));

        LOG_TRACE_L1(m_logger, "Validating candidate groups with flux analysis...");
        const auto [validated_groups, invalidate_groups, validated_group_reactions] = validateGroupsWithFluxAnalysis(ctx, candidate_groups, comp, T9, rho);
        LOG_TRACE_L1(m_logger, "Validated {} group(s) QSE groups. {}", validated_groups.size(), utils::iterable_to_delimited_string(validated_groups));


        LOG_TRACE_L1(m_logger, "Pruning groups based on log abundance-normalized flux analysis...");
        const std::vector<QSEGroup> prunedGroups = pruneValidatedGroups(ctx, validated_groups, validated_group_reactions, comp, T9, rho);
        LOG_TRACE_L1(m_logger, "After Pruning remaining groups are: {}", utils::iterable_to_delimited_string(prunedGroups));

        LOG_TRACE_L1(m_logger, "Re-validating pruned groups with flux analysis...");
        auto [pruned_validated_groups, _, pruned_validated_reactions] = validateGroupsWithFluxAnalysis(ctx, prunedGroups, comp, T9, rho);
        LOG_TRACE_L1(m_logger, "After re-validation, {} QSE groups remain. ({})",pruned_validated_groups.size(), utils::iterable_to_delimited_string(pruned_validated_groups));

        LOG_TRACE_L1(m_logger, "Merging coupled QSE groups...");
        const std::vector<QSEGroup> merged_groups = merge_coupled_groups(ctx, pruned_validated_groups, pruned_validated_reactions);
        LOG_TRACE_L1(m_logger, "After merging, {} QSE groups remain. ({})", merged_groups.size(), utils::iterable_to_delimited_string(merged_groups));

        state->qse_groups = pruned_validated_groups;

        LOG_TRACE_L1(m_logger, "Pushing all identified algebraic species into algebraic set...");
        for (const auto& group : state->qse_groups) {
            // Add algebraic species to the algebraic set
            for (const auto& species : group.algebraic_species) {
                if (std::ranges::find(state->algebraic_species, species) == state->algebraic_species.end()) {
                    state->algebraic_species.push_back(species);
                }
            }
        }
        LOG_TRACE_L1(m_logger, "Algebraic species identified: {}", utils::iterable_to_delimited_string(state->algebraic_species));

        LOG_INFO(
            m_logger,
            "Partitioning complete. Found {} dynamic species, {} algebraic (QSE) species ({}) spread over {} QSE group{}.",
            state->dynamic_species.size(),
            state->algebraic_species.size(),
            utils::iterable_to_delimited_string(state->algebraic_species),
            state->qse_groups.size(),
            state->qse_groups.size() == 1 ? "" : "s"
        );

        // Sort the QSE groups by mean timescale so that fastest groups get equilibrated first (as these may feed slower groups)
        LOG_TRACE_L1(m_logger, "Sorting algebraic set by mean timescale...");
        std::ranges::sort(state->qse_groups, [](const QSEGroup& a, const QSEGroup& b) {
            return a.mean_timescale < b.mean_timescale;
        });

        LOG_TRACE_L1(m_logger, "Finalizing dynamic species list...");
        for (const auto& species : m_baseEngine.getNetworkSpecies(ctx)) {
            const bool involvesAlgebraic = involvesSpeciesInQSE(ctx, species);
            if (std::ranges::find(state->dynamic_species, species) == state->dynamic_species.end() && !involvesAlgebraic) {
                state->dynamic_species.push_back(species);
            }
        }
        LOG_TRACE_L1(m_logger, "Final dynamic species set: {}", utils::iterable_to_delimited_string(state->dynamic_species));

        LOG_TRACE_L1(m_logger, "Creating QSE solvers for each identified QSE group...");
        for (const auto& group : state->qse_groups) {
            std::vector<Species> groupAlgebraicSpecies;
            for (const auto& species : group.algebraic_species) {
                groupAlgebraicSpecies.push_back(species);
            }
            state->qse_solvers.push_back(std::make_unique<QSESolver>(groupAlgebraicSpecies, m_baseEngine, state->sun_ctx));
        }
        LOG_TRACE_L1(m_logger, "{} QSE solvers created.", state->qse_solvers.size());

        LOG_TRACE_L1(m_logger, "Calculating final equilibrated composition...");
        fourdst::composition::Composition result = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);
        LOG_TRACE_L1(m_logger, "Final equilibrated composition calculated...");

        return result;
    }

    void MultiscalePartitioningEngineView::exportToDot(
        scratch::StateBlob &ctx,
        const std::string &filename,
        const fourdst::composition::Composition &comp,
        const double T9,
        const double rho
    ) const {
        std::ofstream dotFile(filename);
        if (!dotFile.is_open()) {
            LOG_ERROR(m_logger, "Failed to open file for writing: {}", filename);
            throw std::runtime_error("Failed to open file for writing: " + filename);
        }

        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);

        const auto& all_species = m_baseEngine.getNetworkSpecies(ctx);
        const auto& all_reactions = m_baseEngine.getNetworkReactions(ctx);

        // --- 1. Pre-computation and Categorization ---

        // Categorize species into algebraic, seed, and core dynamic
        std::unordered_set<Species> algebraic_species;
        std::unordered_set<Species> seed_species;
        for (const auto& group : state->qse_groups) {
            if (group.is_in_equilibrium) {
                algebraic_species.insert(group.algebraic_species.begin(), group.algebraic_species.end());
                seed_species.insert(group.seed_species.begin(), group.seed_species.end());
            }
        }

        const fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, comp, T9, rho, false);
        // Calculate reaction flows and find min/max for logarithmic scaling of transparency
        std::vector<double> reaction_flows;
        reaction_flows.reserve(all_reactions.size());
        double min_log_flow = std::numeric_limits<double>::max();
        double max_log_flow = std::numeric_limits<double>::lowest();

        for (const auto& reaction : all_reactions) {
            double flow = std::abs(m_baseEngine.calculateMolarReactionFlow(ctx, *reaction, qseComposition, T9, rho));
            reaction_flows.push_back(flow);
            if (flow > 1e-99) { // Avoid log(0)
                double log_flow = std::log10(flow);
                min_log_flow = std::min(min_log_flow, log_flow);
                max_log_flow = std::max(max_log_flow, log_flow);
            }
        }
        const double log_flow_range = (max_log_flow > min_log_flow) ? (max_log_flow - min_log_flow) : 1.0;

        // --- 2. Write DOT file content ---

        dotFile << "digraph PartitionedNetwork {\n";
        dotFile << "    graph [rankdir=TB, splines=true, overlap=false, bgcolor=\"#f8fafc\", label=\"Multiscale Partitioned Network View\", fontname=\"Helvetica\", fontsize=16, labeljust=l];\n";
        dotFile << "    node [shape=circle, style=filled, fontname=\"Helvetica\", width=0.8, fixedsize=true];\n";
        dotFile << "    edge [fontname=\"Helvetica\", fontsize=10];\n\n";

        // --- Node Definitions ---
        // Define all species nodes first, so they can be referenced by clusters and ranks later.
        dotFile << "    // --- Species Nodes Definitions ---\n";
        std::map<int, std::vector<std::string>> species_by_mass;
        for (const auto & species : all_species) {
            std::string fillcolor = "#f1f5f9"; // Default: Other/Uninvolved

            // Determine color based on category. A species can be a seed and also in the core dynamic group.
            // The more specific category (algebraic, then seed) takes precedence.
            if (algebraic_species.contains(species)) {
                fillcolor = "#e0f2fe"; // Light Blue: Algebraic (in QSE)
            } else if (seed_species.contains(species)) {
                fillcolor = "#a7f3d0"; // Light Green: Seed (Dynamic, feeds a QSE group)
            } else if (std::ranges::contains(state->dynamic_species, species)) {
                fillcolor = "#dcfce7"; // Pale Green: Core Dynamic
            }
            dotFile << "    \"" << species.name() << "\" [label=\"" << species.name() << "\", fillcolor=\"" << fillcolor << "\"];\n";

            // Group species by mass number for ranked layout.
            // If species.a() returns incorrect values (e.g., 0 for many species), they will be grouped together here.
            species_by_mass[species.a()].emplace_back(species.name());
        }
        dotFile << "\n";

        // --- Layout and Ranking ---
        // Enforce a top-down layout based on mass number.
        dotFile << "    // --- Layout using Ranks ---\n";
        for (const auto &species_list: species_by_mass | std::views::values) {
            dotFile << "    { rank=same; ";
            for (const auto& name : species_list) {
                dotFile << "\"" << name << "\"; ";
            }
            dotFile << "}\n";
        }
        dotFile << "\n";

        // Chain by mass to get top down ordering
        dotFile << "    // --- Chain by Mass ---\n";
        for (const auto& [mass, species_list] : species_by_mass) {
            // Find the next largest mass in the species list
            int minLargestMass = std::numeric_limits<int>::max();
            for (const auto &next_mass: species_by_mass | std::views::keys) {
                if (next_mass > mass && next_mass < minLargestMass) {
                    minLargestMass = next_mass;
                }
            }
            if (minLargestMass != std::numeric_limits<int>::max()) {
                // Connect the current mass to the next largest mass
                dotFile << "    \"" << species_list[0] << "\" -> \"" << species_by_mass[minLargestMass][0] << "\" [style=invis];\n";
            }
        }

        // --- QSE Group Clusters ---
        // Draw a prominent box around the algebraic species of each valid QSE group.
        dotFile << "    // --- QSE Group Clusters ---\n";
        int group_counter = 0;
        for (const auto& group : state->qse_groups) {
            if (!group.is_in_equilibrium || group.algebraic_species.empty()) {
                continue;
            }
            dotFile << "    subgraph cluster_qse_" << group_counter++ << " {\n";
            dotFile << "        label = \"QSE Group " << group_counter << "\";\n";
            dotFile << "        style = \"filled,rounded\";\n";
            dotFile << "        color = \"#38bdf8\";\n"; // A bright, visible blue for the border
            dotFile << "        penwidth = 2.0;\n";      // Thicker border
            dotFile << "        bgcolor = \"#f0f9ff80\";\n"; // Light blue fill with transparency
            dotFile << "        subgraph cluster_seed_" << group_counter << " {\n";
            dotFile << "            label = \"Seed Species\";\n";
            dotFile << "            style = \"filled,rounded\";\n";
            dotFile << "            color = \"#a7f3d0\";\n"; // Light green for seed species
            dotFile << "            penwidth = 1.5;\n"; // Thinner border for seed cluster
            std::vector<std::string> seed_node_ids;
            seed_node_ids.reserve(group.seed_species.size());
            for (const auto& species : group.seed_species) {
                std::stringstream ss;
                ss << "node_" << group_counter << "_seed_" << species.name();
                dotFile << "            " << ss.str() << "    [label=\"" << species.name() << "\"];\n";
                seed_node_ids.push_back(ss.str());
            }
            for (size_t i = 0; i < seed_node_ids.size() - 1; ++i) {
                dotFile << "            " << seed_node_ids[i] << " -> " << seed_node_ids[i + 1] << " [style=invis];\n";
            }
            dotFile << "        }\n";
            dotFile << "        subgraph cluster_algebraic_" << group_counter << " {\n";
            dotFile << "            label = \"Algebraic Species\";\n";
            dotFile << "            style = \"filled,rounded\";\n";
            dotFile << "            color = \"#e0f2fe\";\n"; // Light blue for algebraic species
            dotFile << "            penwidth = 1.5;\n"; // Thinner border for algebraic cluster
            std::vector<std::string> algebraic_node_ids;
            algebraic_node_ids.reserve(group.algebraic_species.size());
            for (const Species& species : group.algebraic_species) {
                std::stringstream ss;
                ss << "node_" << group_counter << "_algebraic_" << species.name();
                dotFile << "            " << ss.str() << "    [label=\"" << species.name() << "\"];\n";
                algebraic_node_ids.push_back(ss.str());
            }
            // Make invisible edges between algebraic indices to keep them in top-down order
            for (size_t i = 0; i < algebraic_node_ids.size() - 1; ++i) {
                dotFile << "            " << algebraic_node_ids[i] << " -> " << algebraic_node_ids[i + 1] << " [style=invis];\n";
            }
            dotFile << "        }\n";
            dotFile << "    }\n";
        }
        dotFile << "\n";


        // --- Legend ---
        // Add a legend to explain colors and conventions.
        dotFile << "    // --- Legend ---\n";
        dotFile << "    subgraph cluster_legend {\n";
        dotFile << "        rank = sink"; // Try to push the legend to the bottom
        dotFile << "        label = \"Legend\";\n";
        dotFile << "        bgcolor = \"#ffffff\";\n";
        dotFile << "        color = \"#e2e8f0\";\n";
        dotFile << "        node [shape=box, style=filled, fontname=\"Helvetica\"];\n";
        dotFile << "        key_core [label=\"Core Dynamic\", fillcolor=\"#dcfce7\"];\n";
        dotFile << "        key_seed [label=\"Seed (Dynamic)\", fillcolor=\"#a7f3d0\"];\n";
        dotFile << "        key_qse [label=\"Algebraic (QSE)\", fillcolor=\"#e0f2fe\"];\n";
        dotFile << "        key_other [label=\"Other\", fillcolor=\"#f1f5f9\"];\n";
        dotFile << "        key_info [label=\"Edge Opacity ~ log(Reaction Flow)\", shape=plaintext];\n";
        dotFile << "        ";// Use invisible edges to stack legend items vertically
        dotFile << "        key_core -> key_seed -> key_qse -> key_other -> key_info [style=invis];\n";
        dotFile << "    }\n\n";

        // --- Reaction Edges ---
        // Draw edges with transparency scaled by the log of the molar reaction flow.
        dotFile << "    // --- Reaction Edges ---\n";
        for (size_t i = 0; i < all_reactions.size(); ++i) {
            const auto& reaction = all_reactions[i];
            const double flow = reaction_flows[i];

            if (flow < 1e-99) continue; // Don't draw edges for negligible flows

            double log_flow_val = std::log10(flow);
            double norm_alpha = (log_flow_val - min_log_flow) / log_flow_range;
            int alpha_val = 0x30 + static_cast<int>(norm_alpha * (0xFF - 0x30)); // Scale from ~20% to 100% opacity
            alpha_val = std::clamp(alpha_val, 0x00, 0xFF);

            std::stringstream alpha_hex;
            alpha_hex << std::setw(2) << std::setfill('0') << std::hex << alpha_val;
            std::string edge_color = "#475569" + alpha_hex.str();

            std::string reactionNodeId = "reaction_" + std::string(reaction.id());
            dotFile << "    \"" << reactionNodeId << "\" [shape=point, fillcolor=black, width=0.05, height=0.05];\n";
            for (const auto& reactant : reaction.reactants()) {
                dotFile << "    \"" << reactant.name() << "\" -> \"" << reactionNodeId << "\" [color=\"" << edge_color << "\", arrowhead=none];\n";
            }
            for (const auto& product : reaction.products()) {
                dotFile << "    \"" << reactionNodeId << "\" -> \"" << product.name() << "\" [color=\"" << edge_color << "\"];\n";
            }
            dotFile << "\n";
        }

        dotFile << "}\n";
        dotFile.close();
    }

    std::vector<Species> MultiscalePartitioningEngineView::getFastSpecies(
        scratch::StateBlob& ctx
    ) const {
        const auto& all_species = m_baseEngine.getNetworkSpecies(ctx);
        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        std::vector<Species> fast_species;
        fast_species.reserve(all_species.size() - state->dynamic_species.size());
        for (const auto& species : all_species) {
            auto it = std::ranges::find(state->dynamic_species, species);
            if (it == state->dynamic_species.end()) {
                fast_species.push_back(species);
            }
        }
        return fast_species;
    }

    const std::vector<Species> & MultiscalePartitioningEngineView::getDynamicSpecies(
        scratch::StateBlob& ctx
    ) {
        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        return state->dynamic_species;
    }

    PrimingReport MultiscalePartitioningEngineView::primeEngine(
        scratch::StateBlob& ctx,
        const NetIn &netIn
    ) const {
        return m_baseEngine.primeEngine(ctx, netIn);
    }


    bool MultiscalePartitioningEngineView::involvesSpecies(
        scratch::StateBlob& ctx,
        const Species &species
    ) {
        if (involvesSpeciesInQSE(ctx, species)) return true; // check this first since the vector is likely to be smaller so short circuit cost is less on average
        if (involvesSpeciesInDynamic(ctx, species)) return true;
        return false;
    }

    bool MultiscalePartitioningEngineView::involvesSpeciesInQSE(
        scratch::StateBlob& ctx,
        const Species &species
    ) {
        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        return std::ranges::find(state->algebraic_species, species) != state->algebraic_species.end();
    }

    bool MultiscalePartitioningEngineView::involvesSpeciesInDynamic(
        scratch::StateBlob& ctx,
        const Species &species
    ) {
        const auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        return std::ranges::find(state->dynamic_species, species) != state->dynamic_species.end();
    }

    fourdst::composition::Composition MultiscalePartitioningEngineView::getNormalizedEquilibratedComposition(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract& comp,
        const double T9,
        const double rho,
        const bool trust
    ) const {
        if (trust) {
            return fourdst::composition::Composition(comp);
        }
        // Caching mechanism to avoid redundant QSE solves
        const std::array<uint64_t, 3> hashes = {
            fourdst::composition::utils::CompositionHash::hash_exact(comp),
            std::hash<double>()(T9),
            std::hash<double>()(rho)
        };

        auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        const uint64_t composite_hash = XXHash64::hash(hashes.begin(), sizeof(uint64_t) * 3, 0);
        if (state->composition_cache.contains(composite_hash)) {
            LOG_TRACE_L3(m_logger, "Cache Hit in Multiscale Partitioning Engine View for composition at T9 = {}, rho = {}.", T9, rho);
            return state->composition_cache.at(composite_hash);
        }
        LOG_TRACE_L3(m_logger, "Cache Miss in Multiscale Partitioning Engine View for composition at T9 = {}, rho = {}. Solving QSE abundances...", T9, rho);

        // Only solve if the composition and thermodynamic conditions have not been cached yet
        fourdst::composition::Composition qseComposition(solveQSEAbundances(ctx, comp, T9, rho));

        state->composition_cache[composite_hash] = qseComposition;

        return qseComposition;
    }

    fourdst::composition::Composition MultiscalePartitioningEngineView::collectComposition(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho
    ) const {
        const fourdst::composition::Composition result = m_baseEngine.collectComposition(ctx, comp, T9, rho);

        fourdst::composition::Composition qseComposition = getNormalizedEquilibratedComposition(ctx, result, T9, rho, false);

        return qseComposition;
    }

    SpeciesStatus MultiscalePartitioningEngineView::getSpeciesStatus(
        scratch::StateBlob& ctx,
        const Species &species
    ) const {
        const SpeciesStatus status = m_baseEngine.getSpeciesStatus(ctx, species);
        if (status == SpeciesStatus::ACTIVE && involvesSpeciesInQSE(ctx, species)) {
            return SpeciesStatus::EQUILIBRIUM;
        }
        return status;
    }

    std::optional<StepDerivatives<double>> MultiscalePartitioningEngineView::getMostRecentRHSCalculation(
        scratch::StateBlob& ctx
    ) const {
        return m_baseEngine.getMostRecentRHSCalculation(ctx);
    }

    std::unique_ptr<scratch::StateBlob> MultiscalePartitioningEngineView::constructStateBlob(const scratch::StateBlob *blob) const {
        std::unique_ptr<scratch::StateBlob> i_blob;
        if (blob) {
            i_blob = blob->clone_structure();
        } else {
            i_blob = std::make_unique<scratch::StateBlob>();
        }
        i_blob->enroll<scratch::MultiscalePartitioningEngineViewScratchPad>();
        auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, false>(*i_blob);
        state->initialize();
        return i_blob;
    }

    size_t MultiscalePartitioningEngineView::getSpeciesIndex(
        scratch::StateBlob& ctx,
        const Species &species
    ) const {
        return m_baseEngine.getSpeciesIndex(ctx, species);
    }


    std::vector<std::vector<Species>> MultiscalePartitioningEngineView::partitionByTimescale(
        scratch::StateBlob& ctx,
        const fourdst::composition::Composition &comp,
        const double T9,
        const double rho
    ) const {
        LOG_TRACE_L1(m_logger, "Partitioning by timescale...");
        const auto destructionTimescale= m_baseEngine.getSpeciesDestructionTimescales(ctx, comp, T9, rho);
        const auto netTimescale = m_baseEngine.getSpeciesTimescales(ctx, comp, T9, rho);

        if (!destructionTimescale || !netTimescale) {
            LOG_CRITICAL(m_logger, "Failed to compute species timescales for partitioning due to base engine error.");
            m_logger -> flush_log();
            throw exceptions::EngineError("Failed to compute species timescales for partitioning in MultiscalePartitioningEngineView due to base engine error.");
        }
        const std::unordered_map<Species, double>& destruction_timescales = destructionTimescale.value();
        [[maybe_unused]] const std::unordered_map<Species, double>& net_timescales = netTimescale.value();

        LOG_TRACE_L3(
            m_logger,
            "{}",
            [&]() -> std::string {
                std::stringstream ss;
                for (const auto& [species, destruction_timescale] : destruction_timescales) {
                    ss << std::format("For {:6} destruction timescale is {:10.4E}s\n", species.name(), destruction_timescale);
                }
                return ss.str();
            }()
        );

        const auto& all_species = m_baseEngine.getNetworkSpecies(ctx);

        std::vector<std::pair<double, Species>> sorted_destruction_timescales;
        for (const auto & species : all_species) {
            double destruction_timescale = destruction_timescales.at(species);
            if (std::isfinite(destruction_timescale) && destruction_timescale > 0) {
                LOG_TRACE_L2(m_logger, "Species {:6} has finite destruction timescale: destruction: {:10.4E} s, net: {:10.4E} s", species.name(), destruction_timescale, net_timescales.at(species));
                sorted_destruction_timescales.emplace_back(destruction_timescale, species);
            } else {
                LOG_TRACE_L2(m_logger, "Species {:6} has infinite or negative destruction timescale: destruction: {:10.4E} s, net: {:10.4E} s", species.name(), destruction_timescale, net_timescales.at(species));
            }
        }

        std::ranges::sort(
            sorted_destruction_timescales,
            [](const auto& a, const auto& b)
            {
                return a.first > b.first;
            }
        );

        std::vector<std::vector<Species>> final_pools;
        if (sorted_destruction_timescales.empty()) {
            return final_pools;
        }

        constexpr double ABSOLUTE_QSE_TIMESCALE_THRESHOLD = 3.156e7; // Absolute threshold for QSE timescale (1 yr)
        constexpr double MIN_GAP_THRESHOLD = 2.0; // Require a 2 order of magnitude gap
        constexpr double MAX_MOLAR_ABUNDANCE_THRESHOLD = 1e-10; // Maximum molar abundance which a fast species is allowed to have (anything more abundant is always considered dynamic)
        constexpr double MIN_MOLAR_ABUNDANCE_THRESHOLD = 1e-50; // Minimum molar abundance to consider a species at all (anything less abundance will be classed as dynamic but with the intent that some latter view will deal with it)

        LOG_TRACE_L2(m_logger, "Found {:5} species with finite timescales.", sorted_destruction_timescales.size());
        LOG_TRACE_L2(m_logger, "Absolute QSE timescale threshold: {:7.3E} seconds ({:7.2E} years).",
            ABSOLUTE_QSE_TIMESCALE_THRESHOLD, ABSOLUTE_QSE_TIMESCALE_THRESHOLD / 3.156e7);
        LOG_TRACE_L2(m_logger, "Minimum gap threshold: {:3} orders of magnitude.", MIN_GAP_THRESHOLD);
        LOG_TRACE_L2(m_logger, "Maximum molar abundance threshold for fast species consideration : {:3}.", MAX_MOLAR_ABUNDANCE_THRESHOLD);
        LOG_TRACE_L2(m_logger, "Minimum molar abundance threshold for species consideration : {:3}.", MIN_MOLAR_ABUNDANCE_THRESHOLD);

        std::vector<Species> dynamic_pool_species;
        std::vector<std::pair<double, Species>> fast_candidates;

        // 1. First Pass: Absolute Timescale Cutoff
        for (const auto& [destruction_timescale, species] : sorted_destruction_timescales) {
            if (destruction_timescale > ABSOLUTE_QSE_TIMESCALE_THRESHOLD) {
                LOG_TRACE_L2(m_logger, "Species {:5} with timescale {:10.4E} is considered dynamic (slower than qse timescale threshold).",
                    species.name(), destruction_timescale);
                dynamic_pool_species.push_back(species);
            } else {
                const double Yi = comp.getMolarAbundance(species);
                if (Yi > MAX_MOLAR_ABUNDANCE_THRESHOLD) {
                    LOG_TRACE_L2(m_logger, "Species {:5} with abundance {:10.4E} is considered dynamic (above minimum abundance threshold of {}).",
                        species.name(), Yi, MAX_MOLAR_ABUNDANCE_THRESHOLD);
                    dynamic_pool_species.push_back(species);
                    continue;
                }
                if (Yi < MIN_MOLAR_ABUNDANCE_THRESHOLD) {
                    LOG_TRACE_L2(m_logger, "Species {:5} with abundance {:10.4E} is considered dynamic (below minimum abundance threshold of {:10.4E}). Likely another network view (such as adaptive engine view) will be needed to deal with this species",
                        species.name(), Yi, MIN_MOLAR_ABUNDANCE_THRESHOLD);
                    dynamic_pool_species.push_back(species);
                    continue;
                }
                LOG_TRACE_L2(m_logger, "Species {:5} with timescale {:10.4E} and molar abundance {:10.4E} is a candidate fast species (faster than qse timescale threshold and less than the molar abundance threshold).",
                    species.name(), destruction_timescale, Yi);
                fast_candidates.emplace_back(destruction_timescale, species);
            }
        }

        if (!dynamic_pool_species.empty()) {
            LOG_TRACE_L2(m_logger, "Found {:5} dynamic species (slower than QSE timescale threshold).", dynamic_pool_species.size());
            final_pools.push_back(dynamic_pool_species);
        }

        if (fast_candidates.empty()) {
            LOG_TRACE_L2(m_logger, "No fast candidates found.");
            return final_pools;
        }

        // 2. Second Pass: Gap Detection on the remaining "fast" candidates
        std::vector<size_t> split_points;
        for (size_t i = 0; i < fast_candidates.size() - 1; ++i) {
            const double t1 = fast_candidates[i].first;
            const double t2 = fast_candidates[i+1].first;
            if (std::log10(t1) - std::log10(t2) > MIN_GAP_THRESHOLD) {
                LOG_TRACE_L2(m_logger, "Detected gap between species {:5} (timescale {:8.2E}) and {:5} (timescale {:10.2E}).",
                    fast_candidates[i].second.name(), t1,
                    fast_candidates[i+1].second.name(), t2);
                split_points.push_back(i + 1);
            }
        }

        size_t last_split = 0;
        for (const size_t split : split_points) {
            std::vector<Species> pool;
            for (size_t i = last_split; i < split; ++i) {
                pool.push_back(fast_candidates[i].second);
            }
            final_pools.push_back(pool);
            last_split = split;
        }

        std::vector<Species> final_fast_pool;
        for (size_t i = last_split; i < fast_candidates.size(); ++i) {
            final_fast_pool.push_back(fast_candidates[i].second);
        }
        final_pools.push_back(final_fast_pool);

        LOG_TRACE_L2(m_logger, "Final partitioned pools: {}",
            [&]() -> std::string {
                std::stringstream ss;
                int oc = 0;
                for (const auto& pool : final_pools) {
                    ss << "[";
                    int ic = 0;
                    for (const auto& species : pool) {
                        ss << species.name();
                        if (ic < pool.size() - 1) {
                            ss << ", ";
                        }
                        ic++;
                    }
                    ss << "]";
                    if (oc < final_pools.size() - 1) {
                        ss << ", ";
                    }
                    oc++;
                }
                return ss.str();
            }());

        LOG_TRACE_L2(
            m_logger,
            "Species Timescales: {}",
            [&]() -> std::string {
                std::stringstream ss;
                size_t poolID = 0;
                for (const auto& pool : final_pools) {
                    ss << "Pool #" << poolID << " [";
                    int ic = 0;
                    for (const auto& species : pool) {
                        ss << species << ": " << destruction_timescales.at(species);
                        if (ic < pool.size() - 1) {
                            ss << ", ";
                        }
                        ic++;
                    }
                    ss << "]";
                    poolID++;
                }
                return ss.str();
            }()
            );
        return final_pools;

    }

    std::pair<bool, reaction::ReactionSet> MultiscalePartitioningEngineView::group_is_a_qse_cluster(
        scratch::StateBlob& ctx,
        const fourdst::composition::Composition &comp,
        const double T9,
        const double rho,
        const QSEGroup &group
    ) const {
        auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);
        double FLUX_RATIO_THRESHOLD = state->flux_coupling_threshold;

        const std::unordered_set<Species> algebraic_group_members(
            group.algebraic_species.begin(),
            group.algebraic_species.end()
        );

        const std::unordered_set<Species> seed_group_members(
            group.seed_species.begin(),
            group.seed_species.end()
        );

        reaction::ReactionSet group_reaction_set;
        double coupling_flux = 0.0;
        double leakage_flux = 0.0;

        for (const auto& reaction: m_baseEngine.getNetworkReactions(ctx)) {
            const double flow = std::abs(m_baseEngine.calculateMolarReactionFlow(ctx, *reaction, comp, T9, rho));
            if (flow == 0.0) {
                continue; // Skip reactions with zero flow
            }
            bool has_internal_algebraic_reactant = false;

            for (const auto& reactant : reaction->reactants()) {
                if (algebraic_group_members.contains(reactant)) {
                    has_internal_algebraic_reactant = true;
                    LOG_TRACE_L3(m_logger, "Adjusting destruction flux (+= {} mol g^-1 s^-1) for QSEGroup due to reactant {} from reaction {}",
                                 flow, reactant.name(), reaction->id());
                }

            }

            bool has_internal_algebraic_product = false;

            for (const auto& product : reaction->products()) {
                if (algebraic_group_members.contains(product)) {
                    has_internal_algebraic_product = true;
                    LOG_TRACE_L3(m_logger, "Adjusting creation flux (+= {} mol g^-1 s^-1) for QSEGroup due to product {} from reaction {}",
                                 flow, product.name(), reaction->id());
                }
            }

            if (!has_internal_algebraic_product && !has_internal_algebraic_reactant) {
                LOG_TRACE_L3(m_logger, "{}: Skipping reaction {} as it has no internal algebraic species in reactants or products.", group.toString(false), reaction->id());
                continue;
            }

            group_reaction_set.add_reaction(reaction->clone());

            LOG_TRACE_L3(
                m_logger,
                "Reaction {:20} has flow {:15.4E} mol g^-1 s^-1 contributing to QSEGroup {:40}",
                reaction->id(),
                flow,
                group.toString(false)
            );

            double algebraic_participants = 0;
            double seed_participants = 0;
            double external_participants = 0;

            std::unordered_set<Species> participants;
            for(const auto& p : reaction->reactants()) participants.insert(p);
            for(const auto& p : reaction->products()) participants.insert(p);

            for (const auto& species : participants) {
                if (algebraic_group_members.contains(species)) {
                    LOG_TRACE_L3(m_logger, "{}: Species {} is an algebraic participant in reaction {}.", group.toString(true), species.name(), reaction->id());
                    algebraic_participants++;
                } else if (seed_group_members.contains(species)) {
                    LOG_TRACE_L3(m_logger, "{}: Species {} is a seed participant in reaction {}.", group.toString(true), species.name(), reaction->id());
                    seed_participants++;
                } else {
                    LOG_TRACE_L3(m_logger, "{}: Species {} is an external participant in reaction {}.", group.toString(true), species.name(), reaction->id());
                    external_participants++;
                }
            }

            const double total_participants = algebraic_participants + seed_participants + external_participants;

            if (total_participants == 0) {
                LOG_CRITICAL(m_logger, "Some catastrophic error has occurred. Reaction {} has no participants.", reaction->id());
                throw std::runtime_error("Some catastrophic error has occurred. Reaction " + std::string(reaction->id()) + " has no participants.");
            }

            const double leakage_fraction = external_participants / total_participants;
            const double coupling_fraction = (algebraic_participants + seed_participants) / total_participants;

            leakage_flux += flow * leakage_fraction;
            coupling_flux += flow * coupling_fraction;

            LOG_TRACE_L2(m_logger, "Reaction {:20} contributes coupling flux {:15.4E} and leakage flux {:15.4E} to QSEGroup {}.",
                reaction->id(),
                flow * coupling_fraction,
                flow * leakage_fraction,
                group.toString(false)
            );


        }

        bool leakage_coupled = (coupling_flux / leakage_flux > FLUX_RATIO_THRESHOLD);
        return std::make_pair(leakage_coupled, group_reaction_set);
    }

    bool MultiscalePartitioningEngineView::group_is_a_qse_pipeline(
        scratch::StateBlob& ctx,
        const fourdst::composition::Composition &comp,
        const double T9,
        const double rho,
        const QSEGroup &group
    ) const {
        // Total fluxes (Standard check)
        double total_prod = 0.0;
        double total_dest = 0.0;

        // Charged-particle only fluxes (Heuristic for fast-neutron regimes)
        double charged_prod = 0.0;
        double charged_dest = 0.0;

        for (const auto& reaction : m_baseEngine.getNetworkReactions(ctx)) {
            const double flow = m_baseEngine.calculateMolarReactionFlow(ctx, *reaction, comp, T9, rho);
            if (std::abs(flow) < 1.0e-99) continue;

            int groupNetStoichiometry = 0;
            for (const auto& species : group.algebraic_species) {
                groupNetStoichiometry += reaction->stoichiometry(species);
            }
            if (groupNetStoichiometry == 0) continue;

            const double flux = flow * groupNetStoichiometry;
            const bool is_neutron_reaction = reaction->contains(n_1);

            if (flux > 0.0) {
                total_prod += flux;
                if (!is_neutron_reaction) charged_prod += flux;
            } else {
                total_dest += -flux;
                if (!is_neutron_reaction) charged_dest += -flux;
            }
        }

        // Check 1: Total Balance
        const double mean_total = (total_prod + total_dest) / 2.0;
        const double diff_total = std::abs(total_prod - total_dest);
        bool total_balanced = (mean_total > 0) && ((diff_total / mean_total) < 0.05);

        // Check 2: Charged-Particle Balance
        // Only valid if there IS charged flow
        const double mean_charged = (charged_prod + charged_dest) / 2.0;
        const double diff_charged = std::abs(charged_prod - charged_dest);
        bool charged_balanced = (mean_charged > 0) && ((diff_charged / mean_charged) < 0.05);

        LOG_TRACE_L2(m_logger, "{} Pipeline Check | Total Pass: {} | Charged Pass: {}",
            group.toString(false), total_balanced, charged_balanced);

        return total_balanced || charged_balanced;
    }


    MultiscalePartitioningEngineView::FluxValidationResult MultiscalePartitioningEngineView::validateGroupsWithFluxAnalysis(
        scratch::StateBlob& ctx,
        const std::vector<QSEGroup> &candidate_groups,
        const fourdst::composition::Composition &comp,
        const double T9, const double rho
    ) const {
        std::vector<QSEGroup> validated_groups;
        std::vector<reaction::ReactionSet> group_reactions;
        std::vector<QSEGroup> invalidated_groups;
        validated_groups.reserve(candidate_groups.size());
        group_reactions.reserve(candidate_groups.size());
        for (auto& group : candidate_groups) {
            // Values for measuring the flux coupling vs leakage
            auto [leakage_coupled, group_reaction_set] = group_is_a_qse_cluster(ctx, comp, T9, rho, group);


            bool is_flow_balanced = group_is_a_qse_pipeline(ctx, comp, T9, rho, group);

            if (leakage_coupled) {
                LOG_TRACE_L1(m_logger, "{} is in equilibrium due to high coupling flux", group.toString(false));
                validated_groups.emplace_back(group);
                validated_groups.back().is_in_equilibrium = true;
                group_reactions.emplace_back(group_reaction_set);
            }
            else if (is_flow_balanced) {
                LOG_TRACE_L1(m_logger, "{} is in equilibrium due to balanced production and destruction fluxes.", group.toString(false));
                validated_groups.emplace_back(group);
                validated_groups.back().is_in_equilibrium = true;
                group_reactions.emplace_back(group_reaction_set);
            }
            else {
                LOG_TRACE_L1(m_logger, "{} is NOT in equilibrium due to high leakage flux and lack of pipeline detection.", group.toString(false));
                invalidated_groups.emplace_back(group);
                invalidated_groups.back().is_in_equilibrium = false;
            }
        }

        LOG_TRACE_L1(m_logger, "Validated {} QSE groups and invalidated {} QSE groups after flux analysis.", validated_groups.size(), invalidated_groups.size());
        return {validated_groups, invalidated_groups, group_reactions};
    }

    fourdst::composition::Composition MultiscalePartitioningEngineView::solveQSEAbundances(
        scratch::StateBlob& ctx,
        const fourdst::composition::CompositionAbstract &comp,
        const double T9,
        const double rho
    ) const {
        LOG_TRACE_L2(m_logger, "Solving for QSE abundances...");
        auto* state = scratch::get_state<scratch::MultiscalePartitioningEngineViewScratchPad, true>(ctx);

        fourdst::composition::Composition outputComposition(comp);

        std::vector<Species> species;
        std::vector<double> abundances;
        species.reserve(state->algebraic_species.size());
        abundances.reserve(state->algebraic_species.size());

        for (const auto& [group, solver]: std::views::zip(state->qse_groups, state->qse_solvers)) {
            const fourdst::composition::Composition& groupResult = solver->solve(ctx, outputComposition, T9, rho);
            for (const auto& [sp, y] : groupResult) {
                if (!std::isfinite(y)) {
                    LOG_CRITICAL(m_logger, "Non-finite abundance {} computed for species {} in QSE group solve at T9 = {}, rho = {}.",
                        y, sp.name(), T9, rho);
                    m_logger->flush_log();
                    throw exceptions::EngineError("Non-finite abundance computed for species " + std::string(sp.name()) + " in QSE group solve.");
                }
                if (y < 0.0 && std::abs(y) < 1e-16) {
                    abundances.push_back(0.0);
                } else if (y < 0 && std::abs(y) >= 1e-16) {
                    throw fourdst::composition::exceptions::InvalidCompositionError(std::format("Computed negative and non-trivial abundance {} for species {} in QSE group solve at T9 = {}, rho = {}. This likely indicates a failure of the QSE solver to converge to a physical solution.", y, sp.name(), T9, rho));
                } else {
                    abundances.push_back(y);
                }
                species.emplace_back(sp);
            }
            // solver->log_diagnostics(group, outputComposition);
        }
        outputComposition.setMolarAbundance(species, abundances);
        LOG_TRACE_L2(m_logger, "Done solving for QSE abundances!");
        return outputComposition;
    }

    size_t MultiscalePartitioningEngineView::identifyMeanSlowestPool(
        scratch::StateBlob& ctx,
        const std::vector<std::vector<Species>> &pools,
        const fourdst::composition::Composition &comp,
        const double T9,
        const double rho
    ) const {
        const auto& result = m_baseEngine.getSpeciesDestructionTimescales(ctx, comp, T9, rho);
        if (!result) {
            LOG_CRITICAL(m_logger, "Failed to get species destruction timescales due base engine failure");
            m_logger->flush_log();
            throw exceptions::EngineError("Failed to get species destruction timescales due base engine failure");
        }
        const std::unordered_map<Species, double> all_timescales = result.value();

        size_t slowest_pool_index = 0; // Default to the first pool if no valid pool is found
        double slowest_mean_timescale = std::numeric_limits<double>::min();
        size_t count = 0;
        for (const auto& pool : pools) {
            double mean_timescale = 0.0;
            for (const auto& species : pool) {
                const double timescale = all_timescales.at(species);
                mean_timescale += timescale;
            }
            mean_timescale = mean_timescale / static_cast<double>(pool.size());
            if (std::isinf(mean_timescale)) {
                LOG_CRITICAL(m_logger, "Encountered infinite mean timescale for pool {} with species: {}",
                    count, [&]() -> std::string {
                        std::stringstream ss;
                        size_t iCount = 0;
                        for (const auto& species : pool) {
                            ss << species.name() << ": " << all_timescales.at(species);
                            if (iCount < pool.size() - 1) {
                                ss << ", ";
                            }
                            iCount++;
                        }
                        return ss.str();
                    }()
                );
                m_logger->flush_log();
                throw std::logic_error("Encountered infinite mean destruction timescale for a pool while identifying the mean slowest pool set, indicating a potential issue with species timescales. Check log file for more details on specific pool composition...");
            }
            if (mean_timescale > slowest_mean_timescale) {
                slowest_mean_timescale = mean_timescale;
                slowest_pool_index = &pool - &pools[0]; // Get the index of the pool
            }
        }
        return slowest_pool_index;
    }

    std::unordered_map<Species, std::vector<Species>> MultiscalePartitioningEngineView::buildConnectivityGraph(
        scratch::StateBlob& ctx,
        const std::vector<Species> &species_pool,
        const fourdst::composition::Composition &comp,
        double T9,
        double rho
    ) const {
        std::unordered_map<Species, std::vector<Species>> connectivity_graph;
        const std::set<Species> pool_set(species_pool.begin(), species_pool.end());

        const std::unordered_set<Species> pool_species = [&]() -> std::unordered_set<Species> {
            std::unordered_set<Species> result;
            for (const auto& species : species_pool) {
                result.insert(species);
            }
            return result;
        }();

        std::map<size_t, std::vector<reaction::LogicalReaclibReaction*>> speciesReactionMap;
        std::vector<const reaction::LogicalReaclibReaction*> candidate_reactions;

        for (const auto& reaction : m_baseEngine.getNetworkReactions(ctx)) {
            const std::vector<Species> &reactants = reaction->reactants();
            const std::vector<Species> &products = reaction->products();

            std::unordered_set<Species> reactant_set(reactants.begin(), reactants.end());
            std::unordered_set<Species> product_set(products.begin(), products.end());

            // Only consider reactions where at least one distinct reactant and product are in the pool
            if (has_distinct_reactant_and_product_species(pool_species, reactant_set, product_set)) {
                std::set<Species> involvedSet;
                involvedSet.insert(reactants.begin(), reactants.end());
                involvedSet.insert(products.begin(), products.end());

                std::vector<Species> intersection;
                intersection.reserve(involvedSet.size());

                for (const auto& s : pool_species) { // Find intersection with pool species
                    if (involvedSet.contains(s)) {
                        intersection.push_back(s);
                    }
                }

                // Add clique
                for (const auto& u : intersection) {
                    for (const auto& v : intersection) {
                        if (u != v) { // Avoid self-loops
                            connectivity_graph[u].push_back(v);
                        }
                    }
                }
            }
        }


        return connectivity_graph;
    }

    std::vector<MultiscalePartitioningEngineView::QSEGroup> MultiscalePartitioningEngineView::constructCandidateGroups(
        scratch::StateBlob& ctx,
        const std::vector<std::vector<Species>> &candidate_pools,
        const fourdst::composition::Composition &comp,
        const double T9,
        const double rho
    ) const {
        const auto& all_reactions = m_baseEngine.getNetworkReactions(ctx);
        const auto& result = m_baseEngine.getSpeciesDestructionTimescales(ctx, comp, T9, rho);
        if (!result) {
            LOG_ERROR(m_logger, "Failed to get species destruction timescales due base engine failure");
            m_logger->flush_log();
            throw exceptions::EngineError("Failed to get species destruction timescales due base engine failure");
        }
        const std::unordered_map<Species, double> destruction_timescales = result.value();

        std::vector<QSEGroup> candidate_groups;
        for (const auto& pool : candidate_pools) {
            if (pool.empty()) continue; // Skip empty pools

            // For each pool first identify all topological bridge connections
            std::vector<std::pair<const reaction::Reaction*, double>> bridge_reactions;
            for (const auto& ash: pool) {
                for (const auto& reaction : all_reactions) {
                    if (reaction->contains(ash)) {
                        // Check to make sure there is at least one reactant that is not in the pool
                        // This lets seed nuclei bring mass into the QSE group.
                        bool has_external_reactant = false;
                        for (const auto& reactant : reaction->reactants()) {
                            if (std::ranges::find(pool, reactant) == pool.end()) {
                                has_external_reactant = true;
                                LOG_TRACE_L3(m_logger, "Found external reactant {} in reaction {} for species {}.", reactant.name(), reaction->id(), ash.name());
                                break; // Found an external reactant, no need to check further
                            }
                        }
                        if (has_external_reactant) {
                            double flow = std::abs(m_baseEngine.calculateMolarReactionFlow(ctx, *reaction, comp, T9, rho));
                            LOG_TRACE_L3(m_logger, "Found bridge reaction {} with flow {} for species {}.", reaction->id(), flow, ash.name());
                            bridge_reactions.emplace_back(reaction.get(), flow);
                        }
                    }
                }
            }
            std::ranges::sort(
                bridge_reactions,
                [](const auto& a, const auto& b) {
                   return a.second > b.second; // Sort by flow in descending order
                });

            constexpr double MIN_GAP_THRESHOLD = 1; // Minimum logarithmic molar flow gap threshold for bridge reactions
            std::vector<size_t> split_points;
            for (size_t i = 0; i < bridge_reactions.size() - 1; ++i) {
                const double f1 = bridge_reactions[i].second;
                const double f2 = bridge_reactions[i + 1].second;
                if (std::log10(f1) - std::log10(f2) > MIN_GAP_THRESHOLD) {
                    LOG_TRACE_L3(m_logger, "Detected gap between bridge reactions with flows {} and {}.", f1, f2);
                    split_points.push_back(i + 1);
                }
            }

            if (split_points.empty()) { // If no split points were found, we consider the whole set of bridge reactions as one group.
                split_points.push_back(bridge_reactions.size() - 1);
            }

            std::vector<Species> seed_species;
            for (auto &reaction: bridge_reactions | std::views::keys) {
                for (const auto& fuel : reaction->reactants()) {
                    // Only add the fuel if it is not already in the pool
                    if (std::ranges::find(pool, fuel) == pool.end()) {
                        seed_species.push_back(fuel);
                    }
                }
            }
            std::set<Species> pool_species(pool.begin(), pool.end());
            for (const auto& species : seed_species) {
                pool_species.insert(species);
            }
            const std::set<Species> poolSet(pool.begin(), pool.end());
            const std::set<Species> seedSet(seed_species.begin(), seed_species.end());

            double mean_timescale = 0.0;
            for (const auto& species : poolSet) {
                if (destruction_timescales.contains(species)) {
                    mean_timescale += std::min(destruction_timescales.at(species), species.halfLife()); // Use the minimum of destruction timescale and half-life
                } else {
                    mean_timescale += species.halfLife();
                }
            }
            mean_timescale /= static_cast<double>(poolSet.size());
            QSEGroup qse_group(false, poolSet, seedSet, mean_timescale);
            candidate_groups.push_back(qse_group);
        }
        return candidate_groups;
    }

     //////////////////////////////////
    /// QSESolver Member Functions ///
    //////////////////////////////////


     bool MultiscalePartitioningEngineView::QSEGroup::contains(const fourdst::atomic::Species &species) const {
        return algebraic_species.contains(species) || seed_species.contains(species);
     }

     bool MultiscalePartitioningEngineView::QSEGroup::containsAlgebraic(const Species &species) const {
        return algebraic_species.contains(species);
     }

     bool MultiscalePartitioningEngineView::QSEGroup::containsSeed(const Species &species) const {
        return seed_species.contains(species);
     }

     MultiscalePartitioningEngineView::QSESolver::QSESolver(
        const std::vector<fourdst::atomic::Species>& species,
        const DynamicEngine& engine,
        const SUNContext sun_ctx
    ) :
    m_N(species.size()),
    m_engine(engine),
    m_species(species),
    m_sun_ctx(sun_ctx) {
        LOG_TRACE_L1(getLogger(), "Initializing QSE Solver with {} species ({})", m_N,
            [&]() -> std::string {
                std::stringstream ss;
                int count = 0;
                for (const auto& sp : species) {
                    ss << sp.name();
                    if (count < species.size() - 1) {
                        ss << ", ";
                    }
                    count++;
                }
                return ss.str();
            }()
        );
        m_Y = utils::init_sun_vector(m_N, m_sun_ctx);
        m_scale = N_VClone(m_Y);
        m_f_scale = N_VClone(m_Y);
        m_constraints = N_VClone(m_Y);
        m_func_tmpl = N_VClone(m_Y);

        if (!m_Y || !m_scale || !m_constraints || !m_func_tmpl) {
            LOG_CRITICAL(getLogger(), "Failed to allocate SUNVectors for QSE solver.");
            throw std::runtime_error("Failed to allocate SUNVectors for QSE solver.");
        }

        for (size_t i = 0; i < m_N; ++i) {
            m_speciesMap[m_species[i]] = i;
        }

        N_VConst(1.0, m_constraints);
        m_kinsol_mem = KINCreate(m_sun_ctx);

        utils::check_sundials_flag(m_kinsol_mem ? 0 : -1, "KINCreate", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);
        utils::check_sundials_flag(KINInit(m_kinsol_mem, sys_func, m_func_tmpl), "KINInit", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);
        utils::check_sundials_flag(KINSetConstraints(m_kinsol_mem, m_constraints), "KINSetConstraints", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        m_J = SUNDenseMatrix(static_cast<sunindextype>(m_N), static_cast<sunindextype>(m_N), m_sun_ctx);
        utils::check_sundials_flag(m_J ? 0 : -1, "SUNDenseMatrix", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        m_LS = SUNLinSol_Dense(m_Y, m_J, m_sun_ctx);
        utils::check_sundials_flag(m_LS ? 0 : -1, "SUNLinSol_Dense", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        utils::check_sundials_flag(KINSetLinearSolver(m_kinsol_mem, m_LS, m_J), "KINSetLinearSolver", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);
        utils::check_sundials_flag(KINSetJacFn(m_kinsol_mem, sys_jac), "KINSetJacFn", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        utils::check_sundials_flag(KINSetMaxSetupCalls(m_kinsol_mem, 20), "KINSetMaxSetupCalls", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        utils::check_sundials_flag(KINSetFuncNormTol(m_kinsol_mem, 1e-8), "KINSetFuncNormTol", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);
        utils::check_sundials_flag(KINSetNumMaxIters(m_kinsol_mem, 200), "KINSetNumMaxIters", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        utils::check_sundials_flag(KINSetScaledStepTol(m_kinsol_mem, 1e-10), "KINSetScaledStepTol", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        // We want to effectively disable this since enormous changes in order of magnitude are realistic for this problem.
        utils::check_sundials_flag(KINSetMaxNewtonStep(m_kinsol_mem, std::numeric_limits<double>::infinity()), "KINSetMaxNewtonStep", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);
        LOG_TRACE_L1(getLogger(), "Finished initializing QSE Solver.");
    }

    MultiscalePartitioningEngineView::QSESolver::~QSESolver() {
        if (m_Y) {
            N_VDestroy(m_Y);
            m_Y = nullptr;
        }
        if (m_scale) {
            N_VDestroy(m_scale);
            m_scale = nullptr;
        }
        if (m_f_scale) {
            N_VDestroy(m_f_scale);
            m_f_scale = nullptr;
        }
        if (m_constraints) {
            N_VDestroy(m_constraints);
            m_constraints = nullptr;
        }
        if (m_func_tmpl) {
            N_VDestroy(m_func_tmpl);
            m_func_tmpl = nullptr;
        }
        if (m_kinsol_mem) {
            KINFree(&m_kinsol_mem);
            m_kinsol_mem = nullptr;
        }
        if (m_J) {
            SUNMatDestroy(m_J);
            m_J = nullptr;
        }
        if (m_LS) {
            SUNLinSolFree(m_LS);
            m_LS = nullptr;
        }
    }

    fourdst::composition::Composition MultiscalePartitioningEngineView::QSESolver::solve(
        scratch::StateBlob& ctx,
        const fourdst::composition::Composition &comp,
        const double T9,
        const double rho
    ) const {

        fourdst::composition::Composition result = comp;
        UserData data {
            m_engine,
            T9,
            rho,
            result,
            m_speciesMap,
            m_species,
            *this,
            ctx
        };

        utils::check_sundials_flag(KINSetUserData(m_kinsol_mem, &data), "KINSetUserData", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        sunrealtype* y_data = N_VGetArrayPointer(m_Y);
        sunrealtype* scale_data = N_VGetArrayPointer(m_scale);

        // It is more cache optimized to do a standard as opposed to range based for-loop here
        for (size_t i = 0; i < m_N; ++i) {
            const auto& species = m_species[i];
            double Y = result.getMolarAbundance(species);

            constexpr double abundance_floor = 1.0e-100;
            Y = std::max(abundance_floor, Y);
            y_data[i] = Y;
            scale_data[i] = 1.0 / Y;
        }

        StepDerivatives<double> rhsGuess;
        auto cached_rhs = m_engine.getMostRecentRHSCalculation(ctx);
        if (!cached_rhs) {
            const auto initial_rhs = m_engine.calculateRHSAndEnergy(ctx, result, T9, rho, false);
            if (!initial_rhs) {
                throw std::runtime_error("In QSE solver failed to calculate initial RHS for caching");
            }
            rhsGuess = initial_rhs.value();
        } else {
            rhsGuess = cached_rhs.value();
        }

        sunrealtype* f_scale_data = N_VGetArrayPointer(m_f_scale);
        for (size_t i = 0; i < m_N; ++i) {
            const auto& species = m_species[i];
            double dydt = std::abs(rhsGuess.dydt.at(species));
            f_scale_data[i] = 1.0 / (dydt + 1e-15);
        }

        if (m_solves > 0 && m_has_jacobian) {
            // After the initial solution we want to allow kinsol to reuse its state
            utils::check_sundials_flag(KINSetNoInitSetup(m_kinsol_mem, SUNTRUE), "KINSetNoInitSetup", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);
        } else {
            utils::check_sundials_flag(KINSetNoInitSetup(m_kinsol_mem, SUNFALSE), "KINSetNoInitSetup", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);
        }

        LOG_TRACE_L2(
            getLogger(),
            "Starting KINSol QSE solver with initial state: {}",
            [&comp, &rhsGuess, &data]() -> std::string {
                std::ostringstream oss;
                oss << "Solve species: <";
                size_t count = 0;
                for (const auto& species : data.qse_solve_species) {
                    oss << species.name();
                    if (count < data.qse_solve_species.size() - 1) {
                        oss << ", ";
                    }
                    count++;
                }
                oss << "> | Initial abundances and rates: ";
                count = 0;
                for (const auto& [species, abundance] : comp) {
                    oss << species.name() << ": Y = " << abundance << ", dY/dt = " << rhsGuess.dydt.at(species);
                    if (count < comp.size() - 1) {
                        oss << ", ";
                    }
                    count++;
                }
                return oss.str();
            }()
        );
        LOG_TRACE_L2(
            getLogger(),
            "Jacobian Prior to KINSol is: {}",
            [&]() -> std::string {
                std::ostringstream oss;
                sunrealtype* J_data = SUNDenseMatrix_Data(m_J);
                const sunindextype N = SUNDenseMatrix_Columns(m_J);
                oss << "[";
                for (size_t i = 0; i < m_N; ++i) {
                    oss << "[";
                    for (size_t j = 0; j < m_N; ++j) {
                        oss << J_data[i * N + j];
                        if (j < m_N - 1) {
                            oss << ", ";
                        }
                    }
                    oss << "]";
                }
                oss << "]";
                return oss.str();
            }()
        );

        const int flag = KINSol(m_kinsol_mem, m_Y, KIN_LINESEARCH, m_scale, m_f_scale);

        if (flag < 0) {
            double fnorm = 0.0;
            constexpr double ACCEPTABLE_FTOL = 1e-4; // Due to row scaling this is a relative threshold. So this is saying that the imbalance is 0.01% of the abundance scale.
            utils::check_sundials_flag(KINGetFuncNorm(m_kinsol_mem, &fnorm), "KINGetFuncNorm", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);
            if (fnorm < ACCEPTABLE_FTOL) {
                LOG_INFO(getLogger(), "KINSol failed to converge within the maximum number of iterations, but achieved acceptable accuracy with function norm {} < {}. Proceeding with solution.",
                    fnorm, ACCEPTABLE_FTOL);
            } else {
                LOG_CRITICAL(getLogger(), "KINSol failed to converge while solving QSE abundances with flag {}. Flag No.: {}, Error (fNorm): {}", utils::kinsol_ret_code_map.at(flag), flag, fnorm);
                LOG_CRITICAL(getLogger(), "State prior to failure: {}",
                    [&comp, &data]() -> std::string {
                        std::ostringstream oss;
                        oss << "Solve species: <";
                        size_t count = 0;
                        for (const auto& species : data.qse_solve_species) {
                            oss << species.name();
                            if (count < data.qse_solve_species.size() - 1) {
                                oss << ", ";
                            }
                            count++;
                        }
                        oss << "> | Abundances and rates at failure: ";
                        count = 0;
                        for (const auto& [species, abundance] : comp) {
                            oss << species.name() << ": Y = " << abundance;
                            if (count < comp.size() - 1) {
                                oss << ", ";
                            }
                            count++;
                        }
                        oss << " | Temperature: " << data.T9 << ", Density: " << data.rho;
                        return oss.str();
                    }()
                );
                throw exceptions::InvalidQSESolutionError("KINSol failed to converge while solving QSE abundances. " + utils::kinsol_ret_code_map.at(flag));
            }
        }



        for (size_t i = 0; i < m_N; ++i) {
            const auto& species = m_species[i];
            result.setMolarAbundance(species, y_data[i]);
        }

        m_solves++;
        return result;
    }

    size_t MultiscalePartitioningEngineView::QSESolver::solves() const {
        return m_solves;
    }

    void MultiscalePartitioningEngineView::QSESolver::log_diagnostics(const QSEGroup &group, const fourdst::composition::Composition &comp) const {
        long int nni, nfe, nje;

        utils::check_sundials_flag(KINGetNumNonlinSolvIters(m_kinsol_mem, &nni), "KINGetNumNonlinSolvIters", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        utils::check_sundials_flag(KINGetNumFuncEvals(m_kinsol_mem, &nfe), "KINGetNumFuncEvals", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        utils::check_sundials_flag(KINGetNumJacEvals(m_kinsol_mem, &nje), "KINGetNumJacEvals", utils::SUNDIALS_RET_CODE_TYPES::KINSOL);

        LOG_TRACE_L1(getLogger(),
            "QSE Stats | Iters: {} | RHS Evals: {} | Jac Evals: {} | Ratio (J/I): {:.2f} | Algebraic Species: {}",
            nni,
            nfe,
            nje,
            static_cast<double>(nje) / static_cast<double>(nni),
            [&group, &comp]() -> std::string {
                std::ostringstream oss;
                size_t count = 0;
                oss << "[";
                for (const auto& species : group.algebraic_species) {
                    oss << species.name() << "(Y = " << comp.getMolarAbundance(species) << ")";
                    if (count < group.algebraic_species.size() - 1) {
                        oss << ", ";
                    }
                    count++;
                }
                oss << "]";
                return oss.str();
            }()
        );
        LOG_TRACE_L2(getLogger(),
            "Jacobian After KINSol is: {}",
            [&]() -> std::string {
                std::ostringstream oss;
                sunrealtype* J_data = SUNDenseMatrix_Data(m_J);
                const sunindextype N = SUNDenseMatrix_Columns(m_J);
                oss << "[";
                for (size_t i = 0; i < m_N; ++i) {
                    oss << "[";
                    for (size_t j = 0; j < m_N; ++j) {
                        oss << J_data[i * N + j];
                        if (j < m_N - 1) {
                            oss << ", ";
                        }
                    }
                    oss << "]";
                }
                oss << "]";
                return oss.str();
            }()
        );
        getLogger()->flush_log(true);
    }

    std::unique_ptr<MultiscalePartitioningEngineView::QSESolver> MultiscalePartitioningEngineView::QSESolver::clone() const {
        auto new_solver = std::make_unique<QSESolver>(m_species, m_engine, m_sun_ctx);
        return new_solver;
    }

    std::unique_ptr<MultiscalePartitioningEngineView::QSESolver> MultiscalePartitioningEngineView::QSESolver::clone(SUNContext sun_ctx) const {
        auto new_solver = std::make_unique<QSESolver>(m_species, m_engine, sun_ctx);
        return new_solver;
    }


    int MultiscalePartitioningEngineView::QSESolver::sys_func(
        const N_Vector y,
        const N_Vector f,
        void *user_data
    ) {
        const auto* data = static_cast<UserData*>(user_data);

        const sunrealtype* y_data = N_VGetArrayPointer(y);
        sunrealtype* f_data = N_VGetArrayPointer(f);

        const auto& map = data->qse_solve_species_index_map;
        for (size_t index = 0; index < data->qse_solve_species.size(); ++index) {
            const auto& species = data->qse_solve_species[index];
            if (!std::isfinite(y_data[index])) {
                std::string msg = std::format("Non-finite abundance {} encountered for species {} in QSE solver sys_func. Attempting recovery...",
                    y_data[index], species.name());
                LOG_ERROR(getLogger(), "{}", msg);
                return 1; // Potentially recoverable error
            }
            data->comp.setMolarAbundance(species, y_data[index]);
        }

        const auto result = data->engine.calculateRHSAndEnergy(data->ctx, data->comp, data->T9, data->rho, false);

        if (!result) {
            return 1; // Potentially recoverable error
        }

        const auto& dydt = result.value().dydt;
#ifndef NDEBUG
        // In debug mode check that all dydt values are finite and not NaN
        for (const auto &species: map | std::views::keys) {
            const double value = dydt.at(species);
            if (!std::isfinite(value)) {
                std::vector<std::pair<Species, double>> invalid_species;
                for (const auto &s: map | std::views::keys) {
                    const double v = dydt.at(s);
                    if (!std::isfinite(v)) {
                        invalid_species.emplace_back(s, v);
                    }
                }
                std::string msg = std::format("Non-finite dydt values encountered for species: {}",
                    [&invalid_species]() -> std::string {
                        std::ostringstream ss;
                        size_t count = 0;
                        for (const auto& [sp, val] : invalid_species) {
                            ss << sp.name() << ": " << val;
                            if (count < invalid_species.size() - 1) {
                                ss << ", ";
                            }
                            count++;
                        }
                        return ss.str();
                    }());
                LOG_CRITICAL(getLogger(), "{}", msg);
                throw exceptions::InvalidQSESolutionError(msg);
            }
        }
#endif

        for (const auto& [species, index] : map) {
            const double val = dydt.at(species);
            f_data[index] = val;
        }

        return 0; // Success
    }


    int MultiscalePartitioningEngineView::QSESolver::sys_jac(
        const N_Vector y,
        N_Vector fy,
        const SUNMatrix J,
        void *user_data,
        N_Vector tmp1,
        N_Vector tmp2
    ) {
        auto* data = static_cast<UserData*>(user_data);

        const sunrealtype* y_data = N_VGetArrayPointer(y);
        const auto& map = data->qse_solve_species_index_map;
        for (const auto& [species, index] : map) {
            data->comp.setMolarAbundance(species, y_data[index]);
        }

        const NetworkJacobian jac = data->engine.generateJacobianMatrix(
            data->ctx,
            data->comp,
            data->T9,
            data->rho,
            data->qse_solve_species
        );

        sunrealtype* J_data = SUNDenseMatrix_Data(J);
        const sunindextype N = SUNDenseMatrix_Columns(J);


        for (const auto& [row_species, row_idx]: map) {
            double max_value = std::numeric_limits<double>::lowest();
            for (const auto &col_species: map | std::views::keys) {
                const double val = std::abs(jac(row_species, col_species));
                if (val > max_value) {
                    max_value = val;
                }
            }

            for (const auto& [col_species, col_idx] : map) {
                J_data[col_idx * N + row_idx] = jac(row_species, col_species);
            }
        }

        LOG_TRACE_L2( getLogger(), "After filling, Jacobian in sys_jac is: {}",
            [&]() -> std::string {
                std::ostringstream oss;
                const sunrealtype* J_log_data = SUNDenseMatrix_Data(J);
                const sunindextype N_log = SUNDenseMatrix_Columns(J);
                oss << "[";
                for (size_t i = 0; i < data->qse_solve_species.size(); ++i) {
                    oss << "[";
                    for (size_t j = 0; j < data->qse_solve_species.size(); ++j) {
                        oss << J_log_data[i * N_log + j];
                        if (j < data->qse_solve_species.size() - 1) {
                            oss << ", ";
                        }
                    }
                    oss << "]";
                    if (i < data->qse_solve_species.size() - 1) {
                        oss << ", ";
                    }
                }
                oss << "]";
                return oss.str();
            }()
        );

        data->instance.m_has_jacobian = true;

        return 0;
    }

      /////////////////////////////////
     /// QSEGroup Member Functions ///
    ////////////////////////////////



    bool MultiscalePartitioningEngineView::QSEGroup::operator==(const QSEGroup &other) const {
       return mean_timescale == other.mean_timescale;
    }

    bool MultiscalePartitioningEngineView::QSEGroup::operator<(const QSEGroup &other) const {
        return mean_timescale < other.mean_timescale;
    }

    bool MultiscalePartitioningEngineView::QSEGroup::operator>(const QSEGroup &other) const {
        return mean_timescale > other.mean_timescale;
    }

    bool MultiscalePartitioningEngineView::QSEGroup::operator!=(const QSEGroup &other) const {
        return !(*this == other);
    }

    std::string MultiscalePartitioningEngineView::QSEGroup::toString(const bool verbose) const {
        if (verbose) {
            return std::format(
                "QSEGroup(Algebraic: [{:40}], Seed [{:40}], Mean Timescale: {:10.4E}, Is in Equilibrium: {:6}",
                utils::iterable_to_delimited_string(algebraic_species),
                utils::iterable_to_delimited_string(seed_species),
                mean_timescale,
                is_in_equilibrium ? "True" : "False"
            );
        }
        return std::format("QSEGroup({})", utils::iterable_to_delimited_string(algebraic_species));
    }

}
