module;

#include <cmath>
#include <concepts>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <vector>

#include <mfem.hpp>

export module mean_field:seed.lane_emden;

export import :dimensions.quantities;
export import :eos.polytrope;
export import :model.typed_stellar;

export namespace mean_field::seed {
    struct DimensionlessLaneEmdenSolution final {
        mfem::Vector coordinate;
        mfem::Vector theta;
        mfem::Vector thetaDerivative;
        std::optional<double> firstZeroCoordinate;
    };

    /*
     * Integrate the dimensionless Lane-Emden equation from the regular center
     * to either the first zero of theta or coordinateLimit, whichever occurs
     * first. This numerical kernel also supports the n = 0 and n = 5 analytic
     * benchmark cases even though they do not both define admissible seeds for
     * the current Polytrope EOS and finite stellar domain.
     */
    [[nodiscard]] DimensionlessLaneEmdenSolution integrateLaneEmden(
        double polytropicIndex,
        double coordinateLimit,
        double integrationStep = 1.0e-3
    );

    struct RadialProfile final {
        mfem::Vector radius;
        mfem::Vector density;
        mfem::Vector specificEnthalpy;

        dimensions::LengthValue stellarRadius;
        dimensions::DensityValue centralDensity;
        dimensions::SpecificEnthalpyValue centralSpecificEnthalpy;
    };

    class LaneEmden final {
    public:
        struct Parameters final {
            std::optional<dimensions::DensityValue> centralDensity{std::nullopt};
            int radialSampleCount{512};
        };

        LaneEmden()
            : m_centralDensity(std::nullopt),
              m_radialSampleCount(512) {
        }

        explicit LaneEmden(const Parameters parameters)
            : m_centralDensity(parameters.centralDensity),
              m_radialSampleCount(parameters.radialSampleCount) {
            if (m_centralDensity.has_value() &&
                (!std::isfinite(m_centralDensity->value()) || m_centralDensity->value() <= 0.0)) {
                throw std::invalid_argument("A Lane-Emden seed central density must be finite and positive.");
            }
            if (m_radialSampleCount < 2) {
                throw std::invalid_argument("A Lane-Emden seed requires at least two radial samples.");
            }
        }

        [[nodiscard]] const std::optional<dimensions::DensityValue> &centralDensity() const noexcept {
            return m_centralDensity;
        }

        [[nodiscard]] int radialSampleCount() const noexcept {
            return m_radialSampleCount;
        }

    private:
        std::optional<dimensions::DensityValue> m_centralDensity;
        int m_radialSampleCount;
    };

    [[nodiscard]] RadialProfile generateLaneEmdenProfile(
        const eos::Polytrope &equationOfState,
        dimensions::DensityValue centralDensity,
        int radialSampleCount
    );

    template <model::StellarModelType Model>
        requires std::remove_cvref_t<Model>::template
    containsSpecification<eos::Polytrope> [[nodiscard]] RadialProfile generateRadialProfile(
        const Model &stellarModel,
        const LaneEmden &strategy
    ) {
        std::optional<dimensions::DensityValue> centralDensity = strategy.centralDensity();

        if (!centralDensity.has_value()) {
            if constexpr (std::remove_cvref_t<Model>::template containsSpecification<models::FixedCentralDensity>) {
                centralDensity = stellarModel.template specification<models::FixedCentralDensity>().targetDensity();
            } else {
                throw std::invalid_argument(
                    "Lane-Emden seed generation requires either FixedCentralDensity or an explicit seed-only central "
                    "density."
                );
            }
        }

        return generateLaneEmdenProfile(
            stellarModel.template specification<eos::Polytrope>(), *centralDensity, strategy.radialSampleCount()
        );
    }

    template <typename Strategy, typename Model>
    concept RadialSeedStrategyFor = requires(const Model &stellarModel, const Strategy &strategy) {
        { generateRadialProfile(stellarModel, strategy) } -> std::same_as<RadialProfile>;
    };
} // namespace mean_field::seed