module;

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

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 namespace mean_field::fem {
    struct FEM {

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

        // =====================================
        // 2. Finite Element Collections
        // =====================================
        // H1 (Continuous): For Gravitational Potential (Phi) and Velocity (v)
        std::unique_ptr<mfem::FiniteElementCollection> H1_fec;

        // L2 (Discontinuous): For Density (rho) to fix O-grid boundary scalloping
        std::unique_ptr<mfem::FiniteElementCollection> L2_fec;

        // H(div)/RT space for gravitational field
        std::unique_ptr<mfem::RT_FECollection> RT_fec;


        // =====================================
        // 3. Finite Element Spaces
        // =====================================
        std::unique_ptr<mfem::ParFiniteElementSpace> H1_fes; // Scalar continuous (Gravity)
        std::unique_ptr<mfem::ParFiniteElementSpace> Vec_H1_fes; // Vector continuous (Velocity field)
        std::unique_ptr<mfem::ParFiniteElementSpace> L2_fes; // Scalar discontinuous (Density)
        std::unique_ptr<mfem::ParFiniteElementSpace> RT_fes; // H(div)/RT space for gravitational field

        // Preconditioning for Gravity
        std::unique_ptr<mfem::ParLORDiscretization> H1_lor_disc;
        const mfem::ParFiniteElementSpace *H1_lor_fes{nullptr};

        // =====================================
        // 4. Domain Mapping
        // =====================================
        std::unique_ptr<mapping::DomainMapper> mapping;

        // =====================================
        // 5. Global System Tracking
        // =====================================
        // [ Velocity | Density | Mapping Parameters (Surface) ]
        mfem::Array<int> block_true_offsets;

        mfem::Array<int> gravity_block_true_offsets;

        // Essential Boundary Conditions for the fluid (e.g., surface stress-free)
        mfem::Array<int> ess_v_tdofs;

        // Elements entirely in the vacuum domain where fluid equations are not solved
        mfem::Array<int> vacuum_tdof_rho;
        mfem::Array<int> vacuum_tdof_v;

        // =====================================
        // 6. Multiphysics State & Integration
        // =====================================
        mfem::Vector com;
        mfem::DenseMatrix Q;

        int int_order{3};
        std::unique_ptr<mfem::IntegrationRule> int_rule;

        physics::GravityContext gravity_context;
        boundary::BoundaryContext boundary_context;

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

        // =====================================
        // 7. Utilities
        // =====================================
        [[nodiscard]] bool okay() const {
            return (mesh != nullptr) &&
                (H1_fec != nullptr) && (L2_fec != nullptr) && (RT_fec != nullptr) &&
                (H1_fes != nullptr) && (Vec_H1_fes != nullptr) && (L2_fes != nullptr) && (RT_fes != nullptr);
        }

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


    FEM setup_fem(const std::string &filename, const utils::Args &args, int extra_refine = 0);
}