#include <algorithm>
#include <array>
#include <bit>
#include <cmath>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <type_traits>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace {
    namespace backend         = mean_field::preconditioning::backend;
    namespace blocks          = mean_field::utils::blocks;
    namespace gravity_context = mean_field::operators::context::gravity_field;
    namespace preconditioning = mean_field::preconditioning;

    using FixedAMG            = backend::HypreBoomerAMG<backend::FixedCycles>;
    using AdaptiveAMG         = backend::HypreBoomerAMG<backend::SolveToTolerance>;
    using FixedGravityLDU =
        preconditioning::GravityFieldBlock<backend::Diagonal, FixedAMG, preconditioning::GravityApproximateLDU>;
    using ChebyshevGravityLDU = preconditioning::
        GravityFieldBlock<backend::MatrixFreeChebyshev, FixedAMG, preconditioning::GravityApproximateLDU>;
    using AdaptiveGravityLDU =
        preconditioning::GravityFieldBlock<backend::Diagonal, AdaptiveAMG, preconditioning::GravityApproximateLDU>;
    using ConsistentGravityLDU = preconditioning::
        GravityFieldBlock<backend::MatrixFreeChebyshev, FixedAMG, preconditioning::GravityConsistentSchurLDU>;
    using AdaptiveConsistentGravityLDU = preconditioning::
        GravityFieldBlock<backend::MatrixFreeChebyshev, AdaptiveAMG, preconditioning::GravityConsistentSchurLDU>;

    using DensityIdentity =
        preconditioning::IdentityBlock<blocks::density::mass::value, blocks::density::mass::residual>;
    using SurfaceIdentity = preconditioning::IdentityBlock<
        blocks::surface_deformation::parameters::value,
        blocks::surface_deformation::shape_equilibrium::residual>;
    using EnthalpyIdentity =
        preconditioning::IdentityBlock<blocks::enthalpy::specific::value, blocks::enthalpy::specific::residual>;
    using MassIdentity = preconditioning::IdentityBlock<
        blocks::fixed_total_mass::mass_normalization::value,
        blocks::fixed_total_mass::mass_normalization::residual>;
    using FixedGravityPlan = preconditioning::
        PreconditionerPlan<DensityIdentity, SurfaceIdentity, FixedGravityLDU, EnthalpyIdentity, MassIdentity>;
    using AdaptiveGravityPlan = preconditioning::
        PreconditionerPlan<DensityIdentity, SurfaceIdentity, AdaptiveGravityLDU, EnthalpyIdentity, MassIdentity>;

    template <typename Policy>
    mfem::Vector applyKnownFactorization(
        Policy policy,
        const mfem::Vector &rightHandSide,
        preconditioning::GravityFactorizationStatistics *statistics = nullptr
    ) {
        mfem::Vector massDiagonal(2);
        massDiagonal     = 1.0;
        auto massInverse = backend::prepare(backend::Diagonal{}, massDiagonal);

        mfem::DenseMatrix schurMatrix(1);
        schurMatrix(0, 0) = 5.0;
        auto schurInverse = backend::prepare(backend::DenseDirect{}, schurMatrix);

        mfem::DenseMatrix divergence(1, 2);
        divergence(0, 0) = 2.0;
        divergence(0, 1) = -1.0;

        preconditioning::GravityFactorizationOperator<Policy> factorization(
            policy, massInverse, schurInverse, divergence
        );
        mfem::Vector action(factorization.Height());
        action = std::numeric_limits<double>::quiet_NaN();
        factorization.Mult(rightHandSide, action);
        if (statistics != nullptr) {
            *statistics = factorization.GetStatistics();
        }
        return action;
    }

    void checkVector(
        const mfem::Vector &computed,
        const std::array<
            double,
            3> &expected
    ) {
        REQUIRE(computed.Size() == static_cast<int>(expected.size()));
        for (int index = 0; index < computed.Size(); ++index) {
            CHECK(computed(index) == Catch::Approx(expected[static_cast<std::size_t>(index)]).margin(2.0e-14));
        }
    }

    template <typename Policy> void checkExactDenseRecovery(Policy policy) {
        mfem::DenseMatrix mass(2);
        mass(0, 0)       = 2.0;
        mass(0, 1)       = 0.5;
        mass(1, 0)       = 0.5;
        mass(1, 1)       = 1.5;
        auto massInverse = backend::prepare(backend::DenseDirect{}, mass);

        mfem::DenseMatrix divergence(1, 2);
        divergence(0, 0) = 1.0;
        divergence(0, 1) = -2.0;

        mfem::Vector divergenceTranspose(2);
        divergenceTranspose(0) = 1.0;
        divergenceTranspose(1) = -2.0;
        mfem::Vector massInverseDivergenceTranspose(2);
        massInverse.Mult(divergenceTranspose, massInverseDivergenceTranspose);

        mfem::DenseMatrix schur(1);
        schur(0, 0)       = divergenceTranspose * massInverseDivergenceTranspose;
        auto schurInverse = backend::prepare(backend::DenseDirect{}, schur);

        preconditioning::GravityFactorizationOperator<Policy> factorization(
            policy, massInverse, schurInverse, divergence
        );

        mfem::Vector exact(3);
        exact(0) = 0.7;
        exact(1) = -1.2;
        exact(2) = 0.4;

        mfem::Vector rightHandSide(3);
        mfem::Vector exactGradient(exact.GetData(), 2);
        mfem::Vector gradientRightHandSide(rightHandSide.GetData(), 2);
        mass.Mult(exactGradient, gradientRightHandSide);
        gradientRightHandSide(0) += divergence(0, 0) * exact(2);
        gradientRightHandSide(1) += divergence(0, 1) * exact(2);
        rightHandSide(2) = divergence(0, 0) * exact(0) + divergence(0, 1) * exact(1);

        mfem::Vector action(3);
        action                            = 0.0;
        const double *const actionStorage = action.GetData();
        factorization.Mult(rightHandSide, action);

        CHECK(action.GetData() == actionStorage);
        for (int index = 0; index < action.Size(); ++index) {
            CHECK(action(index) == Catch::Approx(exact(index)).margin(2.0e-13));
        }

        mfem::Vector repeated(3);
        repeated = 0.0;
        factorization.Mult(rightHandSide, repeated);
        for (int index = 0; index < repeated.Size(); ++index) {
            CHECK(repeated(index) == action(index));
        }
    }

    struct DenseSchurRefinementFixture final {
        static mfem::DenseMatrix makeMass() {
            mfem::DenseMatrix result(3);
            result = 0.0;
            result(0, 0) = 4.0;
            result(1, 1) = 3.0;
            result(2, 2) = 2.0;
            result(0, 1) = result(1, 0) = 1.0;
            result(1, 2) = result(2, 1) = 1.0;
            return result;
        }

        static mfem::DenseMatrix makeDivergence() {
            mfem::DenseMatrix result(2, 3);
            result(0, 0) = 1.0;
            result(0, 1) = 2.0;
            result(0, 2) = -1.0;
            result(1, 0) = 2.0;
            result(1, 1) = -1.0;
            result(1, 2) = 1.0;
            return result;
        }

        static mfem::DenseMatrix makeSurrogate() {
            // B diag(M)^-1 B^T differs from B M^-1 B^T.  Its inverse
            // does not commute with the latter, so a reversed correction
            // order cannot pass the rational-reference checks below.
            mfem::DenseMatrix result(2);
            result(0, 0) = 25.0 / 12.0;
            result(0, 1) = result(1, 0) = -2.0 / 3.0;
            result(1, 1) = 11.0 / 6.0;
            return result;
        }

        mfem::DenseMatrix mass{makeMass()};
        mfem::DenseMatrix divergence{makeDivergence()};
        mfem::DenseMatrix surrogate{makeSurrogate()};
        backend::PreparedDenseDirect massInverse{backend::DenseDirect{}, mass};
        backend::PreparedDenseDirect schurInverse{backend::DenseDirect{}, surrogate};

        static mfem::Vector rightHandSide() {
            mfem::Vector result(5);
            result(0) = 3.0;
            result(1) = -2.0;
            result(2) = 1.0;
            result(3) = 4.0;
            result(4) = -3.0;
            return result;
        }

        mfem::Vector residual(const mfem::Vector &action, const mfem::Vector &rhs) const {
            mfem::Vector result(5);
            mfem::Vector flux(const_cast<mfem::real_t *>(action.GetData()), 3);
            mfem::Vector potential(const_cast<mfem::real_t *>(action.GetData()) + 3, 2);
            mfem::Vector gradientResidual(result.GetData(), 3);
            mfem::Vector potentialResidual(result.GetData() + 3, 2);
            mfem::Vector transposeAction(3);
            mass.Mult(flux, gradientResidual);
            divergence.MultTranspose(potential, transposeAction);
            gradientResidual += transposeAction;
            divergence.Mult(flux, potentialResidual);
            result -= rhs;
            return result;
        }
    };

    struct PreparedGeometry final {
        mean_field::fem::FEM finiteElements;
        gravity_context::GravityFieldGeometryContext context;

        explicit PreparedGeometry(const mean_field::utils::Args &arguments)
            : finiteElements(
                  mean_field::fem::setup_fem(
                      arguments.mesh_file,
                      arguments,
                      0
                  )
              ),
              context(
                  finiteElements,
                  *finiteElements.domainMapperStateless
              ) {
            mfem::Vector displacementTrue(finiteElements.displacementFes->GetTrueVSize());
            displacementTrue                = 0.0;
            const mfem::Vector displacement = context.GetDisplacementMap().gather(displacementTrue);
            context.PreparePrimal(displacement, {.value = 1}, {.value = 1});
        }
    };
} // namespace

TEST_CASE(
    "Gravity Field Blocks Expose Complete Compile-Time Ownership And Backend Contracts",
    tags::preconditioning_gravity_unit
) {
    using Form         = blocks::surface_deformed_stellar_equilibrium_form;
    using JacobianForm = blocks::surface_deformed_stellar_equilibrium_jacobian_form;

    STATIC_CHECK(preconditioning::PreconditionerComponent<FixedGravityLDU>);
    STATIC_CHECK(preconditioning::PreconditionerComponent<ChebyshevGravityLDU>);
    STATIC_CHECK(preconditioning::PreconditionerComponent<AdaptiveGravityLDU>);
    STATIC_CHECK(preconditioning::PreconditionerComponent<ConsistentGravityLDU>);
    STATIC_CHECK(preconditioning::GravityFactorizationPolicy<preconditioning::GravityConsistentSchurLDU>);
    STATIC_CHECK(preconditioning::backend::ArnoldiAdmissible<typename FixedGravityLDU::BackendType>);
    STATIC_CHECK(preconditioning::backend::ArnoldiAdmissible<typename ChebyshevGravityLDU::BackendType>);
    STATIC_CHECK_FALSE(preconditioning::backend::ArnoldiAdmissible<typename AdaptiveGravityLDU::BackendType>);
    STATIC_CHECK(preconditioning::backend::ArnoldiAdmissible<typename ConsistentGravityLDU::BackendType>);
    STATIC_CHECK_FALSE(
        preconditioning::backend::ArnoldiAdmissible<typename AdaptiveConsistentGravityLDU::BackendType>
    );
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<FixedGravityPlan, Form>);
    STATIC_CHECK(preconditioning::CompatiblePreconditionerFor<FixedGravityPlan, Form, JacobianForm>);
    STATIC_CHECK(preconditioning::StationaryLinearPreconditionerPlan<FixedGravityPlan>);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<AdaptiveGravityPlan, Form>);
    STATIC_CHECK(preconditioning::CompatiblePreconditionerFor<AdaptiveGravityPlan, Form, JacobianForm>);
    STATIC_CHECK_FALSE(preconditioning::StationaryLinearPreconditionerPlan<AdaptiveGravityPlan>);
    STATIC_CHECK(FixedGravityLDU::RequiredCouplings::size == 3);
}

TEST_CASE(
    "Gravity Factorization Policies Preserve Their Signed Block Algebra",
    tags::preconditioning_gravity_unit
) {
    mfem::Vector rightHandSide(3);
    rightHandSide(0) = 3.0;
    rightHandSide(1) = 4.0;
    rightHandSide(2) = 7.0;

    checkVector(applyKnownFactorization(preconditioning::GravityBlockDiagonal{}, rightHandSide), {3.0, 4.0, 1.4});
    checkVector(applyKnownFactorization(preconditioning::GravityLowerTriangular{}, rightHandSide), {3.0, 4.0, -1.0});
    checkVector(applyKnownFactorization(preconditioning::GravityUpperTriangular{}, rightHandSide), {5.8, 2.6, -1.4});

    preconditioning::GravityFactorizationStatistics statistics;
    checkVector(
        applyKnownFactorization(preconditioning::GravityApproximateLDU{}, rightHandSide, &statistics), {5.0, 3.0, -1.0}
    );
    CHECK(statistics.applications == 1);
    CHECK(statistics.massInverseApplications == 2);
    CHECK(statistics.potentialSchurApplications == 1);
    CHECK(statistics.divergenceApplications == 1);
    CHECK(statistics.transposeDivergenceApplications == 1);
}

TEST_CASE(
    "Exact Gravity LDU Recovers A Dense Coupled Saddle-Point System",
    tags::preconditioning_gravity_unit
) {
    checkExactDenseRecovery(preconditioning::GravityApproximateLDU{});
    checkExactDenseRecovery(preconditioning::GravityConsistentSchurLDU{.correctionSweeps = 2, .relaxation = 0.5});
}

TEST_CASE(
    "Mass-Consistent Gravity Schur Refinement Preserves The Baseline And Signed Algebra",
    tags::preconditioning_gravity_unit
) {
    DenseSchurRefinementFixture fixture;
    const mfem::Vector rhs = fixture.rightHandSide();
    preconditioning::GravityFactorizationOperator baseline(
        preconditioning::GravityApproximateLDU{}, fixture.massInverse, fixture.schurInverse, fixture.divergence
    );
    mfem::Vector baselineAction(5);
    baseline.Mult(rhs, baselineAction);
    const double baselineResidual = fixture.residual(baselineAction, rhs).Norml2();

    // These rational answers were obtained from the two dense saddle-point
    // blocks, with S_R=[3,-2;-2,59/18] and omega=1/2.  They exercise a genuinely
    // inexact Schur inverse rather than an already-exact scalar example.
    const std::array<std::array<double, 5>, 2> expected{{
        {-458734.0 / 4782969.0, 6011089.0 / 4782969.0, -7518059.0 / 4782969.0,
         -646046.0 / 531441.0, 888176.0 / 531441.0},
        {-631839955.0 / 6973568802.0, 4376327111.0 / 3486784401.0, -10946970773.0 / 6973568802.0,
         -472857412.0 / 387420489.0, 644634595.0 / 387420489.0}
    }};

    double previousResidual = baselineResidual;
    for (const int sweeps : {0, 1, 2}) {
        CAPTURE(sweeps);
        preconditioning::GravityFactorizationOperator refined(
            preconditioning::GravityConsistentSchurLDU{.correctionSweeps = sweeps, .relaxation = 0.5},
            fixture.massInverse, fixture.schurInverse, fixture.divergence
        );
        mfem::Vector action(5);
        action = std::numeric_limits<double>::quiet_NaN();
        const auto *const actionStorage = action.GetData();
        refined.Mult(rhs, action);
        CHECK(action.GetData() == actionStorage);
        for (int index = 0; index < action.Size(); ++index) {
            if (sweeps == 0) {
                CHECK(std::bit_cast<std::uint64_t>(action(index)) ==
                      std::bit_cast<std::uint64_t>(baselineAction(index)));
            } else {
                CHECK(action(index) == Catch::Approx(expected[sweeps - 1][index]).margin(2.0e-13));
            }
        }
        const auto &statistics = refined.GetStatistics();
        CHECK(statistics.applications == 1);
        CHECK(statistics.massInverseApplications == static_cast<std::uint64_t>(2 + sweeps));
        CHECK(statistics.potentialSchurApplications == static_cast<std::uint64_t>(1 + sweeps));
        CHECK(statistics.divergenceApplications == static_cast<std::uint64_t>(1 + sweeps));
        CHECK(statistics.transposeDivergenceApplications == static_cast<std::uint64_t>(1 + sweeps));
        CHECK(statistics.schurRefinementSweeps == static_cast<std::uint64_t>(sweeps));

        const mfem::Vector residual = fixture.residual(action, rhs);
        for (int index = 0; index < 3; ++index) CHECK(std::abs(residual(index)) < 2.0e-13);
        if (sweeps > 0) {
            CHECK(residual.Norml2() < previousResidual);
            previousResidual = residual.Norml2();
        }
    }
}

TEST_CASE(
    "Fixed Gravity Schur Refinement Is Linear And Has No Application History",
    tags::preconditioning_gravity_unit
) {
    DenseSchurRefinementFixture fixture;
    preconditioning::GravityFactorizationOperator refined(
        preconditioning::GravityConsistentSchurLDU{.correctionSweeps = 2, .relaxation = 0.5},
        fixture.massInverse, fixture.schurInverse, fixture.divergence
    );
    const mfem::Vector x = fixture.rightHandSide();
    mfem::Vector y(5), combination(x), px(5), py(5), actual(5), expected(5), repeat(5);
    for (int index = 0; index < y.Size(); ++index) y(index) = std::sin(0.7 * (index + 1));
    combination *= -2.3;
    combination.Add(0.4, y);
    refined.Mult(x, px);
    refined.Mult(y, py);
    refined.Mult(combination, actual);
    expected = px;
    expected *= -2.3;
    expected.Add(0.4, py);
    actual -= expected;
    CHECK(actual.Norml2() < 2.0e-13 * expected.Norml2());
    repeat = std::numeric_limits<double>::quiet_NaN();
    refined.Mult(x, repeat);
    for (int index = 0; index < repeat.Size(); ++index) CHECK(repeat(index) == px(index));

    combination = 0.0;
    refined.Mult(combination, actual);
    CHECK(actual.Norml2() == 0.0);
    CHECK(refined.GetStatistics().applications == 5);
    CHECK(refined.GetStatistics().schurRefinementSweeps == 10);
}

TEST_CASE(
    "Gravity Schur Refinement Rejects Invalid Fixed-Work Options",
    tags::preconditioning_gravity_unit
) {
    DenseSchurRefinementFixture fixture;
    const std::array<preconditioning::GravityConsistentSchurLDU, 5> invalid{{
        {.correctionSweeps = -1, .relaxation = 1.0},
        {.correctionSweeps = 0, .relaxation = 0.0},
        {.correctionSweeps = 1, .relaxation = -0.5},
        {.correctionSweeps = 1, .relaxation = std::numeric_limits<double>::infinity()},
        {.correctionSweeps = 1, .relaxation = std::numeric_limits<double>::quiet_NaN()}
    }};
    for (const auto policy : invalid) {
        CAPTURE(policy.correctionSweeps, policy.relaxation);
        CHECK_THROWS_AS(
            (preconditioning::GravityFactorizationOperator(
                policy, fixture.massInverse, fixture.schurInverse, fixture.divergence
            )),
            std::invalid_argument
        );
    }
}

TEST_CASE(
    "Assembled Gravity Divergence Matches The Prepared Matrix-Free Couplings",
    tags::preconditioning_gravity_integration
) {
    const auto arguments = test_utils::setup_args();
    PreparedGeometry geometry(arguments);

    const auto assembledDivergence = preconditioning::assembleGravityDivergenceSurrogate(geometry.finiteElements);
    const mfem::Operator &preparedDivergence = geometry.context.GetDivergenceOperator();

    const mfem::Vector flux                  = gravity_prepared_test_utils::make_deterministic_vector(
        geometry.finiteElements.gravityFluxFes->GetTrueVSize(), 0.31
    );
    mfem::Vector assembledForward(assembledDivergence->Height());
    mfem::Vector preparedForward(preparedDivergence.Height());
    assembledDivergence->Mult(flux, assembledForward);
    preparedDivergence.Mult(flux, preparedForward);

    const mfem::Vector potential = gravity_prepared_test_utils::make_deterministic_vector(
        geometry.finiteElements.gravityPotentialFes->GetTrueVSize(), 0.73
    );
    mfem::Vector assembledTranspose(assembledDivergence->Width());
    mfem::Vector preparedTranspose(preparedDivergence.Width());
    assembledDivergence->MultTranspose(potential, assembledTranspose);
    preparedDivergence.MultTranspose(potential, preparedTranspose);

    const MPI_Comm communicator = geometry.finiteElements.mesh->GetComm();
    CHECK(gravity_prepared_test_utils::relative_error(assembledForward, preparedForward, communicator) <= 2.0e-12);
    CHECK(gravity_prepared_test_utils::relative_error(assembledTranspose, preparedTranspose, communicator) <= 2.0e-12);

    const auto &gradientMap  = geometry.context.GetMassOperator().GetFluxMap();
    const auto &potentialMap = geometry.context.GetSourceOperator().GetPotentialMap();
    preconditioning::ReducedGravityDivergenceOperator reducedDivergence(preparedDivergence, gradientMap, potentialMap);
    const mfem::Vector reducedFlux =
        gravity_prepared_test_utils::make_deterministic_vector(gradientMap.reduced_size(), 0.47);
    mfem::Vector reducedAction(reducedDivergence.Height());
    reducedDivergence.Mult(reducedFlux, reducedAction);

    const mfem::Vector trueFlux = gradientMap.scatter(reducedFlux);
    mfem::Vector trueAction(potentialMap.full_size());
    preparedDivergence.Mult(trueFlux, trueAction);
    const mfem::Vector expectedReducedAction = potentialMap.gather(trueAction);
    CHECK(gravity_prepared_test_utils::relative_error(reducedAction, expectedReducedAction, communicator) <= 2.0e-14);
}

TEST_CASE(
    "Prepared Gravity Block Diagonal Is Legacy Equivalent And Allocation Stable",
    tags::preconditioning_gravity_integration
) {
    const auto arguments = test_utils::setup_args();
    PreparedGeometry geometry(arguments);

    mean_field::operators::ReducedGravityFieldPreconditioner legacy(geometry.finiteElements, geometry.context);
    const auto block = preconditioning::GravityFieldBlock(
        backend::Diagonal{}, FixedAMG{backend::FixedCycles{.cycles = 1}}, preconditioning::GravityBlockDiagonal{}
    );
    auto prepared                    = preconditioning::prepare(geometry.finiteElements, geometry.context, block);

    const mfem::Vector rightHandSide = gravity_prepared_test_utils::make_deterministic_vector(prepared.Width(), 0.59);
    mfem::Vector legacyAction(prepared.Height());
    mfem::Vector preparedAction(prepared.Height());
    legacyAction                           = 0.0;
    preparedAction                         = 0.0;
    double *const preparedStorage          = preparedAction.GetData();

    const std::uint64_t massPreparations   = geometry.context.GetMassOperator().GetPreparationCount();
    const std::uint64_t sourcePreparations = geometry.context.GetSourceOperator().GetPreparationCount();
    legacy.Mult(rightHandSide, legacyAction);
    prepared.Mult(rightHandSide, preparedAction);

    CHECK(preparedAction.GetData() == preparedStorage);
    CHECK(geometry.context.GetMassOperator().GetPreparationCount() == massPreparations);
    CHECK(geometry.context.GetSourceOperator().GetPreparationCount() == sourcePreparations);
    CHECK(
        gravity_prepared_test_utils::relative_error(
            preparedAction, legacyAction, geometry.finiteElements.mesh->GetComm()
        ) <= 2.0e-12
    );
}

TEST_CASE(
    "Prepared Gravity Blocks Refresh Explicitly Without Repreparing Geometry",
    tags::preconditioning_gravity_integration
) {
    const auto arguments = test_utils::setup_args();
    PreparedGeometry geometry(arguments);

    const auto block = preconditioning::GravityFieldBlock(
        backend::Diagonal{}, FixedAMG{backend::FixedCycles{.cycles = 1}}, preconditioning::GravityApproximateLDU{}
    );
    auto prepared             = preconditioning::prepare(geometry.finiteElements, geometry.context, block);
    const auto chebyshevBlock = preconditioning::GravityFieldBlock(
        backend::MatrixFreeChebyshev{.order = 2, .powerIterations = 10}, FixedAMG{backend::FixedCycles{.cycles = 1}},
        preconditioning::GravityApproximateLDU{}
    );
    auto chebyshevPrepared        = preconditioning::prepare(geometry.finiteElements, geometry.context, chebyshevBlock);
    CHECK(prepared.GetDivergenceAssemblyCount() == 1);
    CHECK(chebyshevPrepared.GetDivergenceAssemblyCount() == 1);

    const auto unchanged          = prepared.Refresh(geometry.finiteElements, geometry.context);
    const auto unchangedChebyshev = chebyshevPrepared.Refresh(geometry.finiteElements, geometry.context);
    CHECK_FALSE(unchanged.DidAnyWork());
    CHECK_FALSE(unchangedChebyshev.DidAnyWork());
    CHECK(prepared.GetStatistics().refreshChecks == 1);
    CHECK(prepared.GetStatistics().noOpRefreshes == 1);

    const mfem::Vector displacementTrue = gravity_prepared_test_utils::make_displacement(geometry.finiteElements, 0.4);
    const mfem::Vector displacement     = geometry.context.GetDisplacementMap().gather(displacementTrue);
    geometry.context.PreparePrimal(displacement, {.value = 1}, {.value = 2});
    CHECK_FALSE(prepared.IsCurrent());
    CHECK_FALSE(chebyshevPrepared.IsCurrent());

    mfem::Vector rightHandSide(prepared.Width());
    mfem::Vector action(prepared.Height());
    rightHandSide = 1.0;
    action        = 0.0;
    CHECK_THROWS_AS(prepared.Mult(rightHandSide, action), std::logic_error);
    CHECK_THROWS_AS(chebyshevPrepared.Mult(rightHandSide, action), std::logic_error);

    const std::uint64_t massPreparations   = geometry.context.GetMassOperator().GetPreparationCount();
    const std::uint64_t sourcePreparations = geometry.context.GetSourceOperator().GetPreparationCount();
    const auto changed                     = prepared.Refresh(geometry.finiteElements, geometry.context);
    const auto changedChebyshev            = chebyshevPrepared.Refresh(geometry.finiteElements, geometry.context);
    CHECK(prepared.GetDivergenceAssemblyCount() == 1);
    CHECK(chebyshevPrepared.GetDivergenceAssemblyCount() == 1);

    CHECK(changed.geometryChanged);
    CHECK_FALSE(changed.discretizationChanged);
    CHECK(changed.rebuiltMassInverse);
    CHECK_FALSE(changed.rebuiltDivergenceBinding);
    CHECK(changed.rebuiltPotentialSchur);
    CHECK(prepared.IsCurrent());
    CHECK(changedChebyshev.geometryChanged);
    CHECK(changedChebyshev.rebuiltMassInverse);
    CHECK(changedChebyshev.rebuiltPotentialSchur);
    CHECK(chebyshevPrepared.IsCurrent());
    CHECK(chebyshevPrepared.GetMassInverse().GetStatistics().setups == 2);
    CHECK(geometry.context.GetMassOperator().GetPreparationCount() == massPreparations);
    CHECK(geometry.context.GetSourceOperator().GetPreparationCount() == sourcePreparations);
    CHECK(prepared.GetStatistics().refreshes == 1);

    mfem::Vector refreshedAction(chebyshevPrepared.Height());
    refreshedAction = 0.0;
    chebyshevPrepared.Mult(rightHandSide, refreshedAction);
    for (int index = 0; index < refreshedAction.Size(); ++index) {
        CHECK(std::isfinite(refreshedAction(index)));
    }

    // A discretization revision reconstructs the matrix-free mass operator.  The
    // owning gravity block must reject every route to its now-stale inverse until
    // refresh has rebound and rebuilt the Chebyshev smoother.
    geometry.context.PreparePrimal(displacement, {.value = 2}, {.value = 2});
    CHECK_FALSE(chebyshevPrepared.IsCurrent());
    CHECK_THROWS_AS(chebyshevPrepared.Mult(rightHandSide, action), std::logic_error);
    CHECK_THROWS_AS(chebyshevPrepared.GetMassInverse(), std::logic_error);

    const auto reconstructed = chebyshevPrepared.Refresh(geometry.finiteElements, geometry.context);
    CHECK(reconstructed.discretizationChanged);
    CHECK_FALSE(reconstructed.geometryChanged);
    CHECK(reconstructed.rebuiltMassInverse);
    CHECK(reconstructed.rebuiltDivergenceBinding);
    CHECK(reconstructed.rebuiltPotentialSchur);
    CHECK(chebyshevPrepared.IsCurrent());
    CHECK(chebyshevPrepared.GetMassInverse().GetStatistics().setups == 3);
    CHECK(chebyshevPrepared.GetDivergenceAssemblyCount() == 2);

    mfem::Vector firstReconstructedAction(chebyshevPrepared.Height());
    mfem::Vector secondReconstructedAction(chebyshevPrepared.Height());
    firstReconstructedAction  = 0.0;
    secondReconstructedAction = 0.0;
    chebyshevPrepared.Mult(rightHandSide, firstReconstructedAction);
    chebyshevPrepared.Mult(rightHandSide, secondReconstructedAction);
    mfem::Vector repeatabilityError(firstReconstructedAction);
    repeatabilityError -= secondReconstructedAction;
    CHECK(repeatabilityError.Norml2() <= 2.0e-14 * std::max(1.0, firstReconstructedAction.Norml2()));
    for (int index = 0; index < firstReconstructedAction.Size(); ++index) {
        CHECK(std::isfinite(firstReconstructedAction(index)));
    }
}

TEST_CASE(
    "Prepared Gravity Schur Refinement Reuses Sparse Couplings Across Geometry Revisions",
    tags::preconditioning_gravity_integration
) {
    const auto arguments = test_utils::setup_args();
    PreparedGeometry geometry(arguments);
    const auto makeBlock = [](const int sweeps) {
        return preconditioning::GravityFieldBlock(
            backend::MatrixFreeChebyshev{.order = 2, .powerIterations = 10},
            FixedAMG{backend::FixedCycles{.cycles = 1}},
            preconditioning::GravityConsistentSchurLDU{.correctionSweeps = sweeps, .relaxation = 0.5}
        );
    };
    auto baseline = preconditioning::prepare(geometry.finiteElements, geometry.context, makeBlock(0));
    auto refined = preconditioning::prepare(geometry.finiteElements, geometry.context, makeBlock(2));
    const MPI_Comm communicator = geometry.finiteElements.mesh->GetComm();
    const auto gradientMap = geometry.context.GetMassOperator().GetFluxMap();
    const auto potentialMap = geometry.context.GetSourceOperator().GetPotentialMap();
    const mfem::Vector rhs = gravity_prepared_test_utils::make_deterministic_vector(refined.Width(), 0.29);
    mfem::Vector actual(refined.Height()), expected(refined.Height());

    const auto compareApplications = [&] {
        const auto sparse = preconditioning::assembleGravityDivergenceSurrogate(geometry.finiteElements);
        preconditioning::ReducedGravityDivergenceOperator sparseReduced(*sparse, gradientMap, potentialMap);
        preconditioning::ReducedGravityDivergenceOperator matrixFreeReduced(
            geometry.context.GetDivergenceOperator(), gradientMap, potentialMap
        );
        for (const double phase : {0.19, 0.53, 0.87}) {
            const auto flux = gravity_prepared_test_utils::make_deterministic_vector(sparseReduced.Width(), phase);
            const auto potential =
                gravity_prepared_test_utils::make_deterministic_vector(sparseReduced.Height(), phase + 0.17);
            mfem::Vector sparseAction(sparseReduced.Height()), matrixFreeAction(matrixFreeReduced.Height());
            mfem::Vector sparseTranspose(sparseReduced.Width()), matrixFreeTranspose(matrixFreeReduced.Width());
            sparseReduced.Mult(flux, sparseAction);
            matrixFreeReduced.Mult(flux, matrixFreeAction);
            sparseReduced.MultTranspose(potential, sparseTranspose);
            matrixFreeReduced.MultTranspose(potential, matrixFreeTranspose);
            CHECK(gravity_prepared_test_utils::relative_error(sparseAction, matrixFreeAction, communicator) < 2.0e-12);
            CHECK(gravity_prepared_test_utils::relative_error(sparseTranspose, matrixFreeTranspose, communicator) <
                  2.0e-12);
        }

        // Reuse precisely the same child solvers, so zero-sweep equivalence
        // checks the path itself rather than independent AMG setup roundoff.
        preconditioning::ReducedFieldSolverAdapter baselineSchur(baseline.GetPotentialSchurInverse(), potentialMap);
        preconditioning::GravityFactorizationOperator baselineReference(
            preconditioning::GravityApproximateLDU{}, baseline.GetMassInverse(), baselineSchur, matrixFreeReduced
        );
        baseline.Mult(rhs, actual);
        baselineReference.Mult(rhs, expected);
        bool baselineIdentical = true;
        int firstDifferentCoefficient = -1;
        double maximumDifference = 0.0;
        for (int index = 0; index < actual.Size(); ++index) {
            const bool identical =
                std::bit_cast<std::uint64_t>(actual(index)) == std::bit_cast<std::uint64_t>(expected(index));
            if (!identical && firstDifferentCoefficient < 0) firstDifferentCoefficient = index;
            baselineIdentical = baselineIdentical && identical;
            maximumDifference = std::max(maximumDifference, std::abs(actual(index) - expected(index)));
        }
        CAPTURE(firstDifferentCoefficient, maximumDifference);
        CHECK(baselineIdentical);

        preconditioning::ReducedFieldSolverAdapter refinedSchur(refined.GetPotentialSchurInverse(), potentialMap);
        preconditioning::GravityFactorizationOperator refinedReference(
            makeBlock(2).factorizationPolicy(), refined.GetMassInverse(), refinedSchur, matrixFreeReduced
        );
        refined.Mult(rhs, actual);
        refinedReference.Mult(rhs, expected);
        CHECK(gravity_prepared_test_utils::relative_error(actual, expected, communicator) < 2.0e-10);
    };

    compareApplications();
    CHECK(refined.GetDivergenceAssemblyCount() == 1);
    CHECK_FALSE(refined.Refresh(geometry.finiteElements, geometry.context).DidAnyWork());
    const mfem::Vector displacementTrue = gravity_prepared_test_utils::make_displacement(geometry.finiteElements, 0.4);
    const mfem::Vector displacement = geometry.context.GetDisplacementMap().gather(displacementTrue);
    geometry.context.PreparePrimal(displacement, {.value = 1}, {.value = 2});
    CHECK_THROWS_AS(refined.Mult(rhs, actual), std::logic_error);
    const auto geometryRefresh = refined.Refresh(geometry.finiteElements, geometry.context);
    (void) baseline.Refresh(geometry.finiteElements, geometry.context);
    CHECK(geometryRefresh.geometryChanged);
    CHECK_FALSE(geometryRefresh.rebuiltDivergenceBinding);
    CHECK(refined.GetDivergenceAssemblyCount() == 1);
    compareApplications();

    geometry.context.PreparePrimal(displacement, {.value = 2}, {.value = 2});
    CHECK_THROWS_AS(refined.Mult(rhs, actual), std::logic_error);
    const auto discretizationRefresh = refined.Refresh(geometry.finiteElements, geometry.context);
    (void) baseline.Refresh(geometry.finiteElements, geometry.context);
    CHECK(discretizationRefresh.rebuiltDivergenceBinding);
    CHECK(refined.GetDivergenceAssemblyCount() == 2);
    compareApplications();

    int rank = 0, rankCount = 0;
    MPI_Comm_rank(communicator, &rank);
    MPI_Comm_size(communicator, &rankCount);
    if (rankCount > 1) {
        const auto checkCollectiveRejection = [&](const preconditioning::GravityConsistentSchurLDU policy) {
            const auto block = preconditioning::GravityFieldBlock(
                backend::Diagonal{}, FixedAMG{backend::FixedCycles{.cycles = 1}}, policy
            );
            CHECK_THROWS_AS(
                preconditioning::prepare(geometry.finiteElements, geometry.context, block), std::invalid_argument
            );
        };
        checkCollectiveRejection({.correctionSweeps = rank == 0 ? -1 : 1, .relaxation = 0.5});
        checkCollectiveRejection({.correctionSweeps = rank == 0 ? 0 : 1, .relaxation = 0.5});
        checkCollectiveRejection({.correctionSweeps = 1, .relaxation = rank == 0 ? 0.5 : 1.0});
    }
}

TEST_CASE(
    "Fixed Gravity Divergence Cache Preserves Changing Mass Weights And Invalidates",
    tags::preconditioning_gravity_integration
) {
    const auto args = test_utils::setup_args();
    auto f = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());
    preconditioning::PreparedGravityDivergenceCache cache;
    mfem::Vector diagonal(f.gravityFluxFes->GetTrueVSize()); diagonal = 1.0;
    const auto x = gravity_prepared_test_utils::make_deterministic_vector(f.gravityPotentialFes->GetTrueVSize(), 0.37);
    auto initial = preconditioning::assembleGravityPotentialSchurSurrogate(f, diagonal, &cache);
    mfem::Vector initialAction(initial->Height()); initial->Mult(x, initialAction);
    REQUIRE(initialAction.Norml2() > 0.0);
    REQUIRE(cache.GetBuildCount() == 1);
    for (double scale : {2.0, 0.5, 1.0}) {
        CAPTURE(scale);
        diagonal = scale;
        auto fast = preconditioning::assembleGravityPotentialSchurSurrogate(f, diagonal, &cache);
        auto reference = preconditioning::assembleGravityPotentialSchurSurrogate(f, diagonal);
        mfem::Vector a(fast->Height()), b(reference->Height()), expected(initialAction);
        fast->Mult(x, a); reference->Mult(x, b); expected /= scale;
        CHECK(gravity_prepared_test_utils::relative_error(a, b, f.mesh->GetComm()) < 2.0e-12);
        CHECK(gravity_prepared_test_utils::relative_error(a, expected, f.mesh->GetComm()) < 2.0e-12);
        CHECK(cache.GetBuildCount() == 1);
        initial->Mult(x, a);
        CHECK(gravity_prepared_test_utils::relative_error(a, initialAction, f.mesh->GetComm()) < 2.0e-14);
    }
    for (double bad : {0.0, -1.0, std::numeric_limits<double>::quiet_NaN()}) {
        diagonal = bad;
        CHECK_THROWS_AS(
            preconditioning::assembleGravityPotentialSchurSurrogate(f, diagonal, &cache), std::invalid_argument);
        CHECK(cache.GetBuildCount() == 1);
    }
    diagonal = 1.0;
    const auto checkRebuild = [&] {
        auto fast = preconditioning::assembleGravityPotentialSchurSurrogate(f, diagonal, &cache);
        auto reference = preconditioning::assembleGravityPotentialSchurSurrogate(f, diagonal);
        mfem::Vector a(fast->Height()), b(reference->Height());
        fast->Mult(x, a); reference->Mult(x, b);
        CHECK(gravity_prepared_test_utils::relative_error(a, b, f.mesh->GetComm()) < 2.0e-12);
    };
    cache.Clear(); checkRebuild();
    CHECK(cache.GetBuildCount() == 2);
    f.mesh->NodesUpdated(); checkRebuild();
    CHECK(cache.GetBuildCount() == 3);
    checkRebuild();
    CHECK(cache.GetBuildCount() == 3);
}