GridFire/src/include/gridfire/reaction/reaction.h

#pragma once

#include <string_view>

#include "fourdst/composition/atomicSpecies.h"
#include "fourdst/logging/logging.h"
#include "quill/Logger.h"
#include <unordered_map>
#include <vector>
#include <unordered_set>


#include "cppad/cppad.hpp"

/**
 * @file reaction.h
 * @brief Defines classes for representing and managing nuclear reactions.
 *
 * This file contains the core data structures for handling nuclear reactions,
 * including individual reactions from specific sources (`Reaction`), collections
 * of reactions (`ReactionSet`), and logical reactions that aggregate rates from
* multiple sources (`LogicalReaclibReaction`, `LogicalReactionSet`).
 */
namespace gridfire::reaction {
    enum class ReactionType {
        WEAK,
        REACLIB,
        LOGICAL_REACLIB,
    };
    /**
     * @struct RateCoefficientSet
     * @brief Holds the seven coefficients for the REACLIB rate equation.
     *
     * This struct stores the parameters (a0-a6) used to calculate reaction rates
     * as a function of temperature.
     */
    struct RateCoefficientSet {
        double a0; ///< Coefficient a0
        double a1; ///< Coefficient a1
        double a2; ///< Coefficient a2
        double a3; ///< Coefficient a3
        double a4; ///< Coefficient a4
        double a5; ///< Coefficient a5
        double a6; ///< Coefficient a6

        /**
         * @brief Overloads the stream insertion operator for easy printing.
         * @param os The output stream.
         * @param r The RateCoefficientSet to print.
         * @return The output stream.
         */
        friend std::ostream& operator<<(std::ostream& os, const RateCoefficientSet& r) {
            os << "[" << r.a0 << ", " << r.a1 << ", " << r.a2 << ", "
               << r.a3 << ", " << r.a4 << ", " << r.a5 << ", " << r.a6 << "]";
            return os;
        }
    };

    /**
     * @class Reaction
     * @brief Represents a single nuclear reaction from a specific data source.
     *
     * This class encapsulates all properties of a single nuclear reaction as defined
     * in formats like REACLIB, including reactants, products, Q-value, and rate
     * coefficients from a particular evaluation (source).
     *
     * Example:
     * @code
     * // Assuming species and rate coefficients are defined
     * Reaction p_gamma_d(
     *     "H_1_H_1_to_H_2", "p(p,g)d", 1, {H_1, H_1}, {H_2}, 5.493, "st08", rate_coeffs
     * );
     * double rate = p_gamma_d.calculate_rate(0.1); // T9 = 0.1
     * @endcode
     */

    class Reaction {
    public:
        virtual ~Reaction() = default;

        [[nodiscard]] virtual double calculate_rate(double T9, double rho, const std::vector<double>& Y) const = 0;
        [[nodiscard]] virtual CppAD::AD<double> calculate_rate(CppAD::AD<double> T9, CppAD::AD<double> rho, const std::vector<CppAD::AD<double>>& Y) const = 0;

        [[nodiscard]] virtual std::string_view id() const = 0;
        [[nodiscard]] virtual const std::vector<fourdst::atomic::Species>& reactants() const = 0;
        [[nodiscard]] virtual const std::vector<fourdst::atomic::Species>& products() const = 0;

        [[nodiscard]] virtual bool contains(const fourdst::atomic::Species& species) const = 0;
        [[nodiscard]] virtual bool contains_reactant(const fourdst::atomic::Species& species) const = 0;
        [[nodiscard]] virtual bool contains_product(const fourdst::atomic::Species& species) const = 0;

        [[nodiscard]] virtual bool is_reverse() const = 0;

        [[nodiscard]] virtual std::unordered_set<fourdst::atomic::Species> all_species() const = 0;
        [[nodiscard]] virtual std::unordered_set<fourdst::atomic::Species> reactant_species() const = 0;
        [[nodiscard]] virtual std::unordered_set<fourdst::atomic::Species> product_species() const = 0;

        [[nodiscard]] virtual size_t num_species() const = 0;

        [[nodiscard]] virtual std::unordered_map<fourdst::atomic::Species, int> stoichiometry() const = 0;
        [[nodiscard]] virtual int stoichiometry(const fourdst::atomic::Species& species) const = 0;

        [[nodiscard]] virtual uint64_t hash(uint64_t seed) const = 0;

        [[nodiscard]] virtual double qValue() const = 0;
        [[nodiscard]] virtual double calculate_forward_rate_log_derivative(double T9, double rho, const std::vector<double>& Y) const = 0;

        [[nodiscard]] virtual ReactionType type() const = 0;

        [[nodiscard]] virtual std::unique_ptr<Reaction> clone() const = 0;

        friend std::ostream& operator<<(std::ostream& os, const Reaction& r) {
            os << "Reaction(ID: " << r.id() << ")";
            return os;
        }
    };
    class ReaclibReaction : public Reaction {
    public:
        ~ReaclibReaction() override = default;

        /**
         * @brief Constructs a Reaction object.
         * @param id A unique identifier for the reaction.
         * @param peName The name in (projectile, ejectile) notation (e.g., "p(p,g)d").
         * @param chapter The REACLIB chapter number, defining reaction structure.
         * @param reactants A vector of reactant species.
         * @param products A vector of product species.
         * @param qValue The Q-value of the reaction in MeV.
         * @param label The source label for the rate data (e.g., "wc12", "st08").
         * @param sets The set of rate coefficients.
         * @param reverse True if this is a reverse reaction rate.
         */
        ReaclibReaction(
            const std::string_view id,
            const std::string_view peName,
            const int chapter,
            const std::vector<fourdst::atomic::Species> &reactants,
            const std::vector<fourdst::atomic::Species> &products,
            const double qValue,
            const std::string_view label,
            const RateCoefficientSet &sets,
            const bool reverse = false);

        /**
         * @brief Calculates the reaction rate for a given temperature.
         * @param T9 The temperature in units of 10^9 K.
         * @param rho Density [Not used in this implementation].
         * @param Y Molar abundances of species [Not used in this implementation].
         * @return The calculated reaction rate.
         */
        [[nodiscard]] double calculate_rate(double T9, double rho, const std::vector<double>& Y) const override;

        /**
         * @brief Calculates the reaction rate for a given temperature using CppAD types.
         * @param T9 The temperature in units of 10^9 K, as a CppAD::AD<double>.
         * @param rho Density, as a CppAD::AD<double> [Not used in this implementation].
         * @param Y Molar abundances of species, as a vector of CppAD::AD<double> [Not used in this implementation].
         * @return The calculated reaction rate, as a CppAD::AD<double>.
         */
        [[nodiscard]] CppAD::AD<double> calculate_rate(CppAD::AD<double> T9, CppAD::AD<double> rho, const std::vector<CppAD::AD<double>>& Y) const override;

        [[nodiscard]] double calculate_forward_rate_log_derivative(double T9, double rho, const std::vector<double>& Y) const override;

        /**
         * @brief Gets the reaction name in (projectile, ejectile) notation.
         * @return The reaction name (e.g., "p(p,g)d").
         */
        [[nodiscard]] virtual std::string_view peName() const { return m_peName; }

        /**
         * @brief Gets the REACLIB chapter number.
         * @return The chapter number.
         */
        [[nodiscard]] int chapter() const { return m_chapter; }

        /**
         * @brief Gets the source label for the rate data.
         * @return The source label (e.g., "wc12w", "st08").
         */
        [[nodiscard]] std::string_view sourceLabel() const { return m_sourceLabel; }

        [[nodiscard]] ReactionType type() const override { return ReactionType::REACLIB; }

        /**
         * @brief Gets the set of rate coefficients.
         * @return A const reference to the RateCoefficientSet.
         */
        [[nodiscard]] const RateCoefficientSet& rateCoefficients() const { return m_rateCoefficients; }

        /**
         * @brief Checks if the reaction involves a given species as a reactant or product.
         * @param species The species to check for.
         * @return True if the species is involved, false otherwise.
         */
        [[nodiscard]] bool contains(const fourdst::atomic::Species& species) const override;

        /**
         * @brief Checks if the reaction involves a given species as a reactant.
         * @param species The species to check for.
         * @return True if the species is a reactant, false otherwise.
         */
        [[nodiscard]] bool contains_reactant(const fourdst::atomic::Species& species) const override;

        /**
         * @brief Checks if the reaction involves a given species as a product.
         * @param species The species to check for.
         * @return True if the species is a product, false otherwise.
         */
        [[nodiscard]] bool contains_product(const fourdst::atomic::Species& species) const override;

        /**
         * @brief Gets a set of all unique species involved in the reaction.
         * @return An unordered_set of all reactant and product species.
         */
        [[nodiscard]] std::unordered_set<fourdst::atomic::Species> all_species() const override;

        /**
         * @brief Gets a set of all unique reactant species.
         * @return An unordered_set of reactant species.
         */
        [[nodiscard]] std::unordered_set<fourdst::atomic::Species> reactant_species() const override;

        /**
         * @brief Gets a set of all unique product species.
         * @return An unordered_set of product species.
         */
        [[nodiscard]] std::unordered_set<fourdst::atomic::Species> product_species() const override;

        /**
         * @brief Gets the number of unique species involved in the reaction.
         * @return The count of unique species.
         */
        [[nodiscard]] size_t num_species() const override;

        /**
         * @brief Calculates the stoichiometric coefficient for a given species.
         * @param species The species for which to find the coefficient.
         * @return The stoichiometric coefficient (negative for reactants, positive for products).
         */
        [[nodiscard]] int stoichiometry(const fourdst::atomic::Species& species) const override;

        /**
         * @brief Gets a map of all species to their stoichiometric coefficients.
         * @return An unordered_map from species to their integer coefficients.
         */
        [[nodiscard]] std::unordered_map<fourdst::atomic::Species, int> stoichiometry() const override;

        /**
         * @brief Gets the unique identifier of the reaction.
         * @return The reaction ID.
         */
        [[nodiscard]] std::string_view id() const override { return m_id; }

        /**
         * @brief Gets the Q-value of the reaction.
         * @return The Q-value in whatever units the reaction was defined in (usually MeV).
         */
        [[nodiscard]] double qValue() const override { return m_qValue; }

        /**
         * @brief Gets the vector of reactant species.
         * @return A const reference to the vector of reactants.
         */
        [[nodiscard]] const std::vector<fourdst::atomic::Species>& reactants() const override { return m_reactants; }

        /**
         * @brief Gets the vector of product species.
         * @return A const reference to the vector of products.
         */
        [[nodiscard]] const std::vector<fourdst::atomic::Species>& products() const override { return m_products; }

        /**
         * @brief Checks if this is a reverse reaction rate.
         * @return True if it is a reverse rate, false otherwise.
         */
        [[nodiscard]] bool is_reverse() const override { return m_reverse; }

        /**
         * @brief Calculates the excess energy from the mass difference of reactants and products.
         * @return The excess energy in MeV.
         */
        [[nodiscard]] double excess_energy() const;

        /**
         * @brief Compares this reaction with another for equality based on their IDs.
         * @param other The other Reaction to compare with.
         * @return True if the reaction IDs are the same.
         */
        bool operator==(const ReaclibReaction& other) const { return m_id == other.m_id; }

        /**
         * @brief Compares this reaction with another for inequality.
         * @param other The other Reaction to compare with.
         * @return True if the reactions are not equal.
         */
        bool operator!=(const ReaclibReaction& other) const { return !(*this == other); }

        /**
         * @brief Computes a hash for the reaction based on its ID.
         * @param seed The seed for the hash function.
         * @return A 64-bit hash value.
         * @details Uses the XXHash64 algorithm on the reaction's ID string.
         */
        [[nodiscard]] uint64_t hash(uint64_t seed) const override;

        [[nodiscard]] std::unique_ptr<Reaction> clone() const override;

        friend std::ostream& operator<<(std::ostream& os, const ReaclibReaction& r) {
            return os << "(ReaclibReaction:" << r.m_id << ")";
        }

    protected:
        quill::Logger* m_logger = fourdst::logging::LogManager::getInstance().getLogger("log");
        std::string m_id; ///< Unique identifier for the reaction (e.g., "h1+h1=>h2+e+nu").
        std::string m_peName; ///< Name of the reaction in (projectile, ejectile) notation (e.g. "p(p,g)d").
        int m_chapter;    ///< Chapter number from the REACLIB database, defining the reaction structure.
        double m_qValue = 0.0; ///< Q-value of the reaction in MeV.
        std::vector<fourdst::atomic::Species> m_reactants; ///< Reactants of the reaction.
        std::vector<fourdst::atomic::Species> m_products; ///< Products of the reaction.
        std::string m_sourceLabel; ///< Source label for the rate data (e.g., "wc12w", "st08").
        RateCoefficientSet m_rateCoefficients; ///< The seven rate coefficients.
        bool m_reverse = false; ///< Flag indicating if this is a reverse reaction rate.
    private:
        /**
         * @brief Template implementation for calculating the reaction rate.
         * @tparam T The numeric type (double or CppAD::AD<double>).
         * @param T9 The temperature in units of 10^9 K.
         * @return The calculated reaction rate.
         * @details The rate is calculated using the standard REACLIB formula:
         * `rate = exp(a0 + a1/T9 + a2/T9^(1/3) + a3*T9^(1/3) + a4*T9 + a5*T9^(5/3) + a6*ln(T9))`
         */
        template <typename T>
        [[nodiscard]] T calculate_rate(const T T9) const {
            const T T913 = CppAD::pow(T9, 1.0/3.0);
            const T T953 = CppAD::pow(T9, 5.0/3.0);
            const T logT9 = CppAD::log(T9);
            const T exponent = m_rateCoefficients.a0 +
                   m_rateCoefficients.a1 / T9 +
                   m_rateCoefficients.a2 / T913 +
                   m_rateCoefficients.a3 * T913 +
                   m_rateCoefficients.a4 * T9 +
                   m_rateCoefficients.a5 * T953 +
                   m_rateCoefficients.a6 * logT9;
            return CppAD::exp(exponent);
        }
    };


    /**
     * @class LogicalReaclibReaction
     * @brief Represents a "logical" reaction that aggregates rates from multiple sources.
     *
     * A LogicalReaclibReaction shares the same reactants and products but combines rates
     * from different evaluations (e.g., "wc12" and "st08" for the same physical
     * reaction). The total rate is the sum of the individual rates.
     * It inherits from Reaction, using the properties of the first provided reaction
     * as its base properties (reactants, products, Q-value, etc.).
     */
    class LogicalReaclibReaction final : public ReaclibReaction {
    public:
        /**
         * @brief Constructs a LogicalReaction from a vector of `Reaction` objects.
         * @param reactions A vector of reactions that represent the same logical process.
         * @throws std::runtime_error if the provided reactions have inconsistent Q-values.
         */
        explicit LogicalReaclibReaction(const std::vector<ReaclibReaction> &reactions);

        /**
         * @brief Adds another `Reaction` source to this logical reaction.
         * @param reaction The reaction to add.
         * @throws std::runtime_error if the reaction has a different `peName`, a duplicate
         *         source label, or an inconsistent Q-value.
         */
        void add_reaction(const ReaclibReaction& reaction);

        /**
         * @brief Gets the number of source rates contributing to this logical reaction.
         * @return The number of aggregated rates.
         */
        [[nodiscard]] size_t size() const { return m_rates.size(); }

        /**
         * @brief Gets the list of source labels for the aggregated rates.
         * @return A vector of source label strings.
         */
        [[nodiscard]] std::vector<std::string> sources() const { return m_sources; }

        /**
         * @brief Calculates the total reaction rate by summing all source rates.
         * @param T9 The temperature in units of 10^9 K.
         * @param rho
         * @param Y
         * @return The total calculated reaction rate.
         */
        [[nodiscard]] double calculate_rate(double T9, double rho, const std::vector<double>& Y) const override;

        [[nodiscard]] double calculate_forward_rate_log_derivative(double T9, double rho, const std::vector<double>& Y) const override;

        [[nodiscard]] ReactionType type() const override { return ReactionType::LOGICAL_REACLIB; }

        [[nodiscard]] std::unique_ptr<Reaction> clone() const override;

        /**
         * @brief Calculates the total reaction rate using CppAD types.
         * @param T9 The temperature in units of 10^9 K, as a CppAD::AD<double>.
         * @param rho
         * @param Y
         * @return The total calculated reaction rate, as a CppAD::AD<double>.
         */
        [[nodiscard]] CppAD::AD<double> calculate_rate(CppAD::AD<double> T9, CppAD::AD<double> rho, const std::vector<CppAD::AD<double>>& Y) const override;

        /** @name Iterators
         *  Provides iterators to loop over the rate coefficient sets.
         */
        ///@{
        auto begin() { return m_rates.begin(); }
        [[nodiscard]] auto begin() const { return m_rates.cbegin(); }
        auto end() { return m_rates.end(); }
        [[nodiscard]] auto end() const { return m_rates.cend(); }
        ///@}
        ///

        friend std::ostream& operator<<(std::ostream& os, const LogicalReaclibReaction& r) {
            os << "(LogicalReaclibReaction: " << r.id() << ", reverse: " << r.is_reverse() << ")";
            return os;
        }

    private:
        std::vector<std::string> m_sources; ///< List of source labels.
        std::vector<RateCoefficientSet> m_rates; ///< List of rate coefficient sets from each source.

    private:
        /**
         * @brief Template implementation for calculating the total reaction rate.
         * @tparam T The numeric type (double or CppAD::AD<double>).
         * @param T9 The temperature in units of 10^9 K.
         * @return The total calculated reaction rate.
         * @details This method iterates through all stored `RateCoefficientSet`s,
         * calculates the rate for each, and returns their sum.
         */
        template <typename T>
        [[nodiscard]] T calculate_rate(const T T9) const {
            T sum = static_cast<T>(0.0);
            const T T913 = CppAD::pow(T9, 1.0/3.0);
            const T T953 = CppAD::pow(T9, 5.0/3.0);
            const T logT9 = CppAD::log(T9);
            // ReSharper disable once CppUseStructuredBinding
            for (const auto& rate : m_rates) {
                const T exponent = rate.a0 +
                       rate.a1 / T9 +
                       rate.a2 / T913 +
                       rate.a3 * T913 +
                       rate.a4 * T9 +
                       rate.a5 * T953 +
                       rate.a6 * logT9;
                sum += CppAD::exp(exponent);
            }
            return sum;
        }
    };

    class ReactionSet final {
    public:
        /**
         * @brief Constructs a ReactionSet from a vector of reactions.
         * @param reactions The initial vector of Reaction objects.
         */
        explicit ReactionSet(std::vector<std::unique_ptr<Reaction>>&& reactions);

        explicit ReactionSet(const std::vector<Reaction*>& reactions);

        ReactionSet();

        /**
         * @brief Copy constructor.
         * @param other The ReactionSet to copy.
         */
        ReactionSet(const ReactionSet& other);

        /**
         * @brief Copy assignment operator.
         * @param other The ReactionSet to assign from.
         * @return A reference to this ReactionSet.
         */
        ReactionSet& operator=(const ReactionSet& other);

        /**
         * @brief Adds a reaction to the set.
         * @param reaction The Reaction to add.
         */
        void add_reaction(const Reaction& reaction);

        void add_reaction(std::unique_ptr<Reaction>&& reaction);

        /**
         * @brief Removes a reaction from the set.
         * @param reaction The Reaction to remove.
         */
        void remove_reaction(const Reaction& reaction);

        /**
         * @brief Checks if the set contains a reaction with the given ID.
         * @param id The ID of the reaction to find.
         * @return True if the reaction is in the set, false otherwise.
         */
        [[nodiscard]] bool contains(const std::string_view& id) const;

        /**
         * @brief Checks if the set contains the given reaction.
         * @param reaction The Reaction to find.
         * @return True if the reaction is in the set, false otherwise.
         */
        [[nodiscard]] bool contains(const Reaction& reaction) const;

        /**
         * @brief Gets the number of reactions in the set.
         * @return The size of the set.
         */
        [[nodiscard]] size_t size() const { return m_reactions.size(); }

        /**
         * @brief Removes all reactions from the set.
         */
        void clear();

        /**
         * @brief Checks if any reaction in the set involves the given species.
         * @param species The species to check for.
         * @return True if the species is involved in any reaction.
         */
        [[nodiscard]] bool contains_species(const fourdst::atomic::Species& species) const;

        /**
         * @brief Checks if any reaction in the set contains the given species as a reactant.
         * @param species The species to check for.
         * @return True if the species is a reactant in any reaction.
         */
        [[nodiscard]] bool contains_reactant(const fourdst::atomic::Species& species) const;

        /**
         * @brief Checks if any reaction in the set contains the given species as a product.
         * @param species The species to check for.
         * @return True if the species is a product in any reaction.
         */
        [[nodiscard]] bool contains_product(const fourdst::atomic::Species& species) const;

        /**
         * @brief Accesses a reaction by its index.
         * @param index The index of the reaction to access.
         * @return A const reference to the Reaction.
         * @throws std::out_of_range if the index is out of bounds.
         */
        [[nodiscard]] const Reaction& operator[](size_t index) const;

        /**
         * @brief Accesses a reaction by its ID.
         * @param id The ID of the reaction to access.
         * @return A const reference to the Reaction.
         * @throws std::out_of_range if no reaction with the given ID exists.
         */
        [[nodiscard]] const Reaction& operator[](const std::string_view& id) const;

        /**
         * @brief Compares this set with another for equality.
         * @param other The other ReactionSet to compare with.
         * @return True if the sets are equal (same size and hash).
         */
        bool operator==(const ReactionSet& other) const;

        /**
         * @brief Compares this set with another for inequality.
         * @param other The other ReactionSet to compare with.
         * @return True if the sets are not equal.
         */
        bool operator!=(const ReactionSet& other) const;

        /**
         * @brief Computes a hash for the entire set.
         * @param seed The seed for the hash function.
         * @return A 64-bit hash value.
         * @details The algorithm computes the hash of each individual reaction,
         * sorts the hashes, and then computes a final hash over the sorted list
         * of hashes. This ensures the hash is order-independent.
         */
        [[nodiscard]] uint64_t hash(uint64_t seed) const;

        /** @name Iterators
         *  Provides iterators to loop over the reactions in the set.
         */
        ///@{
        auto begin() { return m_reactions.begin(); }
        [[nodiscard]] auto begin() const { return m_reactions.cbegin(); }
        auto end() { return m_reactions.end(); }
        [[nodiscard]] auto end() const { return m_reactions.cend(); }
        ///@}
        ///

        [[nodiscard]] std::unordered_set<fourdst::atomic::Species> getReactionSetSpecies() const;

        friend std::ostream& operator<<(std::ostream& os, const ReactionSet& rs) {
            os << "(ReactionSet: {";
            int i = 0;
            for (const auto& reaction : rs.m_reactions) {
                os << *reaction;
                if (i < rs.m_reactions.size() - 1) {
                    os << ", ";
                }
            }
            os << "})";
            return os;
        }
    private:
        quill::Logger* m_logger = fourdst::logging::LogManager::getInstance().getLogger("log");
        std::vector<std::unique_ptr<Reaction>> m_reactions;
        std::string m_id;
        std::unordered_map<std::string, size_t> m_reactionNameMap; ///< Maps reaction IDs to Reaction objects for quick lookup.

    };

    ReactionSet packReactionSet(const ReactionSet& reactionSet);

}