module;

#include <concepts>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <type_traits>
#include <utility>

#include <mfem.hpp>

export module mean_field:operators.stellar_equilibrium_problem;

export import :deformation.domain_deformation;
export import :equilibrium.stellar_discretization;
export import :material.thermodynamic_equations;
export import :model.typed_stellar;
export import :normalization.operators;
export import :operators.prepared_variadic_stellar_equilibrium;
export import :surface.compiler;

export namespace mean_field::equilibrium {
    namespace detail {
        template <
            model::StellarModelType Model,
            bool SymbolicallyCompilable = operators::StellarEquilibriumSystemCompilable<Model>>
        struct StellarSurfaceCompilationAudit {
            static constexpr bool complete = false;
        };

        template <model::StellarModelType Model>
        struct StellarSurfaceCompilationAudit<Model, true> {
        private:
            using ModelType = std::remove_cvref_t<Model>;
            using EquationOfState = typename ModelType::EquationOfStateType;
            using Form = operators::CompiledStellarEquilibriumForm<ModelType>;
            using AvailableEquations = material::StellarEquilibriumThermodynamicEquations;

            static constexpr bool thermodynamicsCompilable =
                material::ThermodynamicEquationsCompilable<EquationOfState, Form, AvailableEquations>;

        public:
            static constexpr bool complete = [] {
                if constexpr (!thermodynamicsCompilable) {
                    return false;
                } else {
                    using ThermodynamicEquations =
                        material::CompiledThermodynamicEquationsT<EquationOfState, Form, AvailableEquations>;
                    using Formulation = typename ThermodynamicEquations::PressureSurfaceFormulation;
                    using CompiledSurface =
                        surface::CompiledPressureSurfaceConstraintT<Formulation, EquationOfState>;
                    return requires(const ModelType &model) {
                        {
                            surface::compilePressureSurfaceConstraint<Formulation>(
                                model.surfaceCondition(),
                                model.equationOfState()
                            )
                        } -> std::same_as<CompiledSurface>;
                    };
                }
            }();
        };
    } // namespace detail

    template <model::StellarModelType Model>
    inline constexpr bool hasStellarEquilibriumSurfaceCompilation =
        detail::StellarSurfaceCompilationAudit<std::remove_cvref_t<Model>>::complete;

    template <typename Candidate>
    concept StellarEquilibriumModel = model::StellarModelType<Candidate> && requires {
        typename std::remove_cvref_t<Candidate>::EquationOfStateType;
        requires(
            std::remove_cvref_t<Candidate>::template specificationRoleCount<
                models::SpecificationRole::boundary_condition> == 1
        );
        requires std::remove_cvref_t<Candidate>::template containsSpecification<models::FixedTotalMass>;
        requires operators::StellarEquilibriumSystemCompilable<std::remove_cvref_t<Candidate>>;
        requires hasStellarEquilibriumSurfaceCompilation<std::remove_cvref_t<Candidate>>;
        requires operators::CompilableRootManifestFor<
            std::remove_cvref_t<Candidate>,
            operators::CompiledStellarEquilibriumForm<std::remove_cvref_t<Candidate>>>;
        requires operators::hasStellarEquilibriumCoreRuntime<std::remove_cvref_t<Candidate>>;
        requires operators::hasCompleteStellarEquilibriumRuntime<std::remove_cvref_t<Candidate>>;
        requires operators::stellarEquilibriumRotationProviderCount<std::remove_cvref_t<Candidate>> <= 1;
    };

    namespace detail {
        template <typename Model, typename Discretization, typename = void>
        struct StellarEquilibriumModelDiscretizationStructureAudit : std::false_type { };

        template <typename Model, typename Discretization>
            requires StellarEquilibriumModel<std::remove_cvref_t<Model>> &&
                     StellarDiscretizationType<std::remove_cvref_t<Discretization>>
        struct StellarEquilibriumModelDiscretizationStructureAudit<
            Model,
            Discretization,
            std::void_t<
                typename std::remove_cvref_t<Discretization>::NormalizationPrescriptionType,
                typename std::remove_cvref_t<Model>::SpecificationTypes,
                operators::CompiledStellarEquilibriumForm<std::remove_cvref_t<Model>>,
                operators::StellarEquilibriumPhysicalCoreType<std::remove_cvref_t<Model>>>>
            : std::bool_constant<normalization::StellarNormalizationRuntimeAvailableFor<
                  typename std::remove_cvref_t<Discretization>::NormalizationPrescriptionType,
                  operators::CompiledStellarEquilibriumForm<std::remove_cvref_t<Model>>,
                  operators::StellarEquilibriumPhysicalCoreType<std::remove_cvref_t<Model>>,
                  typename std::remove_cvref_t<Model>::SpecificationTypes>> { };

        struct StellarEquilibriumProblemFactory;
    } // namespace detail

    template <
        StellarEquilibriumModel Model,
        StellarDiscretizationType Discretization = StellarDiscretization>
        requires detail::StellarEquilibriumModelDiscretizationStructureAudit<
            std::remove_cvref_t<Model>,
            std::remove_cvref_t<Discretization>>::value
    class StellarEquilibriumProblem final {
    public:
        using ModelType = std::remove_cvref_t<Model>;
        using DiscretizationType = std::remove_cvref_t<Discretization>;
        using NormalizationPrescriptionType = typename DiscretizationType::NormalizationPrescriptionType;

        static constexpr bool hasFixedCentralDensity =
            ModelType::template containsSpecification<models::FixedCentralDensity>;
        static constexpr bool hasFixedAngularMomentum =
            ModelType::template containsSpecification<models::FixedAngularMomentum>;
        static constexpr std::size_t generatedRotationProviderCount =
            operators::stellarEquilibriumRotationProviderCount<ModelType>;
        static constexpr bool symbolicallySquare = ModelType::symbolicallySquare;

        using PreparedOperatorType = operators::PreparedVariadicStellarEquilibriumOperator<ModelType>;
        using PhysicalCoreType = typename PreparedOperatorType::PhysicalCoreType;
        using FormType = operators::CompiledStellarEquilibriumForm<ModelType>;
        using JacobianFormType = operators::CompiledStellarEquilibriumJacobianForm<ModelType>;
        using ManifestType = operators::EquilibriumSystemManifest<ModelType, FormType, JacobianFormType>;
        using EquationOfStateType             = model::EquationOfStateType<ModelType>;
        using SurfaceConditionType            = model::SurfaceConditionType<ModelType>;
        using AvailableThermodynamicEquations = material::StellarEquilibriumThermodynamicEquations;
        using ThermodynamicEquationsType =
            material::CompiledThermodynamicEquationsT<EquationOfStateType, FormType, AvailableThermodynamicEquations>;
        using CompiledSurfaceConstraintType = surface::CompiledPressureSurfaceConstraintT<
            typename ThermodynamicEquationsType::PressureSurfaceFormulation,
            EquationOfStateType>;

    private:
        friend struct detail::StellarEquilibriumProblemFactory;

        StellarEquilibriumProblem(
            ModelType stellarModel,
            DiscretizationType discretization
        )
            : m_stellarModel(std::make_shared<ModelType>(std::move(stellarModel))),
              m_discretization(std::move(discretization)),
              m_compiledSurfaceConstraint(CompileSurfaceConstraint(*m_stellarModel)),
              m_preparedOperator(
                  m_discretization.finiteElementModel(),
                  m_discretization.domainMapper(),
                  m_stellarModel,
                  operators::PressureSurfaceConstraintView{m_compiledSurfaceConstraint},
                  CompileDefaultDomainDeformation(m_discretization.finiteElementModel())
              ) {
            VerifyProblem();
        }

        StellarEquilibriumProblem(const StellarEquilibriumProblem &)            = delete;
        StellarEquilibriumProblem &operator=(const StellarEquilibriumProblem &) = delete;
        StellarEquilibriumProblem(StellarEquilibriumProblem &&)                 = delete;
        StellarEquilibriumProblem &operator=(StellarEquilibriumProblem &&)      = delete;

    public:
        [[nodiscard]] const ModelType &GetStellarModel() const noexcept {
            return *m_stellarModel;
        }

        [[nodiscard]] const DiscretizationType &GetDiscretization() const noexcept {
            return m_discretization;
        }

        [[nodiscard]] const NormalizationPrescriptionType &GetNormalizationPrescription() const noexcept {
            return m_discretization.normalizationPrescription();
        }

        [[nodiscard]] const CompiledSurfaceConstraintType &GetCompiledSurfaceConstraint() const noexcept {
            return m_compiledSurfaceConstraint;
        }

        [[nodiscard]] PreparedOperatorType &GetPreparedOperator() noexcept {
            return m_preparedOperator;
        }

        [[nodiscard]] const PreparedOperatorType &GetPreparedOperator() const noexcept {
            return m_preparedOperator;
        }

        [[nodiscard]] const PhysicalCoreType &GetPhysicalOperator() const noexcept {
            return m_preparedOperator.GetPhysicalOperator();
        }

        [[nodiscard]] const auto &GetManifest() const noexcept {
            return m_preparedOperator.GetRootManifest();
        }

        [[nodiscard]] bool IsPrepared() const noexcept {
            return m_preparedOperator.IsPrepared();
        }

        [[nodiscard]] std::uint64_t GetPreparationGeneration() const noexcept {
            return m_preparationGeneration;
        }

        [[nodiscard]] const operators::StellarEquilibriumDependencies &GetLinearizationDependencies() const {
            return GetPhysicalOperator().GetDependencies();
        }

        [[nodiscard]] const operators::StellarEquilibriumDependencyStamp &GetGeometryDependency() const {
            return GetPhysicalOperator().GetGeneratedDisplacementDependency();
        }

        [[nodiscard]] const field::FieldBoundaryDofMap &GetPressureSurfaceRows() const noexcept {
            return GetPhysicalOperator().GetSurfaceConstraintOperator().GetSurfaceRows();
        }

        [[nodiscard]] int StateSize() const noexcept {
            return m_preparedOperator.Width();
        }

        [[nodiscard]] int EquationSize() const noexcept {
            return m_preparedOperator.Height();
        }

        [[nodiscard]] const mfem::Operator &GetLinearizationOperator() const noexcept {
            return m_preparedOperator;
        }

        [[nodiscard]] auto Prepare(
            const mfem::Vector &state,
            const operators::StellarEquilibriumDependencies &dependencies,
            const physics::RigidRotation &rotation
        ) requires(generatedRotationProviderCount == 0) {
            auto report = m_preparedOperator.Prepare(state, dependencies, rotation);
            ++m_preparationGeneration;
            return report;
        }

        [[nodiscard]] auto Prepare(
            const mfem::Vector &state,
            const operators::StellarEquilibriumDependencies &dependencies
        ) requires(generatedRotationProviderCount == 1) {
            auto report = m_preparedOperator.Prepare(state, dependencies);
            ++m_preparationGeneration;
            return report;
        }

        void BuildResidual(mfem::Vector &residual) const {
            m_preparedOperator.BuildResidual(residual);
        }

        void ApplyLinearization(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const {
            m_preparedOperator.Mult(direction, action);
        }

    private:
        [[nodiscard]] static CompiledSurfaceConstraintType CompileSurfaceConstraint(const ModelType &stellarModel) {
            return surface::compilePressureSurfaceConstraint<
                typename ThermodynamicEquationsType::PressureSurfaceFormulation>(
                stellarModel.surfaceCondition(),
                stellarModel.equationOfState()
            );
        }

        [[nodiscard]] static deformation::PreparedDomainDeformationRuntime
        CompileDefaultDomainDeformation(fem::FEM &finiteElementModel) {
            MFEM_VERIFY(
                finiteElementModel.mesh != nullptr,
                "Default stellar domain-deformation compilation requires a physical mesh."
            );
            mfem::Vector referenceCenter(finiteElementModel.mesh->SpaceDimension());
            referenceCenter = 0.0;

            return deformation::PreparedDomainDeformationRuntime{deformation::compileDomainDeformation(
                deformation::NodalRadialSurface{std::move(referenceCenter)},
                deformation::PowerLawRadialInteriorExtension{}, deformation::FixedInfinityRadialVacuumExtension{},
                finiteElementModel
            )};
        }

        void VerifyProblem() const {
            MFEM_VERIFY(symbolicallySquare, "A stellar equilibrium problem must be symbolically square.");
            MFEM_VERIFY(
                StateSize() == EquationSize(),
                "The discretized stellar equilibrium problem has unequal state and equation dimensions."
            );
            MFEM_VERIFY(m_discretization.isCurrent(), "The stellar equilibrium problem has a stale discretization.");
        }

        std::shared_ptr<const ModelType> m_stellarModel;
        DiscretizationType m_discretization;
        CompiledSurfaceConstraintType m_compiledSurfaceConstraint;
        PreparedOperatorType m_preparedOperator;
        std::uint64_t m_preparationGeneration{0};
    };

    template <typename Candidate> struct IsStellarEquilibriumProblem : std::false_type { };

    template <StellarEquilibriumModel Model, StellarDiscretizationType Discretization>
        requires detail::StellarEquilibriumModelDiscretizationStructureAudit<
            std::remove_cvref_t<Model>,
            std::remove_cvref_t<Discretization>>::value
    struct IsStellarEquilibriumProblem<StellarEquilibriumProblem<Model, Discretization>> : std::true_type { };

    template <typename Candidate>
    concept DiscretizedStellarEquilibriumProblem = IsStellarEquilibriumProblem<std::remove_cvref_t<Candidate>>::value;

    namespace detail {
        template <
            typename Model,
            typename Discretization,
            bool StructurallyCompatible =
                StellarEquilibriumModelDiscretizationStructureAudit<
                    std::remove_cvref_t<Model>,
                    std::remove_cvref_t<Discretization>>::value>
        struct StellarEquilibriumModelDiscretizationOperationAudit : std::false_type { };

        template <typename Model, typename Discretization>
        struct StellarEquilibriumModelDiscretizationOperationAudit<
            Model,
            Discretization,
            true> {
        private:
            using ModelType = std::remove_cvref_t<Model>;
            using DiscretizationType = std::remove_cvref_t<Discretization>;
            using Problem = StellarEquilibriumProblem<ModelType, DiscretizationType>;
            using Prescription = typename DiscretizationType::NormalizationPrescriptionType;

        public:
            static constexpr bool value = [] {
                if constexpr (
                    std::same_as<Prescription, normalization::Unnormalized> ||
                    normalization::PhysicalRieszDiagonalPrescription<Prescription>) {
                    return true;
                } else {
                    return normalization::RuntimePreparedNormalizationOperation<Problem>;
                }
            }();
        };
    } // namespace detail

    /*
     * A model and a discretization are separate compile-time choices. Their
     * pairing is valid only when the normalization plan covers the inferred
     * form, the selected physical runtime supports it, and a third-party
     * runtime policy provides its exact preparation operation. Keeping this
     * as a detection-safe public factory boundary rejects incomplete policies
     * at discretize(), before a solver-facing problem can be constructed.
     */
    template <typename Model, typename Discretization>
    concept StellarEquilibriumModelDiscretizationCompatible =
        detail::StellarEquilibriumModelDiscretizationOperationAudit<
            std::remove_cvref_t<Model>,
            std::remove_cvref_t<Discretization>>::value;

    namespace detail {
        /* The structurally formed problem type is needed to probe the ADL
         * operation without a recursive concept. Its constructor remains
         * private, and this factory is the single construction authority after
         * the complete public compatibility contract has succeeded. */
        struct StellarEquilibriumProblemFactory final {
            template <StellarEquilibriumModel Model, StellarDiscretizationType Discretization>
                requires StellarEquilibriumModelDiscretizationCompatible<Model, Discretization>
            [[nodiscard]] static auto Create(
                Model &&stellarModel,
                Discretization discretization
            ) {
                using ModelType = std::remove_cvref_t<Model>;
                using DiscretizationType = std::remove_cvref_t<Discretization>;
                return StellarEquilibriumProblem<ModelType, DiscretizationType>{
                    std::forward<Model>(stellarModel),
                    std::move(discretization)
                };
            }
        };
    } // namespace detail

    template <StellarEquilibriumModel Model, StellarDiscretizationType Discretization>
        requires StellarEquilibriumModelDiscretizationCompatible<Model, Discretization>
    [[nodiscard]] auto discretize(
        Model &&stellarModel,
        Discretization discretization
    ) {
        return detail::StellarEquilibriumProblemFactory::Create(
            std::forward<Model>(stellarModel),
            std::move(discretization)
        );
    }

    template <StellarEquilibriumModel Model>
        requires StellarEquilibriumModelDiscretizationCompatible<Model, StellarDiscretization>
    [[nodiscard]] auto discretize(
        Model &&stellarModel,
        fem::FEM &finiteElementModel
    ) {
        return discretize(std::forward<Model>(stellarModel), StellarDiscretization{finiteElementModel});
    }
} // namespace mean_field::equilibrium