module;

#include <cmath>
#include <memory>
#include <mfem.hpp>
#include <stdexcept>

module mean_field;

import :preconditioning.gravity_field;

namespace mean_field::preconditioning {
    std::unique_ptr<mfem::HypreParMatrix> assembleGravityDivergenceSurrogate(const fem::FEM &f) {
        if (f.mesh == nullptr || f.gravityFluxFes == nullptr || f.gravityPotentialFes == nullptr ||
            f.quadratureFactory == nullptr) {
            throw std::invalid_argument(
                "The gravity divergence surrogate requires its mesh, gravity spaces, and quadrature policy."
            );
        }

        mfem::ParMixedBilinearForm divergence(f.gravityFluxFes.get(), f.gravityPotentialFes.get());
        auto integrator                                   = std::make_unique<mfem::VectorFEDivergenceIntegrator>();

        const mfem::FiniteElement &trialElement           = *f.gravityFluxFes->GetTypicalFE();
        const mfem::FiniteElement &testElement            = *f.gravityPotentialFes->GetTypicalFE();
        const mfem::ElementTransformation &transformation = *f.mesh->GetElementTransformation(0);

        f.quadratureFactory->configure_gravity_divergence(
            *integrator, quadrature::QuadratureRole::preconditioner, trialElement, testElement, transformation,
            utils::DOMAINS::ALL, quadrature::MappingKind::none
        );

        divergence.AddDomainIntegrator(integrator.release());
        divergence.Assemble();
        divergence.Finalize();

        std::unique_ptr<mfem::HypreParMatrix> assembled(divergence.ParallelAssemble());
        if (assembled == nullptr) {
            throw std::runtime_error("MFEM did not assemble the gravity divergence surrogate.");
        }
        return assembled;
    }

    std::unique_ptr<mfem::HypreParMatrix> assembleGravityPotentialSchurSurrogate(
        const fem::FEM &f,
        const mfem::Vector &trueMassDiagonal
    ) {
        if (f.gravityFluxFes == nullptr || trueMassDiagonal.Size() != f.gravityFluxFes->GetTrueVSize()) {
            throw std::invalid_argument(
                "The gravity Schur surrogate requires one mass-diagonal entry per true gravity-gradient DOF."
            );
        }

        mfem::Vector inverseMassDiagonal(trueMassDiagonal);
        for (int index = 0; index < inverseMassDiagonal.Size(); ++index) {
            const double entry = inverseMassDiagonal(index);
            if (!std::isfinite(entry) || entry <= 0.0) {
                throw std::invalid_argument(
                    "The gravity Schur surrogate encountered a non-positive or non-finite mass diagonal."
                );
            }
            inverseMassDiagonal(index) = 1.0 / entry;
        }

        std::unique_ptr<mfem::HypreParMatrix> divergence = assembleGravityDivergenceSurrogate(f);
        std::unique_ptr<mfem::HypreParMatrix> inverseMassDivergenceTranspose(divergence->Transpose());
        inverseMassDivergenceTranspose->ScaleRows(inverseMassDiagonal);

        std::unique_ptr<mfem::HypreParMatrix> schur(
            mfem::ParMult(divergence.get(), inverseMassDivergenceTranspose.get())
        );
        if (schur == nullptr) {
            throw std::runtime_error("MFEM did not assemble the gravity potential-Schur surrogate.");
        }
        return schur;
    }
} // namespace mean_field::preconditioning