module;

#include <memory>
#include <string>

#include <mfem.hpp>
#include <stroid/stroid.h>

export module mean_field:fem;

export import :physics.contexts;
export import :boundary.contexts;
export import :mapping.domain_mapper;
export import :utils.misc;
export import :utils.user;
export import :quadrature.mfem;
export import :field.mfem;

export namespace mean_field::fem {
    using GravityField      = field::Field<field::Gravity>;
    using DisplacementField = field::Field<field::Displacement>;
    using DensityField      = field::Field<field::Density>;
    using EnthalpyField     = field::Field<field::Enthalpy>;

    struct FEM {
        // =====================================================================
        // Mesh
        // =====================================================================

        stroid::StroidMesh smesh;
        std::unique_ptr<mfem::ParMesh> mesh;

        // =====================================================================
        // Compile-time field descriptors
        // =====================================================================

        GravityField gravityField;
        DisplacementField displacementField;
        DensityField densityField;
        EnthalpyField enthalpyField;

        // =====================================================================
        // Gravity field
        //
        // Collection members are declared before their corresponding spaces so
        // that the spaces are destroyed first.
        // =====================================================================

        std::unique_ptr<mfem::FiniteElementCollection> gravityPotentialFec;

        std::unique_ptr<mfem::ParFiniteElementSpace> gravityPotentialFes;

        std::unique_ptr<mfem::FiniteElementCollection> gravityFluxFec;

        std::unique_ptr<mfem::ParFiniteElementSpace> gravityFluxFes;

        // =====================================================================
        // Displacement field
        // =====================================================================

        std::unique_ptr<mfem::FiniteElementCollection> displacementFec;

        std::unique_ptr<mfem::ParFiniteElementSpace> displacementFes;

        std::unique_ptr<mfem::ParGridFunction> displacement;

        // =====================================================================
        // Density field
        // =====================================================================

        std::unique_ptr<mfem::FiniteElementCollection> densityFec;

        std::unique_ptr<mfem::ParFiniteElementSpace> densityFes;

        // =====================================================================
        // Specific-enthalpy field
        // =====================================================================

        std::unique_ptr<mfem::FiniteElementCollection> enthalpyFec;

        std::unique_ptr<mfem::ParFiniteElementSpace> enthalpyFes;

        // =====================================================================
        // Compactification coordinate
        // =====================================================================

        std::unique_ptr<mfem::H1_FECollection> compactificationFec;

        std::unique_ptr<mfem::ParFiniteElementSpace> compactificationFes;

        std::unique_ptr<mfem::ParGridFunction> compactificationCoordinate;

        // =====================================================================
        // Domain mapping
        //
        // These are declared after displacement so that they are destroyed
        // before the displacement grid function to which mapping may refer.
        // DomainMapper is retained only for legacy integrators. New operators
        // use DomainMapperStateless exclusively.
        // =====================================================================

        std::unique_ptr<mapping::DomainMapper> mapping;

        std::unique_ptr<mapping::DomainMapperStateless> domainMapperStateless;

        // =====================================================================
        // Block layouts
        //
        // These arrays are retained only for legacy code. Canonical operator
        // layouts are defined by the compile-time forms in :utils.blocks.
        //
        // Main system:    [Displacement | Density]
        // Gravity system: [Flux | Potential]
        // =====================================================================

        mfem::Array<int> blockTrueOffsets;
        mfem::Array<int> gravityBlockTrueOffsets;

        // =====================================================================
        // Boundary conditions and domain masks
        // =====================================================================

        mfem::Array<int> essentialDisplacementTdofs;
        mfem::Array<int> vacuumDensityTdofs;
        mfem::Array<int> vacuumEnthalpyTdofs;
        mfem::Array<int> vacuumDisplacementTdofs;

        // =====================================================================
        // Global diagnostics
        // =====================================================================

        mfem::Vector com;
        mfem::DenseMatrix Q;

        // =====================================================================
        // Physics and boundary contexts
        // =====================================================================

        physics::GravityContext gravityContext;
        boundary::BoundaryContext boundaryContext;

        std::unique_ptr<quadrature::RuleFactory> quadratureFactory;

        // =====================================================================
        // Validation
        // =====================================================================

        [[nodiscard]] bool okay() const {
            return mesh != nullptr &&

                   gravityPotentialFec != nullptr && gravityPotentialFes != nullptr && gravityFluxFec != nullptr &&
                   gravityFluxFes != nullptr &&

                   displacementFec != nullptr && displacementFes != nullptr && displacement != nullptr &&

                   densityFec != nullptr && densityFes != nullptr &&

                   enthalpyFec != nullptr && enthalpyFes != nullptr &&

                   compactificationFec != nullptr && compactificationFes != nullptr &&
                   compactificationCoordinate != nullptr &&

                   mapping != nullptr && domainMapperStateless != nullptr && quadratureFactory != nullptr &&

                   blockTrueOffsets.Size() == 3 && gravityBlockTrueOffsets.Size() == 3;
        }

        [[nodiscard]] bool has_mapping() const {
            return mapping != nullptr;
        }
    };

    FEM setup_fem(
        const std::string &filename,
        const utils::Args &args,
        int extraRefine = 0
    );
} // namespace mean_field::fem