module;

#include <compare>
#include <concepts>
#include <cstdint>
#include <type_traits>
#include <utility>

#include <mfem.hpp>

export module mean_field:operators.prepared_stellar_equilibrium;

export import :eos.polytrope;
export import :fem;
export import :field.mfem;
export import :mapping.domain_mapper;
export import :model.stellar;
export import :model.typed_stellar;
export import :operators.context.gravity_field;
export import :operators.gravity_field;
export import :operators.gravity_field_jacobian;
export import :operators.prepared_barotropic_closure;
export import :operators.prepared_displacement_residual;
export import :operators.prepared_hydrostatic_equilibrium;
export import :operators.prepared_mass_normalization;
export import :operators.prepared_surface_constraint;
export import :operators.root_manifest;
export import :physics.rigid_rotation;
export import :utils.blocks;

export namespace mean_field::operators {
    struct StellarEquilibriumDependencyStamp final {
        std::uint64_t identity{0};
        std::uint64_t revision{0};

        constexpr auto operator<=>(const StellarEquilibriumDependencyStamp &) const = default;
    };

    struct StellarEquilibriumDependencies final {
        StellarEquilibriumDependencyStamp discretization;
        StellarEquilibriumDependencyStamp density;
        StellarEquilibriumDependencyStamp surfaceDeformation;
        StellarEquilibriumDependencyStamp gravityGradient;
        StellarEquilibriumDependencyStamp gravityPotential;
        StellarEquilibriumDependencyStamp enthalpy;
        StellarEquilibriumDependencyStamp bernoulliConstant;
        StellarEquilibriumDependencyStamp rotation;
        StellarEquilibriumDependencyStamp targetMass;

        constexpr auto operator<=>(const StellarEquilibriumDependencies &) const = default;
    };

    struct PreparedStellarEquilibriumReport final {
        context::gravity_field::GravityFieldPreparationReport gravity;
        PreparedBarotropicClosureReport barotropicClosure;
        PreparedHydrostaticEquilibriumReport hydrostatic;
        PreparedDisplacementResidualReport displacement;
        PreparedMassNormalizationReport massNormalization;
        PreparedSurfaceConstraintReport surfaceConstraint;
        deformation::DomainDeformationGeometryReport generatedGeometry;
        StellarEquilibriumDependencyStamp generatedDisplacement;
        bool generatedVolumeDisplacement{false};
        bool assembledResidual{false};

        [[nodiscard]] bool DidAnyChildWork() const noexcept {
            return gravity.DidAnyWork() || barotropicClosure.DidAnyWork() || hydrostatic.DidAnyWork() ||
                   displacement.DidAnyWork() || massNormalization.DidAnyWork() || surfaceConstraint.DidAnyWork();
        }

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return DidAnyChildWork() || generatedVolumeDisplacement || assembledResidual;
        }
    };

    struct PreparedStellarEquilibriumStatistics final {
        std::uint64_t residualAssemblies{0};
        std::uint64_t residualApplications{0};
        std::uint64_t jacobianApplications{0};
        std::uint64_t generatedGeometryBuilds{0};

        constexpr auto operator<=>(const PreparedStellarEquilibriumStatistics &) const = default;
    };

    using StellarEquilibriumLayout =
        utils::blocks::form_layout<utils::blocks::surface_deformed_stellar_equilibrium_form>;

    using StellarEquilibriumSpecificationModel =
        model::StellarModel<models::SpecificationSet<eos::Polytrope, models::FixedTotalMass, surface::Isobaric>>;

    using StellarEquilibriumSystemManifest = EquilibriumSystemManifest<
        StellarEquilibriumSpecificationModel,
        utils::blocks::surface_deformed_stellar_equilibrium_form,
        utils::blocks::surface_deformed_stellar_equilibrium_jacobian_form>;

    using StellarEquilibriumRootManifest = StellarEquilibriumSystemManifest;

    class PreparedStellarEquilibriumOperator final : public mfem::Operator {
    public:
        template <models::StellarModelType Model>
            requires std::same_as<
                         typename std::remove_cvref_t<Model>::EquationOfStateType,
                         eos::Polytrope> &&
                     SingleFieldPressureSurfaceConstraintFor<
                         typename std::remove_cvref_t<Model>::SurfaceConstraintType,
                         field::Enthalpy> &&
                     std::is_lvalue_reference_v<Model &&>
        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            Model &&stellarModel
        )
            : PreparedStellarEquilibriumOperator(
                  f,
                  domainMapper,
                  stellarModel.equationOfState(),
                  models::compileConstraint(models::FixedTotalMass{dimensions::MassValue{stellarModel.targetMass()}}),
                  PressureSurfaceConstraintView{stellarModel.compiledSurfaceConstraint()},
                  deformation::PreparedDomainDeformationRuntime{stellarModel.compileDomainDeformation(f)}
              ) {
        }

        /*
         * Authoritative construction path for a compiled equilibrium system.
         * The caller owns the EOS and compiled surface constraint for this
         * operator's lifetime; the remaining compiled contributions are
         * transferred into the operator.
         */
        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            models::CompiledFixedMass fixedMassConstraint,
            PressureSurfaceConstraintView surfaceConstraint,
            deformation::PreparedDomainDeformationRuntime domainDeformation
        );

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

        PreparedStellarEquilibriumReport 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]] double GetTargetMass() const noexcept;
        [[nodiscard]] const StellarEquilibriumLayout &GetLayout() const noexcept;
        [[nodiscard]] const StellarEquilibriumRootManifest &GetRootManifest() const noexcept;
        [[nodiscard]] RootStateView<utils::blocks::surface_deformed_stellar_equilibrium_form>
        GetRootStateView(const mfem::Vector &state) const;
        [[nodiscard]] ResidualView<utils::blocks::surface_deformed_stellar_equilibrium_form>
        GetResidualView(mfem::Vector &residual) const;
        [[nodiscard]] RootConstraintReport GetFixedMassReport() const;
        [[nodiscard]] const StellarEquilibriumDependencies &GetDependencies() const;
        [[nodiscard]] const PreparedStellarEquilibriumStatistics &GetStatistics() const noexcept;

        [[nodiscard]] const context::gravity_field::GravityFieldLinearizationContext &
        GetGravityContext() const noexcept;
        [[nodiscard]] const GravityFieldOperator &GetGravityOperator() const noexcept;
        [[nodiscard]] const GravityFieldJacobianOperator &GetGravityJacobianOperator() const noexcept;
        [[nodiscard]] const PreparedBarotropicClosureOperator &GetBarotropicClosureOperator() const noexcept;
        [[nodiscard]] const context::barotropic::BarotropicClosureLinearizationContext &
        GetBarotropicClosureContext() const noexcept;
        [[nodiscard]] const PreparedHydrostaticEquilibriumOperator &GetHydrostaticOperator() const noexcept;
        [[nodiscard]] const PreparedDisplacementResidualOperator &GetDisplacementOperator() const noexcept;
        [[nodiscard]] const PreparedMassNormalizationOperator &GetMassNormalizationOperator() const noexcept;
        [[nodiscard]] const PreparedPressureSurfaceConstraint &GetSurfaceConstraintOperator() const noexcept;
        [[nodiscard]] const deformation::PreparedDomainDeformationRuntime &GetDomainDeformation() const noexcept;
        [[nodiscard]] const mfem::Vector &GetSurfaceDeformationParameters() const;
        [[nodiscard]] const mfem::Vector &GetGeneratedVolumeDisplacement() const;
        [[nodiscard]] const mfem::Vector &GetFullMechanicalResidual() const;
        [[nodiscard]] const StellarEquilibriumDependencyStamp &GetGeneratedDisplacementDependency() const;

    private:
        struct ConstructionData;

        static ConstructionData MakeConstructionData(
            fem::FEM &f,
            deformation::PreparedDomainDeformationRuntime domainDeformation
        );

        PreparedStellarEquilibriumOperator(
            fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            models::CompiledFixedMass fixedMassConstraint,
            PressureSurfaceConstraintView surfaceConstraint,
            ConstructionData constructionData
        );

        void AssembleResidual();
        void VerifyPrepared() const;

        StellarEquilibriumRootManifest m_rootManifest;
        mfem::Array<int> m_gravityStateOffsets;

        context::gravity_field::GravityFieldLinearizationContext m_gravityContext;
        GravityFieldJacobianOperator m_gravityJacobianOperator;
        GravityFieldOperator m_gravityOperator;

        PreparedBarotropicClosureOperator m_barotropicClosureOperator;
        PreparedHydrostaticEquilibriumOperator m_hydrostaticOperator;
        PreparedDisplacementResidualOperator m_displacementOperator;
        PreparedMassNormalizationOperator m_massNormalizationOperator;
        PreparedPressureSurfaceConstraint m_surfaceConstraintOperator;
        deformation::PreparedDomainDeformationRuntime m_domainDeformation;

        StellarEquilibriumDependencies m_preparedDependencies;
        StellarEquilibriumDependencyStamp m_generatedDisplacementDependency;
        deformation::DomainDeformationGeometryReport m_generatedGeometryReport;
        mfem::Vector m_surfaceDeformationParameters;
        mfem::Vector m_generatedVolumeDisplacement;
        mfem::Vector m_fullMechanicalResidual;
        mfem::Vector m_cachedResidual;
        models::CompiledFixedMass m_fixedMassConstraint;

        mutable PreparedStellarEquilibriumStatistics m_statistics;
        bool m_isPrepared{false};

        mfem::Vector m_gravityState;

        mutable mfem::Vector m_gravityDirection;
        mutable mfem::Vector m_volumeDisplacementDirection;
        mutable mfem::Vector m_fullMechanicalAction;
        mutable mfem::Vector m_surfaceShapeAction;
        mutable mfem::Vector m_pullbackDerivativeAction;
    };
} // namespace mean_field::operators