module;
#include <string>
#include <memory>

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

module mean_field;
import :boundary.contexts;
import :mapping.coefficients;
import :utils.misc;
import :utils.user;


namespace mean_field::fem {
    FEM setup_fem(const std::string &filename, const utils::Args &args, const int extra_refine) {
        FEM fem;

        //==================================================================
        // Section 1: Mesh and FE Space Setup
        //==================================================================
        fem.smesh = stroid::IO::LoadStroidMesh(filename).value();
        if (extra_refine > 0) {
            stroid::refinement::UniformRefinement(fem.smesh, extra_refine);
        }

        fem.mesh = std::make_unique<mfem::ParMesh>(MPI_COMM_WORLD, *fem.smesh.mesh);
        fem.mesh->EnsureNodes();

        const int geom_order = utils::get_mesh_order(*fem.mesh);
        const int dim = fem.mesh->Dimension();

        const int v_order = 2;
        const int rho_order = 2;
        const int p = rho_order ;

        const int cb_type = mfem::BasisType::GaussLobatto;
        const int ob_type = mfem::BasisType::IntegratedGLL;

        fem.RT_fec = std::make_unique<mfem::RT_FECollection>(p, dim, cb_type, ob_type);
        fem.RT_fes = std::make_unique<mfem::ParFiniteElementSpace>(fem.mesh.get(), fem.RT_fec.get());

        fem.H1_fec = std::make_unique<mfem::H1_FECollection>(v_order, dim);
        fem.L2_fec = std::make_unique<mfem::L2_FECollection>(rho_order, dim);

        // Gravity (Scalar H1) and Velocity (Vector H1)
        fem.H1_fes = std::make_unique<mfem::ParFiniteElementSpace>(fem.mesh.get(), fem.H1_fec.get());
        fem.Vec_H1_fes = std::make_unique<mfem::ParFiniteElementSpace>(fem.mesh.get(), fem.H1_fec.get(), dim,
                                                                       mfem::Ordering::byNODES);

        // Density & Pressure (Scalar Discontinuous L2)
        fem.L2_fes = std::make_unique<mfem::ParFiniteElementSpace>(fem.mesh.get(), fem.L2_fec.get());

        //==================================================================
        // Section 2: Domain Mapping
        //==================================================================
        auto [r_star_ref, r_inf_ref] = utils::discover_bounds(fem.mesh.get(), 3)
                .or_else([](const boundary::BoundsError &err)-> std::expected<boundary::Bounds, boundary::BoundsError> {
                    throw std::runtime_error("Unable to determine vacuum domain reference boundary...");
                }).value();

        fem.mapping = std::make_unique<mapping::DomainMapper>(r_star_ref, r_inf_ref);

        //==================================================================
        // Section 3: Multi-physics Block-offsets
        //==================================================================
        fem.block_true_offsets.SetSize(3);
        fem.block_true_offsets[0] = 0;
        fem.block_true_offsets[1] = fem.Vec_H1_fes->GetTrueVSize();
        fem.block_true_offsets[2] = fem.block_true_offsets[1] + fem.L2_fes->GetTrueVSize();

        fem.gravity_block_true_offsets.SetSize(3);
        fem.gravity_block_true_offsets[0] = 0;
        fem.gravity_block_true_offsets[1] = fem.RT_fes->GetTrueVSize();
        fem.gravity_block_true_offsets[2] = fem.gravity_block_true_offsets[1] + fem.L2_fes->GetTrueVSize();

        //==================================================================
        // Section 4: Multipole BC setup.
        //==================================================================
        fem.com.SetSize(dim);
        fem.com = 0.0;
        fem.Q.SetSize(dim, dim);
        fem.Q = 0.0;

        //==================================================================
        // Section 5: Integration Rules
        //==================================================================
        MFEM_ASSERT(fem.mesh->GetElementGeometry(0) == mfem::Geometry::CUBE,
                    "Currently only hexahedral meshes are supported");
        const int element_order = fem.H1_fes->GetMaxElementOrder();
        fem.int_order = 2 * element_order + geom_order - 2 + args.quad_boost;

        fem.int_rule = std::make_unique<mfem::IntegrationRule>(mfem::IntRules.Get(mfem::Geometry::CUBE, fem.int_order));

        //==================================================================
        // Section 6: Essential Boundaries & Domain Masks
        //==================================================================
        fem.ess_v_tdofs.SetSize(0);

        populate_element_mask(fem.mesh.get(), utils::DOMAINS::STELLAR, fem.gravity_context.stellar_mask);

        const int n_bdr_attrs = fem.mesh->bdr_attributes.Max();
        fem.boundary_context.inf_bounds.SetSize(n_bdr_attrs);
        fem.boundary_context.stellar_bounds.SetSize(n_bdr_attrs);

        fem.boundary_context.inf_bounds = 0;
        fem.boundary_context.stellar_bounds = 0;

        fem.boundary_context.inf_bounds[static_cast<int>(boundary::Boundaries::INF_SURFACE) - 1] = 1;
        fem.boundary_context.stellar_bounds[static_cast<int>(boundary::Boundaries::STELLAR_SURFACE) - 1] = 1;

        //==================================================================
        // Section 7: Gravity Context Setup
        //==================================================================
        fem.gravity_context.minres = std::make_unique<mfem::MINRESSolver>(fem.mesh->GetComm());
        fem.gravity_context.minres->SetRelTol(1e-12);
        fem.gravity_context.minres->SetAbsTol(1e-12);
        fem.gravity_context.minres->SetMaxIter(1000);
        fem.gravity_context.minres->SetPrintLevel(0);

        fem.gravity_context.prec_Phi = std::make_unique<mfem::HypreBoomerAMG>();
        fem.gravity_context.prec_Phi->SetPrintLevel(0);

        fem.gravity_context.block_prec = std::make_unique<mfem::BlockDiagonalPreconditioner>(fem.gravity_block_true_offsets);

        fem.gravity_context.minres->SetPreconditioner(*fem.gravity_context.block_prec);
        //=========================================================
        // Section 10: Set All vacuum elements true degrees of freedom
        //=========================================================
        {
            mfem::Array<int> vacuum_mask;
            utils::populate_element_mask(fem.mesh.get(), utils::DOMAINS::VACUUM, vacuum_mask);

            utils::populate_domain_tdofs(fem.Vec_H1_fes.get(), vacuum_mask, fem.vacuum_tdof_v);
            utils::populate_domain_tdofs(fem.L2_fes.get(), vacuum_mask, fem.vacuum_tdof_rho);
        }

        const quadrature::QuadratureOptions& quadrature_options = args.quadrature;

        if (quadrature_options.validation.reject_negative_boosts && quadrature_options.global_boost < 0) {
            throw std::invalid_argument("Global quadrature boost cannot be negative.");
        }

        quadrature::RuleSet quadrature_rule_set = quadrature::make_rule_set(quadrature_options.mode, quadrature_options.global_boost);

        if (quadrature_options.fallback_fixed_order.has_value()) {
            if (*quadrature_options.fallback_fixed_order < 0) {
                throw std::invalid_argument("Fallback quadrature order cannot be negative.");
            }
            quadrature_rule_set.fallback.fixed_order = quadrature_options.fallback_fixed_order;
        }

        auto apply_quadrature_options = [&quadrature_options](quadrature::RuleControl& rule_control, const quadrature::QuadratureTermOptions& term_options) {
            if (term_options.fixed_order.has_value() && *term_options.fixed_order < 0) {
                throw std::invalid_argument("Fixed quadrature order cannot be negative.");
            }

            if (quadrature_options.validation.reject_negative_boosts && term_options.additional_boost < 0) {
                throw std::invalid_argument("Term quadrature boost cannot be negative.");
            }

            rule_control.boost += term_options.additional_boost;

            if (term_options.fixed_order.has_value()) {
                rule_control.fixed_order = term_options.fixed_order;
            }
        };

        apply_quadrature_options(quadrature_rule_set.gravity_hdiv_mass, quadrature_options.gravity_hdiv_mass);
        apply_quadrature_options(quadrature_rule_set.gravity_divergence, quadrature_options.gravity_divergence);
        apply_quadrature_options(quadrature_rule_set.gravity_source, quadrature_options.gravity_source);
        apply_quadrature_options(quadrature_rule_set.gravity_boundary, quadrature_options.gravity_boundary);
        apply_quadrature_options(quadrature_rule_set.density_projection, quadrature_options.density_projection);
        apply_quadrature_options(quadrature_rule_set.mass_conservation, quadrature_options.mass_conservation);
        apply_quadrature_options(quadrature_rule_set.center_of_mass, quadrature_options.center_of_mass);
        apply_quadrature_options(quadrature_rule_set.quadrupole, quadrature_options.quadrupole);
        apply_quadrature_options(quadrature_rule_set.gravitational_energy, quadrature_options.gravitational_energy);
        apply_quadrature_options(quadrature_rule_set.virial, quadrature_options.virial);
        apply_quadrature_options(quadrature_rule_set.error_norm, quadrature_options.error_norm);

        apply_quadrature_options(quadrature_rule_set.roles.discretization, quadrature_options.roles.discretization);
        apply_quadrature_options(quadrature_rule_set.roles.preconditioner, quadrature_options.roles.preconditioner);
        apply_quadrature_options(quadrature_rule_set.roles.diagnostic, quadrature_options.roles.diagnostic);
        apply_quadrature_options(quadrature_rule_set.roles.projection, quadrature_options.roles.projection);

        fem.quadrature_factory = std::make_unique<quadrature::RuleFactory>(quadrature::Policy(std::move(quadrature_rule_set)));
        return fem;
    }
}