GridFire/src/lib/engine/procedures/construction.cpp

#include "gridfire/engine/procedures/construction.h"
#include "gridfire/reaction/weak/weak_interpolator.h"
#include "gridfire/reaction/weak/weak.h"
#include "gridfire/reaction/weak/weak_types.h"

#include <ranges>
#include <stdexcept>
#include <memory>

#include "gridfire/reaction/reaction.h"
#include "gridfire/reaction/reaclib.h"

#include "fourdst/composition/composition.h"

#include "fourdst/logging/logging.h"

#include "quill/Logger.h"
#include "quill/LogMacros.h"

namespace gridfire {
    using reaction::ReactionSet;
    using reaction::Reaction;
    using fourdst::composition::Composition;
    using fourdst::atomic::Species;


    ReactionSet build_nuclear_network(
        const Composition& composition,
        const rates::weak::WeakRateInterpolator& weakInterpolator,
        BuildDepthType maxLayers,
        bool reverse
    ) {
        int depth;
        if (std::holds_alternative<NetworkBuildDepth>(maxLayers)) {
            depth = static_cast<int>(std::get<NetworkBuildDepth>(maxLayers));
        } else {
            depth = std::get<int>(maxLayers);
        }
        auto logger = fourdst::logging::LogManager::getInstance().getLogger("log");
        if (depth == 0) {
            LOG_ERROR(logger, "Network build depth is set to 0. No reactions will be collected.");
            throw std::logic_error("Network build depth is set to 0. No reactions will be collected.");
        }

        // --- Step 1: Create a single master pool that owns all possible reactions ---
        ReactionSet master_reaction_pool;

        // Clone all relevant REACLIB reactions into the master pool
        const auto& allReaclibReactions = reaclib::get_all_reaclib_reactions();
        for (const auto& reaction : allReaclibReactions) {
            if (reaction->is_reverse() == reverse) {
                master_reaction_pool.add_reaction(reaction->clone());
            }
        }

        // --- Clone all possible weak reactions into the master reaction pool ---
        // for (const auto& parent_species: weakInterpolator.available_isotopes()) {
        //     std::expected<Species, fourdst::atomic::SpeciesErrorType> upProduct = fourdst::atomic::az_to_species(
        //         parent_species.a(),
        //         parent_species.z() + 1
        //     );
        //     std::expected<Species, fourdst::atomic::SpeciesErrorType> downProduct = fourdst::atomic::az_to_species(
        //         parent_species.a(),
        //         parent_species.z() - 1
        //     );
        //     if (downProduct.has_value()) { // Only add the reaction if the Species map contains the product
        //         master_reaction_pool.add_reaction(
        //             std::make_unique<rates::weak::WeakReaction>(
        //                 parent_species,
        //                 rates::weak::WeakReactionType::BETA_PLUS_DECAY,
        //                 weakInterpolator
        //             )
        //         );
        //         master_reaction_pool.add_reaction(
        //             std::make_unique<rates::weak::WeakReaction>(
        //                 parent_species,
        //                 rates::weak::WeakReactionType::ELECTRON_CAPTURE,
        //                 weakInterpolator
        //             )
        //         );
        //     }
        //     if (upProduct.has_value()) { // Only add the reaction if the Species map contains the product
        //         master_reaction_pool.add_reaction(
        //             std::make_unique<rates::weak::WeakReaction>(
        //                 parent_species,
        //                 rates::weak::WeakReactionType::BETA_MINUS_DECAY,
        //                 weakInterpolator
        //             )
        //         );
        //         master_reaction_pool.add_reaction(
        //             std::make_unique<rates::weak::WeakReaction>(
        //                 parent_species,
        //                 rates::weak::WeakReactionType::POSITRON_CAPTURE,
        //                 weakInterpolator
        //             )
        //         );
        //     }
        // } TODO: Remove comments, weak reactions have been disabled for testing

        // --- Step 2: Use non-owning raw pointers for the fast build algorithm ---
        std::vector<Reaction*> remainingReactions;
        remainingReactions.reserve(master_reaction_pool.size());
        for(const auto& reaction : master_reaction_pool) {
            remainingReactions.push_back(reaction.get());
        }

        if (depth == static_cast<int>(NetworkBuildDepth::Full)) {
            LOG_INFO(logger, "Building full nuclear network with a total of {} reactions.", remainingReactions.size());
            return master_reaction_pool; // No need to build layer by layer if we want the full network
        }

        // --- Step 3: Execute the layered network build using observing pointers ---
        std::unordered_set<Species> availableSpecies;
        for (const auto& entry : composition | std::views::values) {
            if (entry.mass_fraction() > 0.0) {
                availableSpecies.insert(entry.isotope());
            }
        }

        std::vector<Reaction*> collectedReactionPtrs;
        LOG_INFO(logger, "Starting network construction with {} available species.", availableSpecies.size());

        for (int layer = 0; layer < depth && !remainingReactions.empty(); ++layer) {
            size_t collectedThisLayer = 0;
            size_t collectedStrong = 0;
            size_t collectedWeak = 0;
            LOG_TRACE_L1(logger, "Collecting reactions for layer {} with {} remaining reactions. Currently there are {} available species", layer, remainingReactions.size(), availableSpecies.size());
            std::vector<Reaction*> reactionsForNextPass;
            std::unordered_set<Species> newProductsThisLayer;
            bool newReactionsAdded = false;

            reactionsForNextPass.reserve(remainingReactions.size());

            for (Reaction* reaction : remainingReactions) {
                bool allReactantsAvailable = true;
                for (const auto& reactant : reaction->reactants()) {
                    if (!availableSpecies.contains(reactant)) {
                        allReactantsAvailable = false;
                        break;
                    }
                }

                if (allReactantsAvailable) {
                    collectedReactionPtrs.push_back(reaction);
                    if (reaction->type() == reaction::ReactionType::WEAK) {
                        collectedWeak++;
                    } else {
                        collectedStrong++;
                    }
                    collectedThisLayer++;
                    newReactionsAdded = true;

                    for (const auto& product : reaction->products()) {
                        newProductsThisLayer.insert(product);
                    }
                } else {
                    reactionsForNextPass.push_back(reaction);
                }
            }

            if (!newReactionsAdded) {
                LOG_INFO(logger, "No new reactions added in layer {}. Stopping network construction.", layer);
                break;
            }

            size_t oldProductCount = availableSpecies.size();
            availableSpecies.insert(newProductsThisLayer.begin(), newProductsThisLayer.end());
            size_t newProductCount = availableSpecies.size() - oldProductCount;
            LOG_TRACE_L1(
                logger,
                "Layer {}: Collected {} new reactions ({} strong, {} weak). New products this layer: {}",
                layer,
                collectedThisLayer,
                collectedStrong,
                collectedWeak,
                newProductCount
            );
            remainingReactions = std::move(reactionsForNextPass);
        }

        // --- Step 4: Construct the final ReactionSet by moving ownership ---
        LOG_INFO(logger, "Network construction completed. Assembling final set of {} reactions.", collectedReactionPtrs.size());
        ReactionSet finalReactionSet;

        std::unordered_set<Reaction*> collectedPtrSet(collectedReactionPtrs.begin(), collectedReactionPtrs.end());

        for (auto& reaction_ptr : master_reaction_pool) {
            if (reaction_ptr && collectedPtrSet.contains(reaction_ptr.get())) {
                finalReactionSet.add_reaction(std::move(reaction_ptr));
            }
        }

        return finalReactionSet;
    }


}