#pragma once
#include <algorithm>
#include <utility>
#include <vector>
#include <string>
#include <iostream>
#include <unordered_map>
#include "atomicSpecies.h"

#include "cppad/cppad.hpp"

namespace gridfire::reaclib {
    /**
     * @struct RateFitSet
     * @brief Holds the seven fitting parameters for a single REACLIB rate set.
     * @details The thermonuclear reaction rate for a single set is calculated as:
     * rate = exp(a0 + a1/T9 + a2/T9^(-1/3) + a3*T9^(1/3) + a4*T9 + a5*T9^(5/3) + a6*ln(T9))
     * where T9 is the temperature in billions of Kelvin. The total rate for a
     * reaction is the sum of the rates from all its sets.
     */
    struct RateFitSet {
        double a0, a1, a2, a3, a4, a5, a6;
    };

    /**
     * @struct REACLIBReaction
     * @brief Represents a single nuclear reaction from the JINA REACLIB database.
     * @details This struct is designed to be constructed at compile time (constexpr) from
     * the data parsed by the Python generation script. It stores all necessary
     * information to identify a reaction and calculate its rate.
     */
    class REACLIBReaction {
    private:
        std::string m_id; ///< Unique identifier for the reaction, generated by the Python script.
        int m_chapter;    ///< Chapter number from the REACLIB database, defining the reaction structure.
        std::vector<fourdst::atomic::Species> m_reactantNames; ///< Names of the reactant species involved in the reaction.
        std::vector<fourdst::atomic::Species> m_productNames;  ///< Names of the product species produced by the reaction.
        double m_qValue; ///< Q-value of the reaction in MeV, representing the energy released or absorbed.
        std::string m_sourceLabel; ///< Source label for the rate data, indicating the origin of the rate coefficients (e.g., "wc12w", "st08").
        RateFitSet m_rateSets; ///< Fitting coefficients for the reaction rate, containing the seven parameters used in the rate calculation.
        bool m_reverse = false; ///< indicates if the reaction is reversed

        friend bool operator==(const REACLIBReaction& lhs, const REACLIBReaction& rhs);
        friend bool operator!=(const REACLIBReaction& lhs, const REACLIBReaction& rhs);
    public:
        /**
         * @brief Constructs a REACLIBReaction object at compile time.
         * @param id A unique string identifier generated by the Python script.
         * @param chapter The REACLIB chapter number, defining the reaction structure.
         * @param reactants A vector of strings with the names of the reactant species.
         * @param products A vector of strings with the names of the product species.
         * @param qValue The Q-value of the reaction in MeV.
         * @param label The source label for the rate data (e.g., "wc12w", "st08").
         * @param sets A vector of RateFitSet, containing the fitting coefficients for the rate.
         * @param reverse A boolean indicating if the reaction is reversed (default is false).
         */
        REACLIBReaction(
            std::string id,
            int chapter,
            std::vector<fourdst::atomic::Species> reactants,
            std::vector<fourdst::atomic::Species> products,
            double qValue,
            std::string label,
            RateFitSet sets,
            bool reverse = false)
            : m_id(std::move(id)),
              m_chapter(chapter),
              m_reactantNames(std::move(reactants)),
              m_productNames(std::move(products)),
              m_qValue(qValue),
              m_sourceLabel(std::move(label)),
              m_rateSets(sets),
              m_reverse(reverse) {}

        template <typename GeneralScalarType>
        [[nodiscard]] GeneralScalarType calculate_rate(const GeneralScalarType T9) const {
            const GeneralScalarType T913 = CppAD::pow(T9, 1.0/3.0);
            const GeneralScalarType rateExponent = m_rateSets.a0 +
                m_rateSets.a1 / T9 +
                m_rateSets.a2 / T913 +
                m_rateSets.a3 * T913 +
                m_rateSets.a4 * T9 +
                m_rateSets.a5 * CppAD::pow(T9, 5.0/3.0) +
                m_rateSets.a6 * CppAD::log(T9);
            return CppAD::exp(rateExponent);
        }

        [[nodiscard]] const std::string& id() const { return m_id; }

        [[nodiscard]] int chapter() const { return m_chapter; }

        [[nodiscard]] const std::vector<fourdst::atomic::Species>& reactants() const { return m_reactantNames; }

        [[nodiscard]] const std::vector<fourdst::atomic::Species>& products() const { return m_productNames; }

        [[nodiscard]] double qValue() const { return m_qValue; }

        [[nodiscard]] std::string sourceLabel() const { return m_sourceLabel; }

        [[nodiscard]] bool is_reverse() const { return m_reverse; }

        friend std::ostream& operator<<(std::ostream& os, const REACLIBReaction& reaction) {
            os << "REACLIBReaction(" << reaction.m_id << ", "
               << "Chapter: " << reaction.m_chapter << ")";
            return os;
        }

        friend bool operator==(const REACLIBReaction& lhs, const REACLIBReaction& rhs);

        friend bool operator!=(const REACLIBReaction& lhs, const REACLIBReaction& rhs);
    };

     class REACLIBReactionSet {
     private:
        std::vector<REACLIBReaction> m_reactions;
     public:
        REACLIBReactionSet() = default;

        explicit REACLIBReactionSet(std::vector<REACLIBReaction> r) : m_reactions(std::move(r)) {}

        void add_reaction(const REACLIBReaction& reaction) {
            if (not contains(reaction.id())) {
                m_reactions.push_back(reaction);
            }
        }

        void remove_reaction(const REACLIBReaction& reaction) {
            if (contains(reaction.id())) {
                m_reactions.erase(std::remove_if(m_reactions.begin(), m_reactions.end(),
                    [&reaction](const REACLIBReaction& r) { return r.id() == reaction.id(); }), m_reactions.end());
            }
        }

        [[nodiscard]] bool contains(const std::string& id) const {
            for (const auto& reaction : m_reactions) {
                if (reaction.id() == id) {
                    return true;
                }
            }
            return false;
        }

        [[nodiscard]] bool contains(const REACLIBReaction& reaction) const {
            for (const auto& r : m_reactions) {
                if (r == reaction) {
                    return true;
                }
            }
            return false;
        }

        [[nodiscard]] size_t size() const {
            return m_reactions.size();
        }

        friend std::ostream& operator<<(std::ostream& os, const REACLIBReactionSet& set) {
            os << "REACLIBReactionSet(";
            int reactionNo = 0;
            for (const auto& reaction : set.m_reactions) {
                reactionNo++;
                os << reaction.id();
                if (reactionNo < set.m_reactions.size()) {
                    os << ", ";
                }
            }
            os << ")";
            return os;
        }

        void sort(double T9=1.0) {
            // Sort based on the evaluation of each reaction's rate function
            std::sort(m_reactions.begin(), m_reactions.end(),
                [T9](const REACLIBReaction& a, const REACLIBReaction& b) {
                    return a.calculate_rate(T9) > b.calculate_rate(T9);
                });
        }

        auto begin() {
            return m_reactions.begin();
        }

        auto end() {
            return m_reactions.end();
        }

        bool containsSpecies(const fourdst::atomic::Species &species) const {
            for (const auto& reaction : m_reactions) {
                if (std::ranges::find(reaction.reactants(), species) != reaction.reactants().end() ||
                    std::ranges::find(reaction.products(), species) != reaction.products().end()) {
                    return true;
                }
            }
            return false;
        }

        [[nodiscard]] const REACLIBReaction& operator[](size_t index) const {
            if (index >= m_reactions.size()) {
                throw std::out_of_range("Index out of range in REACLIBReactionSet.");
            }
            return m_reactions[index];
        }

        [[nodiscard]] std::vector<REACLIBReaction> get_reactions() const {
            return m_reactions;
        }

        [[nodiscard]] auto begin() const {
            return m_reactions.cbegin();
        }

        [[nodiscard]] auto end() const {
            return m_reactions.cend();
        }

    };
    static std::unordered_map<std::string, REACLIBReaction> s_all_reaclib_reactions;
    static bool s_initialized = false;

    inline bool operator==(const REACLIBReaction& lhs, const REACLIBReaction& rhs) {
        return lhs.m_id == rhs.m_id;
    }
    inline bool operator!=(const REACLIBReaction& lhs, const REACLIBReaction& rhs) {
        return !(lhs == rhs);
    }

    inline bool operator==(const REACLIBReactionSet& lhs, const REACLIBReactionSet& rhs) {
        if (lhs.size() != rhs.size()) return false;
        for (const auto& reaction : lhs) {
            if (!rhs.contains(reaction)) return false;
        }
        return true;
    }

    inline bool operator!=(const REACLIBReactionSet& lhs, const REACLIBReactionSet& rhs) {
        return !(lhs == rhs);
    }
}

namespace std {
    template<>
    struct hash<gridfire::reaclib::REACLIBReaction> {
        size_t operator()(const gridfire::reaclib::REACLIBReaction& r) const noexcept {
            return std::hash<std::string>()(r.id());
        }
    };
} // namespace std