module;

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

#include <mfem.hpp>

export module mean_field:operators.prepared_central_density_stellar_equilibrium;

export import :model.compiled_fixed_central_density;
export import :operators.prepared_central_density;
export import :operators.prepared_stellar_equilibrium;

export namespace mean_field::operators {
    using CentralDensityStellarEquilibriumSpecificationModel = model::StellarModel<
        models::
            SpecificationSet<eos::Polytrope, models::FixedTotalMass, surface::Isobaric, models::FixedCentralDensity>>;

    using CentralDensityStellarEquilibriumForm = utils::blocks::central_density_bordered_stellar_equilibrium_form;
    using CentralDensityStellarEquilibriumJacobianForm =
        utils::blocks::central_density_bordered_stellar_equilibrium_jacobian_form;
    using CentralDensityStellarEquilibriumLayout = utils::blocks::form_layout<CentralDensityStellarEquilibriumForm>;
    using CentralDensityStellarEquilibriumSystemManifest = EquilibriumSystemManifest<
        CentralDensityStellarEquilibriumSpecificationModel,
        CentralDensityStellarEquilibriumForm,
        CentralDensityStellarEquilibriumJacobianForm>;

    using CentralDensityStellarEquilibriumRootManifest = CentralDensityStellarEquilibriumSystemManifest;

    struct PreparedCentralDensityStellarEquilibriumReport final {
        PreparedStellarEquilibriumReport physical;
        PreparedCentralDensityReport phase;
        bool assembledResidual{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return physical.DidAnyWork() || phase.DidAnyWork() || assembledResidual;
        }
    };

    class PreparedCentralDensityStellarEquilibriumOperator final : public mfem::Operator {
    public:
        PreparedCentralDensityStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            models::CompiledFixedMass fixedMassConstraint,
            PressureSurfaceConstraintView surfaceConstraint,
            deformation::PreparedDomainDeformationRuntime domainDeformation,
            models::CompiledFixedCentralDensity centralDensity
        )
            : PreparedCentralDensityStellarEquilibriumOperator(
                  f,
                  std::make_unique<PreparedStellarEquilibriumOperator>(
                      f,
                      domainMapper,
                      equationOfState,
                      std::move(fixedMassConstraint),
                      surfaceConstraint,
                      std::move(domainDeformation)
                  ),
                  std::move(centralDensity),
                  MakeCenterDofMap(f)
              ) {
        }

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

        PreparedCentralDensityStellarEquilibriumReport Prepare(
            const mfem::Vector &state,
            const StellarEquilibriumDependencies &dependencies,
            const physics::RigidRotation &rotation
        );

        void BuildResidual(mfem::Vector &residual) const;

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override;

        [[nodiscard]] bool IsPrepared() const noexcept;
        [[nodiscard]] const CentralDensityStellarEquilibriumLayout &GetLayout() const noexcept;
        [[nodiscard]] const CentralDensityStellarEquilibriumRootManifest &GetRootManifest() const noexcept;
        [[nodiscard]] const PreparedStellarEquilibriumOperator &GetPhysicalOperator() const noexcept;
        [[nodiscard]] const PreparedCentralDensityConstraint &GetCentralDensityConstraint() const noexcept;
        [[nodiscard]] RootConstraintReport GetFixedMassReport() const;
        [[nodiscard]] CentralDensityConstraintReport GetCentralDensityReport() const;

    private:
        static field::FieldPointDofMap MakeCenterDofMap(const fem::FEM &f);

        PreparedCentralDensityStellarEquilibriumOperator(
            fem::FEM &f,
            std::unique_ptr<PreparedStellarEquilibriumOperator> physicalOperator,
            models::CompiledFixedCentralDensity centralDensity,
            field::FieldPointDofMap centerDof
        );

        void AssembleResidual();
        void VerifyPrepared() const;

        std::unique_ptr<PreparedStellarEquilibriumOperator> m_physicalOperator;
        models::CompiledFixedCentralDensity m_centralDensity;
        PreparedCentralDensityConstraint m_phaseConstraint;
        CentralDensityStellarEquilibriumRootManifest m_rootManifest;
        mfem::Vector m_cachedResidual;
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators