module;

#include <concepts>
#include <cstddef>
#include <cstdint>
#include <expected>
#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>> { };

    } // 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 Report                             = typename PreparedOperatorType::Report;
        using PreparationResult                  = typename PreparedOperatorType::PreparationResult;
        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,
            fem::FEM &finiteElementModel
        )
            : m_stellarModel(std::make_shared<ModelType>(std::move(stellarModel))),
              m_discretization(std::move(discretization)),
              m_compiledSurfaceConstraint(CompileSurfaceConstraint(*m_stellarModel)),
              m_preparedOperator(
                  finiteElementModel,
                  m_discretization.domainMapper(),
                  m_stellarModel,
                  operators::PressureSurfaceConstraintView{m_compiledSurfaceConstraint},
                  CompileDefaultDomainDeformation(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]] MPI_Comm GetCommunicator() const & {
            return m_discretization.communicator();
        }

        [[nodiscard]] MPI_Comm GetCommunicator() const && = delete;

        [[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]] PreparationResult TryPrepare(
            const mfem::Vector &state,
            const operators::StellarEquilibriumDependencies &dependencies,
            const physics::RigidRotation &rotation
        )
            requires(generatedRotationProviderCount == 0)
        {
            auto result = m_preparedOperator.TryPrepare(state, dependencies, rotation);
            if (!result.has_value()) {
                return std::unexpected(result.error());
            }
            ++m_preparationGeneration;
            return result;
        }

        [[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;
        }

        [[nodiscard]] PreparationResult TryPrepare(
            const mfem::Vector &state,
            const operators::StellarEquilibriumDependencies &dependencies
        )
            requires(generatedRotationProviderCount == 1)
        {
            auto result = m_preparedOperator.TryPrepare(state, dependencies);
            if (!result.has_value()) {
                return std::unexpected(result.error());
            }
            ++m_preparationGeneration;
            return result;
        }

        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);
        }

        void BuildVolumeDisplacementDirection(
            const mfem::Vector &stateDirection,
            mfem::Vector &volumeDisplacementDirection
        ) const {
            m_preparedOperator.BuildVolumeDisplacementDirection(stateDirection, volumeDisplacementDirection);
        }

    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 mean_field::equilibrium

namespace mean_field::equilibrium::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>;
            fem::FEM &finiteElementModel = discretization.MutableFiniteElementModelForAssembly();
            return StellarEquilibriumProblem<ModelType, DiscretizationType>{
                std::forward<Model>(stellarModel), std::move(discretization), finiteElementModel
            };
        }

        template <
            StellarEquilibriumModel Model,
            StellarDiscretizationType Discretization>
            requires StellarEquilibriumModelDiscretizationCompatible<
                Model,
                Discretization>
        [[nodiscard]] static auto CreateOwned(
            Model &&stellarModel,
            Discretization discretization
        ) {
            using ModelType              = std::remove_cvref_t<Model>;
            using DiscretizationType     = std::remove_cvref_t<Discretization>;
            using ProblemType            = StellarEquilibriumProblem<ModelType, DiscretizationType>;
            fem::FEM &finiteElementModel = discretization.MutableFiniteElementModelForAssembly();
            return std::unique_ptr<ProblemType>{
                new ProblemType{std::forward<Model>(stellarModel), std::move(discretization), finiteElementModel}
            };
        }

        template <DiscretizedStellarEquilibriumProblem Problem>
        [[nodiscard]] static fem::FEM &MutableFiniteElementModelForProjection(Problem &problem) {
            return problem.m_discretization.MutableFiniteElementModelForAssembly();
        }

        template <DiscretizedStellarEquilibriumProblem Problem>
        [[nodiscard]] static const fem::FEM &FiniteElementModel(const Problem &problem) {
            return problem.m_discretization.RequireFiniteElementModel();
        }
    };
} // namespace mean_field::equilibrium::detail

export namespace mean_field::equilibrium {
    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{std::move(finiteElementModel)});
    }

    template <StellarEquilibriumModel Model>
        requires StellarEquilibriumModelDiscretizationCompatible<
                     Model,
                     StellarDiscretization>
    StellarEquilibriumProblem<
        std::remove_cvref_t<Model>,
        StellarDiscretization>
    discretize(
        Model &&,
        fem::FEM &
    ) = delete;
} // namespace mean_field::equilibrium