module;

#include <cstdint>
#include <mfem.hpp>
#include <vector>

export module mean_field:operators.prepared_barotropic_closure;

export import :fem;
export import :mapping.domain_mapper;
export import :physics.barotrope;

export namespace mean_field::operators {
    class PreparedBarotropicClosureOperator final : public mfem::Operator {
    public:
        PreparedBarotropicClosureOperator(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper,
            const physics::PolytropicBarotrope &barotrope
        );

        void Prepare(
            const mfem::Vector &baseDensityTrue,
            const mfem::Vector &baseEnthalpyTrue,
            const mfem::Vector &displacementTrue
        );

        void Mult(
            const mfem::Vector &densityVariationTrue,
            const mfem::Vector &enthalpyVariationTrue,
            const mfem::Vector &displacementVariationTrue,
            mfem::Vector &action
        ) const;

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

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

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

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

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

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

    private:
        void VerifyPrepared() const;

        void Mult(
            const mfem::Vector &densityVariationTrue,
            const mfem::Vector &enthalpyVariationTrue,
            mfem::Vector &action
        ) const;

        struct ElementPAData {
            mfem::Array<int> densityDofs;
            mfem::Array<int> enthalpyDofs;

            mfem::DofTransformation *densityDofTransformation{nullptr};
            mfem::DofTransformation *enthalpyDofTransformation{nullptr};

            mfem::DenseMatrix densityBasis;
            mfem::DenseMatrix enthalpyBasis;

            mfem::Vector weightedResidual;
            mfem::Vector quadratureWeights;
            mfem::Vector weightedEnthalpyDerivative;
        };

        const fem::FEM &m_fem;
        const mapping::DomainMapperStateless &m_domainMapper;
        const physics::PolytropicBarotrope &m_barotrope;

        std::vector<ElementPAData> m_elements;

        mfem::Vector m_baseDensityTrue;
        mfem::Vector m_baseEnthalpyTrue;
        mfem::Vector m_baseDisplacementTrue;

        int m_densitySize{0};
        int m_enthalpySize{0};

        std::uint64_t m_preparationCount{0};
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators