module;

#include <concepts>
#include <cstddef>
#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 :model.typed_stellar;
export import :operators.prepared_central_density_stellar_equilibrium;
export import :surface.compiler;

export namespace mean_field::equilibrium {
    template <typename Candidate>
    concept StellarEquilibriumModel = model::StellarModelType<Candidate> && requires {
        requires std::remove_cvref_t<Candidate>::template containsSpecification<eos::Polytrope>;
        requires std::remove_cvref_t<Candidate>::template containsSpecification<surface::Isobaric>;
        requires std::remove_cvref_t<Candidate>::template containsSpecification<models::FixedTotalMass>;
        requires std::remove_cvref_t<Candidate>::specificationCount ==
                     3 + static_cast<std::size_t>(
                             std::remove_cvref_t<Candidate>::template containsSpecification<models::FixedCentralDensity>
                         );
    };

    template <StellarEquilibriumModel Model> class StellarEquilibriumProblem final {
    public:
        using ModelType = std::remove_cvref_t<Model>;

        static constexpr bool hasFixedCentralDensity =
            ModelType::template containsSpecification<models::FixedCentralDensity>;
        static constexpr bool symbolicallySquare = ModelType::symbolicallySquare;

        using PreparedOperatorType               = std::conditional_t<
            hasFixedCentralDensity,
            operators::PreparedCentralDensityStellarEquilibriumOperator,
            operators::PreparedStellarEquilibriumOperator>;
        using CompiledSurfaceConstraintType =
            surface::CompiledPressureSurfaceConstraintT<surface::BarotropicSurfaceFormulation, eos::Polytrope>;

        StellarEquilibriumProblem(
            ModelType stellarModel,
            const StellarDiscretization discretization
        )
            requires(!hasFixedCentralDensity)
            : m_stellarModel(std::move(stellarModel)),
              m_discretization(discretization),
              m_compiledSurfaceConstraint(CompileSurfaceConstraint(m_stellarModel)),
              m_preparedOperator(
                  m_discretization.finiteElementModel(),
                  m_discretization.domainMapper(),
                  m_stellarModel.template specification<eos::Polytrope>(),
                  models::compileConstraint(m_stellarModel.template specification<models::FixedTotalMass>()),
                  operators::PressureSurfaceConstraintView{m_compiledSurfaceConstraint},
                  CompileDefaultDomainDeformation(m_discretization.finiteElementModel())
              ) {
            VerifyProblem();
        }

        StellarEquilibriumProblem(
            ModelType stellarModel,
            const StellarDiscretization discretization
        )
            requires hasFixedCentralDensity
            : m_stellarModel(std::move(stellarModel)),
              m_discretization(discretization),
              m_compiledSurfaceConstraint(CompileSurfaceConstraint(m_stellarModel)),
              m_preparedOperator(
                  m_discretization.finiteElementModel(),
                  m_discretization.domainMapper(),
                  m_stellarModel.template specification<eos::Polytrope>(),
                  models::compileConstraint(m_stellarModel.template specification<models::FixedTotalMass>()),
                  operators::PressureSurfaceConstraintView{m_compiledSurfaceConstraint},
                  CompileDefaultDomainDeformation(m_discretization.finiteElementModel()),
                  models::compileConstraint(
                      m_stellarModel.template specification<models::FixedCentralDensity>(),
                      m_stellarModel.template specification<eos::Polytrope>()
                  )
              ) {
            VerifyProblem();
        }

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

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

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

        [[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 auto &GetManifest() const noexcept {
            return m_preparedOperator.GetRootManifest();
        }

        [[nodiscard]] const field::FieldBoundaryDofMap &GetPressureSurfaceRows() const noexcept {
            if constexpr (hasFixedCentralDensity) {
                return m_preparedOperator.GetPhysicalOperator().GetSurfaceConstraintOperator().GetSurfaceRows();
            } else {
                return m_preparedOperator.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
        ) {
            return m_preparedOperator.Prepare(state, dependencies, rotation);
        }

        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<surface::BarotropicSurfaceFormulation>(
                stellarModel.template specification<surface::Isobaric>(),
                stellarModel.template specification<eos::Polytrope>()
            );
        }

        [[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.");
        }

        ModelType m_stellarModel;
        StellarDiscretization m_discretization;
        CompiledSurfaceConstraintType m_compiledSurfaceConstraint;
        PreparedOperatorType m_preparedOperator;
    };

    template <StellarEquilibriumModel Model>
    [[nodiscard]] auto discretize(
        Model &&stellarModel,
        const StellarDiscretization discretization
    ) {
        using ModelType = std::remove_cvref_t<Model>;
        return StellarEquilibriumProblem<ModelType>{std::forward<Model>(stellarModel), discretization};
    }

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