module;

#include <compare>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <vector>

#include <mfem.hpp>

export module mean_field:operators.prepared_hydrostatic_equilibrium;

export import :fem;
export import :mapping.domain_mapper;
export import :operators.context.hydrostatic_equilibrium;
export import :physics.rigid_rotation;

export namespace mean_field::operators {
    struct PreparedHydrostaticEquilibriumReport {
        context::hydrostatic::HydrostaticPreparationReport contextReport;

        bool updatedRotation{false};
        bool preparedAlgebraicJacobianBlocks{false};
        bool preparedDisplacementJacobianData{false};
        bool preparedResidual{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return contextReport.DidAnyWork() || updatedRotation || preparedAlgebraicJacobianBlocks ||
                   preparedDisplacementJacobianData || preparedResidual;
        }
    };

    struct PreparedHydrostaticAlgebraicJacobianStatistics {
        std::uint64_t preparations{0};
        std::uint64_t enthalpyApplications{0};
        std::uint64_t gravityPotentialApplications{0};
        std::uint64_t bernoulliConstantApplications{0};
        std::uint64_t combinedApplications{0};

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

    struct PreparedHydrostaticDisplacementJacobianStatistics {
        std::uint64_t preparations{0};
        std::uint64_t applications{0};

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

    struct PreparedHydrostaticCompleteJacobianStatistics {
        std::uint64_t applications{0};

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

    enum class HydrostaticJacobianInputBlock : int {
        enthalpy          = 0,
        gravityPotential  = 1,
        bernoulliConstant = 2,
        displacement      = 3
    };

    class HydrostaticJacobianBlockLayout final {
    public:
        explicit HydrostaticJacobianBlockLayout(const fem::FEM &f);

        [[nodiscard]] int Offset(HydrostaticJacobianInputBlock block) const;

        [[nodiscard]] int Size(HydrostaticJacobianInputBlock block) const;

        [[nodiscard]] int GetTotalSize() const noexcept;

        [[nodiscard]] int GetResidualSize() const noexcept;

    private:
        int m_enthalpySize{0};
        int m_gravityPotentialSize{0};
        int m_displacementSize{0};
        int m_totalSize{0};
        int m_residualSize{0};
    };

    class PreparedHydrostaticEquilibriumOperator final {
    public:
        PreparedHydrostaticEquilibriumOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper
        );

        PreparedHydrostaticEquilibriumOperator(const PreparedHydrostaticEquilibriumOperator &)            = delete;

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

        PreparedHydrostaticEquilibriumOperator(PreparedHydrostaticEquilibriumOperator &&)                 = delete;

        PreparedHydrostaticEquilibriumOperator &operator=(PreparedHydrostaticEquilibriumOperator &&)      = delete;

        PreparedHydrostaticEquilibriumReport Prepare(
            const context::hydrostatic::HydrostaticEquilibriumStateView &state,
            const context::hydrostatic::HydrostaticEquilibriumDependencies &dependencies,
            const physics::RigidRotation &rotation
        );

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

        void ApplyEnthalpyJacobianAction(
            const mfem::Vector &enthalpyVariation,
            mfem::Vector &action
        ) const;

        void ApplyGravityPotentialJacobianAction(
            const mfem::Vector &gravityPotentialVariation,
            mfem::Vector &action
        ) const;

        void ApplyBernoulliConstantJacobianAction(
            double bernoulliConstantVariation,
            mfem::Vector &action
        ) const;

        void ApplyAlgebraicJacobianAction(
            const mfem::Vector &enthalpyVariation,
            const mfem::Vector &gravityPotentialVariation,
            double bernoulliConstantVariation,
            mfem::Vector &action
        ) const;

        void ApplyDisplacementJacobianAction(
            const mfem::Vector &displacementVariation,
            mfem::Vector &action
        ) const;

        void ApplyCompleteJacobianAction(
            const mfem::Vector &enthalpyVariation,
            const mfem::Vector &gravityPotentialVariation,
            double bernoulliConstantVariation,
            const mfem::Vector &displacementVariation,
            mfem::Vector &action
        ) const;

        [[nodiscard]] bool IsPrepared() const noexcept;

        [[nodiscard]] const context::hydrostatic::HydrostaticPreparationStatistics &
        GetContextPreparationStatistics() const noexcept;

        [[nodiscard]] std::uint64_t GetResidualPreparationCount() const noexcept;

        [[nodiscard]] std::uint64_t GetResidualApplicationCount() const noexcept;

        [[nodiscard]] const PreparedHydrostaticAlgebraicJacobianStatistics &
        GetAlgebraicJacobianStatistics() const noexcept;

        [[nodiscard]] const PreparedHydrostaticDisplacementJacobianStatistics &
        GetDisplacementJacobianStatistics() const noexcept;

        [[nodiscard]] const PreparedHydrostaticCompleteJacobianStatistics &
        GetCompleteJacobianStatistics() const noexcept;

        [[nodiscard]] std::size_t GetStellarElementCount() const noexcept;

        [[nodiscard]] const fem::FEM &GetFEM() const noexcept;

        [[nodiscard]] const field::FieldDofMap &GetEnthalpyMap() const noexcept;

        [[nodiscard]] const field::FieldDofMap &GetGravityPotentialMap() const noexcept;

        [[nodiscard]] const field::FieldDofMap &GetDisplacementMap() const noexcept;

    private:
        struct ElementPAData {
            int elementId{-1};

            mfem::Array<int> enthalpyDofs;
            mfem::Array<int> gravityPotentialDofs;
            mfem::Array<int> displacementDofs;

            mfem::DofTransformation *enthalpyDofTransformation{nullptr};

            mfem::DofTransformation *gravityPotentialDofTransformation{nullptr};

            mfem::DofTransformation *displacementDofTransformation{nullptr};

            const mfem::IntegrationRule *integrationRule{nullptr};

            // Rows are quadrature points and columns are element DOFs.
            mfem::DenseMatrix enthalpyBasis;
            mfem::DenseMatrix gravityPotentialBasis;

            // Rows are quadrature points and columns are physical components.
            mfem::DenseMatrix physicalPositions;

            mfem::Vector quadratureWeights;
            std::vector<mapping::VolumeMappingContext> baseMappingContexts;

            std::optional<mapping::ElementDisplacementData> baseDisplacementData;

            std::optional<mapping::ElementCompactificationData> compactificationData;

            mfem::Vector rotationPotential;
            mfem::DenseMatrix rotationGradient;
            mfem::Vector hydrostaticImbalance;
            mfem::Vector weightedResidual;

            // Geometry-dependent algebraic Jacobian blocks.
            mfem::DenseMatrix enthalpyJacobian;
            mfem::DenseMatrix gravityPotentialJacobian;
            mfem::Vector bernoulliConstantJacobian;
        };

        void PrepareStaticPlan();
        void PrepareGeometry();
        void PrepareAlgebraicJacobianBlocks();
        void PrepareRotation();
        void PrepareBaseState();
        void FinalizeDisplacementJacobianPreparation();
        void AssembleCachedResidual();
        void VerifyPrepared() const;

        const fem::FEM &m_fem;
        const mapping::DomainMapper &m_domainMapper;

        context::hydrostatic::HydrostaticEquilibriumContext m_context;

        std::optional<physics::RigidRotation> m_rotation;
        std::vector<ElementPAData> m_elements;
        mfem::Vector m_cachedResidual;

        mutable mfem::Vector m_enthalpyVariationTrue;
        mutable mfem::Vector m_gravityPotentialVariationTrue;
        mutable mfem::Vector m_displacementVariationTrue;
        mutable mfem::Vector m_fullEnthalpyAction;

        std::uint64_t m_residualPreparationCount{0};

        mutable std::uint64_t m_residualApplicationCount{0};

        mutable PreparedHydrostaticAlgebraicJacobianStatistics m_algebraicJacobianStatistics;

        mutable PreparedHydrostaticDisplacementJacobianStatistics m_displacementJacobianStatistics;

        mutable PreparedHydrostaticCompleteJacobianStatistics m_completeJacobianStatistics;

        bool m_isPrepared{false};
    };

    class PreparedHydrostaticEquilibriumJacobianOperator final : public mfem::Operator {
    public:
        PreparedHydrostaticEquilibriumJacobianOperator(
            const fem::FEM &f,
            const PreparedHydrostaticEquilibriumOperator &preparedOperator
        );

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

        [[nodiscard]] const HydrostaticJacobianBlockLayout &GetLayout() const noexcept;

    private:
        HydrostaticJacobianBlockLayout m_layout;

        const PreparedHydrostaticEquilibriumOperator &m_preparedOperator;
    };
} // namespace mean_field::operators