module;

#include <concepts>
#include <stdexcept>
#include <type_traits>
#include <utility>

#include <mfem.hpp>

export module mean_field:seed.stellar_equilibrium_projection;

export import :operators.stellar_equilibrium_problem;
export import :physics.gravity;
export import :seed.lane_emden;

export namespace mean_field::seed {
    struct StellarEquilibriumProjectionOptions final {
        physics::GravitySolveOptions gravity{};
        double surfaceRadiusRelativeTolerance{5.0e-4};
    };

    template <equilibrium::StellarEquilibriumModel Model> struct ProjectedEquilibriumState final {
        using ModelType = std::remove_cvref_t<Model>;

        mfem::Vector values;
    };

    namespace detail {
        struct ProjectedRadialFields final {
            mfem::Vector density;
            mfem::Vector gravityGradient;
            mfem::Vector gravityPotential;
            mfem::Vector specificEnthalpy;
            double bernoulliConstant;
        };

        [[nodiscard]] ProjectedRadialFields projectRadialFields(
            const equilibrium::StellarDiscretization &discretization,
            const RadialProfile &profile,
            dimensions::MassValue targetMass,
            dimensions::PressureValue targetSurfacePressure,
            const StellarEquilibriumProjectionOptions &options
        );

        inline void assignProjectedBlock(
            mfem::Vector destination,
            const mfem::Vector &source,
            const char *name
        ) {
            if (destination.Size() != source.Size()) {
                throw std::invalid_argument(name);
            }
            destination = source;
        }
    } // namespace detail

    template <equilibrium::StellarEquilibriumModel Model>
    [[nodiscard]] ProjectedEquilibriumState<Model> projectRadialProfile(
        const equilibrium::StellarEquilibriumProblem<Model> &problem,
        const RadialProfile &profile,
        const StellarEquilibriumProjectionOptions &options = {}
    ) {
        const detail::ProjectedRadialFields fields = detail::projectRadialFields(
            problem.GetDiscretization(), profile,
            problem.GetStellarModel().template specification<models::FixedTotalMass>().targetMass(),
            problem.GetStellarModel().template specification<surface::Isobaric>().targetPressure(), options
        );

        mfem::Vector values(problem.StateSize());
        values               = 0.0;
        const auto stateView = problem.GetManifest().stateView(values);

        detail::assignProjectedBlock(
            stateView.block(utils::blocks::density_field.mass_term), fields.density,
            "The projected density does not match the compiled equilibrium-state block."
        );
        stateView.block(utils::blocks::surface_deformation_field.parameters_term) = 0.0;
        detail::assignProjectedBlock(
            stateView.block(utils::blocks::gravity_field.gradient_term), fields.gravityGradient,
            "The projected gravity gradient does not match the compiled equilibrium-state block."
        );
        detail::assignProjectedBlock(
            stateView.block(utils::blocks::gravity_field.poisson_term), fields.gravityPotential,
            "The projected gravity potential does not match the compiled equilibrium-state block."
        );
        detail::assignProjectedBlock(
            stateView.block(utils::blocks::enthalpy_field.specific_term), fields.specificEnthalpy,
            "The projected specific enthalpy does not match the compiled equilibrium-state block."
        );

        /*
         * Projection of a continuous spherical profile onto a faceted
         * reference mesh generally leaves a small trace error on the physical
         * surface. The pressure condition replaces these carrier rows in the
         * compiled equilibrium problem, so impose its required carrier value
         * exactly after bulk projection instead of treating that geometric
         * mismatch as part of the initial residual.
         */
        mfem::Vector enthalpy = stateView.block(utils::blocks::enthalpy_field.specific_term);
        const dimensions::SpecificEnthalpyValue requiredSurfaceEnthalpy =
            eos::evaluate<dimensions::quantity::SpecificEnthalpy>(
                problem.GetStellarModel().template specification<eos::Polytrope>(),
                problem.GetStellarModel().template specification<surface::Isobaric>().targetPressure()
            );
        for (const int surfaceRow : problem.GetPressureSurfaceRows().reduced_dofs()) {
            enthalpy(surfaceRow) = requiredSurfaceEnthalpy.value();
        }

        mfem::Vector fixedMassCoordinate =
            stateView.block(utils::blocks::fixed_total_mass_constraint.mass_normalization_term);
        if (fixedMassCoordinate.Size() != 1) {
            throw std::invalid_argument("FixedTotalMass must generate exactly one equilibrium-state coordinate.");
        }
        fixedMassCoordinate(0) = fields.bernoulliConstant;

        if constexpr (std::remove_cvref_t<Model>::template containsSpecification<models::FixedCentralDensity>) {
            stateView.block(utils::blocks::fixed_central_density_phase.central_value_term) = 0.0;
        }

        return {.values = std::move(values)};
    }

    template <
        equilibrium::StellarEquilibriumModel Model,
        typename Strategy>
        requires RadialSeedStrategyFor<
            Strategy,
            typename equilibrium::StellarEquilibriumProblem<Model>::ModelType>
    [[nodiscard]] ProjectedEquilibriumState<Model> makeProjectedEquilibriumState(
        const equilibrium::StellarEquilibriumProblem<Model> &problem,
        const Strategy &strategy,
        const StellarEquilibriumProjectionOptions &options = {}
    ) {
        return projectRadialProfile(problem, generateRadialProfile(problem.GetStellarModel(), strategy), options);
    }
} // namespace mean_field::seed