module;

#include <cmath>
#include <memory>
#include <mfem.hpp>
#include <mpi.h>
#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;
    }

    const mfem::HypreParMatrix &PreparedGravityDivergenceCache::Get(const fem::FEM &f) {
        if (!f.mesh || !f.gravityFluxFes || !f.gravityPotentialFes || !f.quadratureFactory)
            throw std::invalid_argument("Fixed divergence caching requires a complete gravity discretization.");
        const auto *mesh = f.mesh.get();
        const auto *flux = f.gravityFluxFes.get(), *potential = f.gravityPotentialFes.get();
        const int localChanged = !m_matrix || m_mesh != mesh || m_flux != flux || m_potential != potential ||
            m_policy != f.quadratureFactory.get() || m_meshSequence != mesh->GetSequence() ||
            m_nodesSequence != mesh->GetNodesSequence() || m_fluxSequence != flux->GetSequence() ||
            m_potentialSequence != potential->GetSequence();
        int changed = 0;
        if (MPI_Allreduce(&localChanged, &changed, 1, MPI_INT, MPI_MAX, f.mesh->GetComm()) != MPI_SUCCESS)
            throw std::runtime_error("Could not synchronize fixed divergence cache validity.");
        if (changed) {
            auto matrix = assembleGravityDivergenceSurrogate(f);
            m_matrix = std::move(matrix);
            m_mesh = mesh; m_flux = flux; m_potential = potential; m_policy = f.quadratureFactory.get();
            m_meshSequence = mesh->GetSequence(); m_nodesSequence = mesh->GetNodesSequence();
            m_fluxSequence = flux->GetSequence(); m_potentialSequence = potential->GetSequence();
            ++m_builds;
        }
        return *m_matrix;
    }

    std::unique_ptr<mfem::HypreParMatrix> assembleGravityPotentialSchurSurrogate(
        const fem::FEM &f,
        const mfem::Vector &trueMassDiagonal,
        PreparedGravityDivergenceCache *divergenceCache
    ) {
        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> uncachedDivergence;
        if (!divergenceCache) uncachedDivergence = assembleGravityDivergenceSurrogate(f);
        const mfem::HypreParMatrix &divergence = divergenceCache ? divergenceCache->Get(f) : *uncachedDivergence;
        std::unique_ptr<mfem::HypreParMatrix> inverseMassDivergenceTranspose(divergence.Transpose());
        inverseMassDivergenceTranspose->ScaleRows(inverseMassDiagonal);

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