module;

#include <cmath>
#include <concepts>
#include <cstddef>
#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;
    };

    /*
     * The public radial-projection extension boundary deliberately speaks in
     * physical state names. A specification author supplies one small rule
     * and opts in with
     *
     *   using RadialProjection = seed::projection::Use<MyProjectionPhysics>;
     *
     * There is no registry ordinal and no model-combination specialization.
     */
    struct RadialProjectionScales final {
        dimensions::MassValue targetMass;
        dimensions::LengthValue stellarRadius;
        double bernoulliConstant;
        double sphericalMomentOfInertia;
        const mfem::Array<int> *surfaceCarrierRows;
    };

    struct RadialProjectionState final {
        mfem::Vector density;
        mfem::Vector surfaceShape;
        mfem::Vector gravityGradient;
        mfem::Vector gravityPotential;
        mfem::Vector specificEnthalpy;
    };

    namespace projection {
        template <typename Physics> struct Use final {
            using PhysicsType = Physics;
        };

        /* An explicit opt-in for a specification that leaves a radial seed unchanged. */
        struct NoStateChange {
            static constexpr bool registered                         = true;
            static constexpr bool providesRadialMass                 = false;

            template <typename Model> static constexpr bool supports = true;

            template <
                typename Specification,
                typename Model>
            static void validate(
                const Specification &,
                const Model &,
                const RadialProfile &,
                const StellarEquilibriumProjectionOptions &
            ) noexcept {
            }

            template <
                typename Specification,
                typename Model>
            static void initialize(
                const Specification &,
                const Model &,
                const RadialProjectionScales &,
                RadialProjectionState &,
                mfem::Vector
            ) noexcept {
            }
        };
    } // namespace projection

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

        [[nodiscard]] ProjectedRadialFields projectRadialFields(
            fem::FEM &finiteElementModel,
            const RadialProfile &profile,
            dimensions::MassValue targetMass,
            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;
        }

        struct UnavailableRadialProjectionPhysics final {
            static constexpr bool registered                         = false;
            static constexpr bool providesRadialMass                 = false;

            template <typename Model> static constexpr bool supports = false;
        };

        struct FixedTotalMassRadialProjectionPhysics final {
            static constexpr bool registered                         = true;
            static constexpr bool providesRadialMass                 = true;

            template <typename Model> static constexpr bool supports = true;

            [[nodiscard]] static dimensions::MassValue targetMass(const models::FixedTotalMass &specification) {
                return specification.targetMass();
            }

            template <typename Model>
            static void validate(
                const models::FixedTotalMass &,
                const Model &,
                const RadialProfile &,
                const StellarEquilibriumProjectionOptions &
            ) noexcept {
            }

            template <typename Model>
            static void initialize(
                const models::FixedTotalMass &,
                const Model &,
                const RadialProjectionScales &scales,
                RadialProjectionState &,
                mfem::Vector coordinate
            ) {
                if (coordinate.Size() != 1) {
                    throw std::invalid_argument(
                        "A fixed-total-mass radial projection requires exactly one generated multiplier."
                    );
                }
                coordinate(0) = scales.bernoulliConstant;
            }
        };

        struct IsobaricRadialProjectionPhysics final {
            static constexpr bool registered         = true;
            static constexpr bool providesRadialMass = false;

            template <typename Model>
            static constexpr bool supports = requires(const Model &model, const surface::Isobaric &condition) {
                {
                    eos::evaluate<dimensions::quantity::SpecificEnthalpy>(
                        model.equationOfState(), condition.targetPressure()
                    )
                } -> std::same_as<dimensions::SpecificEnthalpyValue>;
            };

            template <typename Model>
            static void validate(
                const surface::Isobaric &condition,
                const Model &,
                const RadialProfile &,
                const StellarEquilibriumProjectionOptions &
            ) {
                if (condition.targetPressure().value() != 0.0) {
                    throw std::invalid_argument("A Lane-Emden radial seed requires a zero-pressure isobaric surface.");
                }
            }

            template <typename Model>
            static void initialize(
                const surface::Isobaric &condition,
                const Model &model,
                const RadialProjectionScales &scales,
                RadialProjectionState &state,
                mfem::Vector coordinate
            ) {
                if (coordinate.Size() != 0) {
                    throw std::logic_error("An isobaric surface must not generate a radial-seed coordinate.");
                }
                if (scales.surfaceCarrierRows == nullptr) {
                    throw std::logic_error("An isobaric radial projection requires compiled surface-carrier rows.");
                }

                const dimensions::SpecificEnthalpyValue requiredSurfaceEnthalpy =
                    eos::evaluate<dimensions::quantity::SpecificEnthalpy>(
                        model.equationOfState(), condition.targetPressure()
                    );
                for (const int surfaceRow : *scales.surfaceCarrierRows) {
                    state.specificEnthalpy(surfaceRow) = requiredSurfaceEnthalpy.value();
                }
            }
        };

        struct FixedCentralDensityRadialProjectionPhysics final {
            static constexpr bool registered                         = true;
            static constexpr bool providesRadialMass                 = false;

            template <typename Model> static constexpr bool supports = true;

            template <typename Model>
            static void validate(
                const models::FixedCentralDensity &,
                const Model &,
                const RadialProfile &,
                const StellarEquilibriumProjectionOptions &
            ) noexcept {
            }

            template <typename Model>
            static void initialize(
                const models::FixedCentralDensity &,
                const Model &,
                const RadialProjectionScales &,
                RadialProjectionState &,
                mfem::Vector coordinate
            ) {
                if (coordinate.Size() != 1) {
                    throw std::invalid_argument(
                        "A fixed-central-density radial projection requires exactly one generated phase coordinate."
                    );
                }
                coordinate(0) = 0.0;
            }
        };

        struct FixedAngularMomentumRadialProjectionPhysics final {
            static constexpr bool registered                         = true;
            static constexpr bool providesRadialMass                 = false;

            template <typename Model> static constexpr bool supports = true;

            template <typename Model>
            static void validate(
                const models::FixedAngularMomentum &,
                const Model &,
                const RadialProfile &,
                const StellarEquilibriumProjectionOptions &
            ) noexcept {
            }

            template <typename Model>
            static void initialize(
                const models::FixedAngularMomentum &constraint,
                const Model &,
                const RadialProjectionScales &scales,
                RadialProjectionState &,
                mfem::Vector coordinate
            ) {
                if (coordinate.Size() != 1) {
                    throw std::logic_error(
                        "A fixed-angular-momentum radial projection requires one angular-velocity coordinate."
                    );
                }

                double centerSquared   = 0.0;
                double centerAlongAxis = 0.0;
                for (std::size_t component = 0; component < constraint.center().size(); ++component) {
                    centerSquared += constraint.center()[component] * constraint.center()[component];
                    centerAlongAxis += constraint.center()[component] * constraint.axis()[component];
                }
                const double parallelAxisCorrection =
                    scales.targetMass.value() * (centerSquared - centerAlongAxis * centerAlongAxis);
                const double momentOfInertia = scales.sphericalMomentOfInertia + parallelAxisCorrection;
                if (!std::isfinite(momentOfInertia) || momentOfInertia <= 0.0) {
                    throw std::runtime_error("The radial seed has no finite, positive axial moment of inertia.");
                }
                coordinate(0) = constraint.targetAngularMomentum().value() / momentOfInertia;
            }
        };

        template <typename Specification> struct BuiltinRadialProjectionPhysics {
            using Type = UnavailableRadialProjectionPhysics;
        };

        template <> struct BuiltinRadialProjectionPhysics<eos::Polytrope> {
            using Type = projection::NoStateChange;
        };

        template <> struct BuiltinRadialProjectionPhysics<surface::Isobaric> {
            using Type = IsobaricRadialProjectionPhysics;
        };

        template <> struct BuiltinRadialProjectionPhysics<models::FixedTotalMass> {
            using Type = FixedTotalMassRadialProjectionPhysics;
        };

        template <> struct BuiltinRadialProjectionPhysics<models::FixedCentralDensity> {
            using Type = FixedCentralDensityRadialProjectionPhysics;
        };

        template <> struct BuiltinRadialProjectionPhysics<models::FixedAngularMomentum> {
            using Type = FixedAngularMomentumRadialProjectionPhysics;
        };

        template <typename Candidate> struct UnwrapRadialProjectionPhysics {
            using Type                  = UnavailableRadialProjectionPhysics;
            static constexpr bool valid = false;
        };

        template <typename Physics> struct UnwrapRadialProjectionPhysics<projection::Use<Physics>> {
            using Type                  = Physics;
            static constexpr bool valid = true;
        };

        template <typename Specification, typename = void> struct SelectRadialProjectionPhysics {
            using Type = typename BuiltinRadialProjectionPhysics<Specification>::Type;
        };

        template <typename Specification>
        struct SelectRadialProjectionPhysics<Specification, std::void_t<typename Specification::RadialProjection>> {
        private:
            using Wrapped = UnwrapRadialProjectionPhysics<typename Specification::RadialProjection>;

        public:
            using Type = std::conditional_t<Wrapped::valid, typename Wrapped::Type, UnavailableRadialProjectionPhysics>;
        };

        template <typename Physics> [[nodiscard]] consteval bool radialProjectionPhysicsRegistered() {
            if constexpr (requires {
                              { Physics::registered } -> std::convertible_to<bool>;
                          }) {
                return static_cast<bool>(Physics::registered);
            } else {
                return false;
            }
        }

        template <typename Physics> [[nodiscard]] consteval bool radialProjectionPhysicsProvidesMass() {
            if constexpr (requires {
                              { Physics::providesRadialMass } -> std::convertible_to<bool>;
                          }) {
                return static_cast<bool>(Physics::providesRadialMass);
            } else {
                return false;
            }
        }

        template <
            typename Specification,
            typename Model>
        [[nodiscard]] consteval bool radialProjectionPhysicsIsComplete() {
            using Physics = typename SelectRadialProjectionPhysics<Specification>::Type;
            if constexpr (!radialProjectionPhysicsRegistered<Physics>()) {
                return false;
            } else if constexpr (!requires {
                                     { Physics::template supports<Model> } -> std::convertible_to<bool>;
                                 }) {
                return false;
            } else if constexpr (!static_cast<bool>(Physics::template supports<Model>)) {
                return false;
            } else if constexpr (!requires(
                                     const Specification &specification, const Model &model,
                                     const RadialProfile &profile, const StellarEquilibriumProjectionOptions &options,
                                     const RadialProjectionScales &scales, RadialProjectionState &state,
                                     mfem::Vector coordinate
                                 ) {
                                     Physics::validate(specification, model, profile, options);
                                     Physics::initialize(specification, model, scales, state, coordinate);
                                 }) {
                return false;
            } else if constexpr (radialProjectionPhysicsProvidesMass<Physics>()) {
                return requires(const Specification &specification) {
                    { Physics::targetMass(specification) } -> std::same_as<dimensions::MassValue>;
                };
            } else {
                return true;
            }
        }

        template <models::ModelSpecification Specification, models::GeneratedStateKind Kind>
        struct RadialProjectionCoordinateTerm;

        template <models::ModelSpecification Specification>
        struct RadialProjectionCoordinateTerm<Specification, models::GeneratedStateKind::multiplier> final {
            using value = utils::blocks::generated_value_block<models::MultiplierFor<Specification>>;
        };

        template <models::ModelSpecification Specification>
        struct RadialProjectionCoordinateTerm<Specification, models::GeneratedStateKind::physical_coordinate> final {
            using value = utils::blocks::generated_value_block<models::PhysicalCoordinateFor<Specification>>;
        };

        template <models::ModelSpecification Specification>
        struct RadialProjectionCoordinateTerm<Specification, models::GeneratedStateKind::solver_border> final {
            using value = utils::blocks::generated_value_block<models::BorderFor<Specification>>;
        };

        template <typename Model, typename SpecificationSet> struct CompileRadialProjection;

        template <model::StellarModelType Model, models::ModelSpecification... Specifications>
        struct CompileRadialProjection<Model, models::detail::SpecificationSetStorage<Specifications...>> {
            using ModelType = std::remove_cvref_t<Model>;

            static constexpr std::size_t radialMassProviderCount =
                (std::size_t{0} + ... +
                 (radialProjectionPhysicsProvidesMass<typename SelectRadialProjectionPhysics<Specifications>::Type>()
                      ? std::size_t{1}
                      : std::size_t{0}));
            static constexpr bool complete =
                radialMassProviderCount == 1 && (radialProjectionPhysicsIsComplete<Specifications, ModelType>() && ...);

            [[nodiscard]] static dimensions::MassValue targetMass(const ModelType &model)
                requires complete
            {
                dimensions::MassValue result{0.0};
                (
                    [&] {
                        using Physics = typename SelectRadialProjectionPhysics<Specifications>::Type;
                        if constexpr (radialProjectionPhysicsProvidesMass<Physics>()) {
                            result = Physics::targetMass(model.template specification<Specifications>());
                        }
                    }(),
                    ...);
                return result;
            }

            static void validate(
                const ModelType &model,
                const RadialProfile &profile,
                const StellarEquilibriumProjectionOptions &options
            )
                requires complete
            {
                (
                    [&] {
                        using Physics = typename SelectRadialProjectionPhysics<Specifications>::Type;
                        Physics::validate(model.template specification<Specifications>(), model, profile, options);
                    }(),
                    ...);
            }

            template <typename StateView>
            static void initialize(
                const ModelType &model,
                const RadialProjectionScales &scales,
                RadialProjectionState &state,
                const StateView &stateView
            )
                requires complete
            {
                (
                    [&] {
                        using Contribution = models::SpecificationContribution<Specifications>;
                        using Physics      = typename SelectRadialProjectionPhysics<Specifications>::Type;
                        if constexpr (Contribution::generatedValueArity == 0) {
                            Physics::initialize(
                                model.template specification<Specifications>(), model, scales, state, mfem::Vector{}
                            );
                        } else {
                            static_assert(
                                Contribution::generatedValueArity == 1,
                                "Radial projection currently requires each specification contribution to generate at "
                                "most one scalar coordinate."
                            );
                            using Term =
                                RadialProjectionCoordinateTerm<Specifications, Contribution::generatedStateKind>;
                            Physics::initialize(
                                model.template specification<Specifications>(), model, scales, state,
                                stateView.block(Term{})
                            );
                        }
                    }(),
                    ...);
            }
        };

        template <typename Candidate, bool = model::StellarModelType<Candidate>>
        struct RadialProjectionCompilationAudit {
            static constexpr bool complete = false;
        };

        template <typename Candidate>
        struct RadialProjectionCompilationAudit<Candidate, true>
            : CompileRadialProjection<Candidate, typename Candidate::SpecificationTypes> { };
    } // namespace detail

    template <typename Candidate>
    inline constexpr bool radialProjectionIsCompilable =
        detail::RadialProjectionCompilationAudit<std::remove_cvref_t<Candidate>>::complete;

    template <typename Candidate>
    concept RadialProfileProjectableModel =
        model::StellarModelType<Candidate> && radialProjectionIsCompilable<std::remove_cvref_t<Candidate>>;

    template <
        equilibrium::StellarEquilibriumModel Model,
        equilibrium::StellarDiscretizationType Discretization>
        requires RadialProfileProjectableModel<Model>
    [[nodiscard]] ProjectedEquilibriumState<Model> projectRadialProfile(
        equilibrium::StellarEquilibriumProblem<
            Model,
            Discretization> &problem,
        const RadialProfile &profile,
        const StellarEquilibriumProjectionOptions &options = {}
    ) {
        using Projection = detail::CompileRadialProjection<
            std::remove_cvref_t<Model>, typename std::remove_cvref_t<Model>::SpecificationTypes>;
        const auto &stellarModel = problem.GetStellarModel();
        Projection::validate(stellarModel, profile, options);
        const dimensions::MassValue targetMass     = Projection::targetMass(stellarModel);
        const detail::ProjectedRadialFields fields = detail::projectRadialFields(
            equilibrium::detail::StellarEquilibriumProblemFactory::MutableFiniteElementModelForProjection(problem),
            profile, targetMass, 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."
        );

        /*
         * Each specification now initializes only its inferred contribution.
         * In particular, the surface rule imposes the exact carrier trace and
         * generated constraints obtain their coordinate by type, not by a
         * hard-coded whole-model layout.
         */
        RadialProjectionState projectedState{
            .density          = stateView.block(utils::blocks::density_field.mass_term),
            .surfaceShape     = stateView.block(utils::blocks::surface_deformation_field.parameters_term),
            .gravityGradient  = stateView.block(utils::blocks::gravity_field.gradient_term),
            .gravityPotential = stateView.block(utils::blocks::gravity_field.poisson_term),
            .specificEnthalpy = stateView.block(utils::blocks::enthalpy_field.specific_term)
        };
        const RadialProjectionScales scales{
            .targetMass               = targetMass,
            .stellarRadius            = profile.stellarRadius,
            .bernoulliConstant        = fields.bernoulliConstant,
            .sphericalMomentOfInertia = fields.sphericalMomentOfInertia,
            .surfaceCarrierRows       = &problem.GetPressureSurfaceRows().reduced_dofs()
        };
        Projection::initialize(stellarModel, scales, projectedState, stateView);

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

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