#include <algorithm>
#include <chrono>
#include <cmath>
#include <iostream>
#include <limits>
#include <map>
#include <string>
#include <utility>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>
#include <mpi.h>

import experiment;
import mean_field;
import test_helpers;

namespace {
    using Clock               = std::chrono::steady_clock;

    namespace backend         = mean_field::preconditioning::backend;
    namespace preconditioning = mean_field::preconditioning;

    [[nodiscard]] const char *buildConfiguration() noexcept {
#ifdef NDEBUG
        return "release";
#else
        return "debug";
#endif
    }

    [[nodiscard]] double maximumRankSeconds(
        const Clock::time_point start,
        const MPI_Comm communicator
    ) {
        const double localSeconds = std::chrono::duration<double>(Clock::now() - start).count();
        double maximumSeconds     = 0.0;
        MPI_Allreduce(&localSeconds, &maximumSeconds, 1, MPI_DOUBLE, MPI_MAX, communicator);
        return maximumSeconds;
    }

    [[nodiscard]] double globalNorm(
        const mfem::Vector &vector,
        const MPI_Comm communicator
    ) {
        const double localSquared = vector * vector;
        double globalSquared      = 0.0;
        MPI_Allreduce(&localSquared, &globalSquared, 1, MPI_DOUBLE, MPI_SUM, communicator);
        return std::sqrt(std::max(globalSquared, 0.0));
    }

    [[nodiscard]] double globalDot(
        const mfem::Vector &left,
        const mfem::Vector &right,
        const MPI_Comm communicator
    ) {
        const double localDot = left * right;
        double result         = 0.0;
        MPI_Allreduce(&localDot, &result, 1, MPI_DOUBLE, MPI_SUM, communicator);
        return result;
    }

    void announce(
        const MPI_Comm communicator,
        const std::string &message
    ) {
        int rank = 0;
        MPI_Comm_rank(communicator, &rank);
        if (rank == 0) {
            std::cout << message << std::endl;
        }
    }

    class ReducedGravityOperator final : public mfem::Operator {
    public:
        explicit ReducedGravityOperator(
            const mean_field::operators::context::gravity_field::GravityFieldGeometryContext &context
        )
            : mfem::Operator(
                  context.GetMassOperator().GetFluxMap().reduced_size() +
                  context.GetSourceOperator().GetPotentialMap().reduced_size()
              ),
              m_mass(&context.GetMassOperator()),
              m_divergence(
                  context.GetDivergenceOperator(),
                  context.GetMassOperator().GetFluxMap(),
                  context.GetSourceOperator().GetPotentialMap()
              ),
              m_offsets(3),
              m_gradientWorkspace(context.GetMassOperator().GetFluxMap().reduced_size()) {
            m_offsets[0] = 0;
            m_offsets[1] = context.GetMassOperator().GetFluxMap().reduced_size();
            m_offsets[2] = Height();
        }

        void Mult(
            const mfem::Vector &state,
            mfem::Vector &residual
        ) const override {
            if (state.Size() != Width() || residual.Size() != Height()) {
                throw std::invalid_argument("The reduced gravity experiment requires preallocated compatible vectors.");
            }

            const mfem::Vector gradient(
                const_cast<mfem::real_t *>(state.GetData()) + m_offsets[0], m_offsets[1] - m_offsets[0]
            );
            const mfem::Vector potential(
                const_cast<mfem::real_t *>(state.GetData()) + m_offsets[1], m_offsets[2] - m_offsets[1]
            );
            mfem::Vector gradientResidual(residual.GetData() + m_offsets[0], m_offsets[1] - m_offsets[0]);
            mfem::Vector potentialResidual(residual.GetData() + m_offsets[1], m_offsets[2] - m_offsets[1]);

            m_mass->Mult(gradient, gradientResidual);
            m_divergence.MultTranspose(potential, m_gradientWorkspace);
            gradientResidual += m_gradientWorkspace;
            m_divergence.Mult(gradient, potentialResidual);
        }

    private:
        const mfem::Operator *m_mass;
        preconditioning::ReducedGravityDivergenceOperator m_divergence;
        mfem::Array<int> m_offsets;
        mutable mfem::Vector m_gradientWorkspace;
    };

    [[nodiscard]] std::map<
        std::string,
        std::string>
    commonParameters(
        const std::string &candidate,
        const std::string &measurement,
        const int dimension
    ) {
        return {
            {"build_configuration", buildConfiguration()},
            {"candidate", candidate},
            {"experiment_schema", "p4_reduced_gravity_v1"},
            {"factorization", candidate},
            {"measurement", measurement},
            {"mesh_file", test_utils::setup_args().mesh_file},
            {"operator", "reduced_gravity_saddle_point"},
            {"preconditioned_product", "G M^-1"},
            {"root_dimension", std::to_string(dimension)}
        };
    }

    void recordSpectrum(
        const std::string &candidate,
        const mean_field::solver::ArnoldiSpectralMeasurement &spectrum,
        const int dimension,
        const double setupSeconds
    ) {
        experiment::record_experiment_result(
            "gravity_preconditioning_p4", candidate + "_spectrum",
            commonParameters(candidate, "arnoldi_summary", dimension),
            {{"setup_seconds_maximum_rank", setupSeconds},
             {"requested_dimension", static_cast<double>(spectrum.requestedDimension)},
             {"achieved_dimension", static_cast<double>(spectrum.achievedDimension)},
             {"operator_applications", static_cast<double>(spectrum.operatorApplications)},
             {"measurement_seconds_maximum_rank", spectrum.measurementSecondsMaximumRank},
             {"operator_application_seconds_maximum_rank", spectrum.operatorApplicationSecondsMaximumRank},
             {"projected_condition_proxy", spectrum.projectedConditionProxy},
             {"projected_largest_singular_value", spectrum.projectedLargestSingularValue},
             {"projected_smallest_singular_value", spectrum.projectedSmallestSingularValue},
             {"centroid_real_part", spectrum.centroidRealPart},
             {"rms_distance_from_one", spectrum.rmsDistanceFromOne},
             {"rms_cluster_radius", spectrum.rmsClusterRadius},
             {"minimum_magnitude", spectrum.minimumMagnitude},
             {"maximum_magnitude", spectrum.maximumMagnitude},
             {"minimum_real_part", spectrum.minimumRealPart},
             {"maximum_real_part", spectrum.maximumRealPart},
             {"maximum_absolute_imaginary_part", spectrum.maximumAbsoluteImaginaryPart},
             {"negative_real_part_count", static_cast<double>(spectrum.negativeRealPartCount)},
             {"converged_ritz_value_count", static_cast<double>(spectrum.convergedRitzValueCount)},
             {"conjugate_pair_defect", spectrum.conjugatePairDefect},
             {"projected_departure_from_normality", spectrum.projectedDepartureFromNormality},
             {"field_of_values_minimum_real_part", spectrum.projectedFieldOfValuesMinimumRealPart},
             {"field_of_values_maximum_real_part", spectrum.projectedFieldOfValuesMaximumRealPart}}
        );

        for (std::size_t index = 0; index < spectrum.ritzValues.size(); ++index) {
            const auto &value = spectrum.ritzValues[index];
            experiment::record_experiment_result(
                "gravity_preconditioning_p4", candidate + "_ritz_" + std::to_string(index),
                commonParameters(candidate, "ritz_value", dimension),
                {{"ritz_index", static_cast<double>(index)},
                 {"real_part", value.realPart},
                 {"imaginary_part", value.imaginaryPart},
                 {"magnitude", value.magnitude},
                 {"distance_from_one", value.distanceFromOne},
                 {"residual_estimate", value.residualEstimate},
                 {"relative_residual_estimate", value.relativeResidualEstimate},
                 {"converged", value.converged ? 1.0 : 0.0}}
            );
        }
    }

    void measureCandidate(
        const std::string &candidate,
        mfem::Solver &inversePreconditioner,
        const double setupSeconds,
        const ReducedGravityOperator &gravityOperator,
        const mfem::Vector &rightHandSide,
        const mfem::Vector &arnoldiDirection,
        const MPI_Comm communicator
    ) {
        constexpr int arnoldiDimension = 32;

        mean_field::solver::InstrumentedOperator instrumentedGravity(gravityOperator);
        mean_field::solver::InstrumentedPreconditioner instrumentedPreconditioner(inversePreconditioner);
        mean_field::solver::ResidualHistoryMonitor monitor;
        mfem::FGMRESSolver krylov(communicator);
        krylov.SetPreconditioner(instrumentedPreconditioner);
        krylov.SetOperator(instrumentedGravity);
        krylov.SetMonitor(monitor);
        krylov.SetRelTol(1.0e-8);
        krylov.SetAbsTol(1.0e-12);
        krylov.SetMaxIter(100);
        krylov.SetKDim(30);
        krylov.SetPrintLevel(0);

        mfem::Vector solution(gravityOperator.Width());
        solution = 0.0;
        announce(communicator, "P4 reduced gravity: solving with " + candidate);
        const Clock::time_point solveStart = Clock::now();
        krylov.Mult(rightHandSide, solution);
        const double solveSeconds = maximumRankSeconds(solveStart, communicator);

        mfem::Vector reconstructed(rightHandSide.Size());
        gravityOperator.Mult(solution, reconstructed);
        reconstructed -= rightHandSide;
        const double relativeResidual =
            globalNorm(reconstructed, communicator) /
            std::max(globalNorm(rightHandSide, communicator), std::numeric_limits<double>::epsilon());

        const auto jacobianStatistics       = instrumentedGravity.GetStatistics();
        const auto preconditionerStatistics = instrumentedPreconditioner.GetStatistics();
        REQUIRE(std::isfinite(relativeResidual));
        experiment::record_experiment_result(
            "gravity_preconditioning_p4", candidate + "_linear_solve",
            commonParameters(candidate, "linear_solve", gravityOperator.Width()),
            {{"setup_seconds_maximum_rank", setupSeconds},
             {"solver_converged", krylov.GetConverged() ? 1.0 : 0.0},
             {"outer_iterations", static_cast<double>(krylov.GetNumIterations())},
             {"true_relative_residual", relativeResidual},
             {"solve_seconds_maximum_rank", solveSeconds},
             {"gravity_applications", static_cast<double>(jacobianStatistics.applications)},
             {"gravity_application_seconds", jacobianStatistics.totalSeconds},
             {"preconditioner_applications", static_cast<double>(preconditionerStatistics.applications)},
             {"preconditioner_application_seconds", preconditionerStatistics.totalSeconds},
             {"preconditioner_maximum_application_seconds", preconditionerStatistics.maximumSeconds}}
        );

        instrumentedGravity.ResetStatistics();
        instrumentedPreconditioner.ResetStatistics();
        mean_field::solver::FixedRightPreconditionedOperator product(instrumentedGravity, instrumentedPreconditioner);
        announce(communicator, "P4 reduced gravity: measuring " + candidate + " Arnoldi spectrum");
        const auto spectrum = mean_field::solver::measureArnoldiSpectrum(
            product, arnoldiDirection, communicator,
            {.krylovDimension                  = arnoldiDimension,
             .breakdownRelativeTolerance       = 1.0e-13,
             .ritzConvergenceRelativeTolerance = 1.0e-7,
             .reorthogonalize                  = true}
        );
        recordSpectrum(candidate, spectrum, gravityOperator.Width(), setupSeconds);
    }

    template <
        preconditioning::GravityFactorizationPolicy Policy,
        backend::Registered MassBackend = backend::Diagonal>
        requires backend::Compatible<
            MassBackend,
            preconditioning::GravityMassInverseCharacteristics>
    void prepareAndMeasureTypedCandidate(
        const std::string &candidate,
        Policy policy,
        const mean_field::fem::FEM &finiteElements,
        const mean_field::operators::context::gravity_field::GravityFieldGeometryContext &geometryContext,
        const ReducedGravityOperator &gravityOperator,
        const mfem::Vector &rightHandSide,
        const mfem::Vector &arnoldiDirection,
        const MPI_Comm communicator,
        const int amgCycles     = 1,
        MassBackend massBackend = {}
    ) {
        const Clock::time_point setupStart = Clock::now();
        const auto block                   = preconditioning::GravityFieldBlock(
            std::move(massBackend), backend::HypreBoomerAMG{backend::FixedCycles{.cycles = amgCycles}}, policy
        );
        auto prepared            = preconditioning::prepare(finiteElements, geometryContext, block);
        const double setupTime   = maximumRankSeconds(setupStart, communicator);

        const auto &massOperator = geometryContext.GetMassOperator();
        const mfem::Vector firstMassRightHandSide =
            gravity_prepared_test_utils::make_deterministic_vector(massOperator.Width(), 0.41);
        const mfem::Vector secondMassRightHandSide =
            gravity_prepared_test_utils::make_deterministic_vector(massOperator.Width(), 1.17);
        mfem::Vector firstMassAction(massOperator.Width());
        mfem::Vector secondMassAction(massOperator.Width());
        prepared.GetMassInverse().Mult(firstMassRightHandSide, firstMassAction);
        prepared.GetMassInverse().Mult(secondMassRightHandSide, secondMassAction);
        mfem::Vector recoveredMassRightHandSide(massOperator.Height());
        massOperator.Mult(firstMassAction, recoveredMassRightHandSide);
        recoveredMassRightHandSide -= firstMassRightHandSide;
        const double massRecoveryDefect =
            globalNorm(recoveredMassRightHandSide, communicator) / globalNorm(firstMassRightHandSide, communicator);
        const double firstSecond = globalDot(firstMassRightHandSide, secondMassAction, communicator);
        const double secondFirst = globalDot(secondMassRightHandSide, firstMassAction, communicator);
        const double massSymmetryDefect =
            std::abs(firstSecond - secondFirst) / std::max({1.0, std::abs(firstSecond), std::abs(secondFirst)});
        const double massPositiveRayleigh = globalDot(firstMassRightHandSide, firstMassAction, communicator) /
                                            std::max(
                                                globalDot(firstMassRightHandSide, firstMassRightHandSide, communicator),
                                                std::numeric_limits<double>::min()
                                            );

        const auto &schurOperator         = prepared.GetPotentialSchurSurrogate();
        const mfem::Vector schurRightHandSide =
            gravity_prepared_test_utils::make_deterministic_vector(schurOperator.Width(), 0.73);
        mfem::Vector schurAction(schurOperator.Width());
        prepared.GetPotentialSchurInverse().Mult(schurRightHandSide, schurAction);
        mfem::Vector recoveredSchurRightHandSide(schurOperator.Height());
        schurOperator.Mult(schurAction, recoveredSchurRightHandSide);
        recoveredSchurRightHandSide -= schurRightHandSide;
        const double schurRecoveryDefect =
            globalNorm(recoveredSchurRightHandSide, communicator) / globalNorm(schurRightHandSide, communicator);

        experiment::record_experiment_result(
            "gravity_preconditioning_p4", candidate + "_block_quality",
            commonParameters(candidate, "block_inverse_quality", gravityOperator.Width()),
            {{"amg_cycles", static_cast<double>(amgCycles)},
             {"mass_inverse_recovery_defect", massRecoveryDefect},
             {"mass_inverse_symmetry_defect", massSymmetryDefect},
             {"mass_inverse_positive_rayleigh", massPositiveRayleigh},
             {"potential_schur_inverse_recovery_defect", schurRecoveryDefect}}
        );
        measureCandidate(
            candidate, prepared, setupTime, gravityOperator, rightHandSide, arnoldiDirection, communicator
        );
    }

    [[nodiscard]] int firstReportedThresholdIteration(
        const std::vector<mean_field::solver::IterationResidualMeasurement> &history,
        const double initialNorm,
        const double relativeThreshold
    ) {
        if (!std::isfinite(initialNorm) || initialNorm <= 0.0) {
            return -1;
        }
        for (const auto &sample : history) {
            if (std::abs(sample.reportedNorm) / initialNorm <= relativeThreshold) {
                return sample.iteration;
            }
        }
        return -1;
    }
} // namespace

TEST_CASE(
    "Reduced Gravity P4 Factorization Comparison",
    "[preconditioning][gravity][diagnostics][experiment][spectrum]"
) {
    const auto arguments                = test_utils::setup_args();
    mean_field::fem::FEM finiteElements = mean_field::fem::setup_fem(arguments.mesh_file, arguments, 0);
    const MPI_Comm communicator         = finiteElements.mesh->GetComm();
    using GeometryContext               = mean_field::operators::context::gravity_field::GravityFieldGeometryContext;
    GeometryContext geometryContext(finiteElements, *finiteElements.domainMapperStateless);

    mfem::Vector displacementTrue(finiteElements.displacementFes->GetTrueVSize());
    displacementTrue                = 0.0;
    const mfem::Vector displacement = geometryContext.GetDisplacementMap().gather(displacementTrue);
    geometryContext.PreparePrimal(displacement, {.value = 1}, {.value = 1});

    ReducedGravityOperator gravityOperator(geometryContext);
    const mfem::Vector exact = gravity_prepared_test_utils::make_deterministic_vector(gravityOperator.Width(), 0.37);
    mfem::Vector rightHandSide(gravityOperator.Height());
    gravityOperator.Mult(exact, rightHandSide);
    const mfem::Vector arnoldiDirection =
        gravity_prepared_test_utils::make_deterministic_vector(gravityOperator.Width(), 0.83);

    const Clock::time_point legacySetupStart = Clock::now();
    mean_field::operators::ReducedGravityFieldPreconditioner legacy(finiteElements, geometryContext);
    const double legacySetupTime = maximumRankSeconds(legacySetupStart, communicator);
    measureCandidate(
        "legacy_block_diagonal", legacy, legacySetupTime, gravityOperator, rightHandSide, arnoldiDirection, communicator
    );

    prepareAndMeasureTypedCandidate(
        "typed_block_diagonal", preconditioning::GravityBlockDiagonal{}, finiteElements, geometryContext,
        gravityOperator, rightHandSide, arnoldiDirection, communicator
    );
    prepareAndMeasureTypedCandidate(
        "lower_triangular", preconditioning::GravityLowerTriangular{}, finiteElements, geometryContext, gravityOperator,
        rightHandSide, arnoldiDirection, communicator
    );
    prepareAndMeasureTypedCandidate(
        "upper_triangular", preconditioning::GravityUpperTriangular{}, finiteElements, geometryContext, gravityOperator,
        rightHandSide, arnoldiDirection, communicator
    );
    prepareAndMeasureTypedCandidate(
        "approximate_ldu", preconditioning::GravityApproximateLDU{}, finiteElements, geometryContext, gravityOperator,
        rightHandSide, arnoldiDirection, communicator
    );
}

TEST_CASE(
    "Reduced Gravity P4 Fixed AMG Cycle Sweep",
    "[preconditioning][gravity][diagnostics][experiment][amg_cycle_sweep]"
) {
    const auto arguments                = test_utils::setup_args();
    mean_field::fem::FEM finiteElements = mean_field::fem::setup_fem(arguments.mesh_file, arguments, 0);
    const MPI_Comm communicator         = finiteElements.mesh->GetComm();
    using GeometryContext               = mean_field::operators::context::gravity_field::GravityFieldGeometryContext;
    GeometryContext geometryContext(finiteElements, *finiteElements.domainMapperStateless);

    mfem::Vector displacementTrue(finiteElements.displacementFes->GetTrueVSize());
    displacementTrue                = 0.0;
    const mfem::Vector displacement = geometryContext.GetDisplacementMap().gather(displacementTrue);
    geometryContext.PreparePrimal(displacement, {.value = 1}, {.value = 1});

    ReducedGravityOperator gravityOperator(geometryContext);
    const mfem::Vector exact = gravity_prepared_test_utils::make_deterministic_vector(gravityOperator.Width(), 0.37);
    mfem::Vector rightHandSide(gravityOperator.Height());
    gravityOperator.Mult(exact, rightHandSide);
    const mfem::Vector arnoldiDirection =
        gravity_prepared_test_utils::make_deterministic_vector(gravityOperator.Width(), 0.83);

    for (const int cycles : {1, 2, 3, 4, 6, 8}) {
        prepareAndMeasureTypedCandidate(
            "approximate_ldu_amg_cycles_" + std::to_string(cycles), preconditioning::GravityApproximateLDU{},
            finiteElements, geometryContext, gravityOperator, rightHandSide, arnoldiDirection, communicator, cycles
        );
    }
    for (const int order : {2, 3, 4, 5}) {
        for (const int cycles : {1, 2, 3}) {
            prepareAndMeasureTypedCandidate(
                "approximate_ldu_chebyshev_" + std::to_string(order) + "_amg_cycles_" + std::to_string(cycles),
                preconditioning::GravityApproximateLDU{}, finiteElements, geometryContext, gravityOperator,
                rightHandSide, arnoldiDirection, communicator, cycles,
                backend::MatrixFreeChebyshev{.order = order, .powerIterations = 20}
            );
        }
    }
}

TEST_CASE(
    "Reduced Gravity P4 LDU Extended FGMRES Convergence",
    "[preconditioning][gravity][diagnostics][experiment][p4_followup][extended_solve]"
) {
    constexpr int maximumIterations     = 200;
    constexpr int restartDimension      = 30;

    const auto arguments                = test_utils::setup_args();
    mean_field::fem::FEM finiteElements = mean_field::fem::setup_fem(arguments.mesh_file, arguments, 0);
    const MPI_Comm communicator         = finiteElements.mesh->GetComm();
    using GeometryContext               = mean_field::operators::context::gravity_field::GravityFieldGeometryContext;
    GeometryContext geometryContext(finiteElements, *finiteElements.domainMapperStateless);

    mfem::Vector displacementTrue(finiteElements.displacementFes->GetTrueVSize());
    displacementTrue                = 0.0;
    const mfem::Vector displacement = geometryContext.GetDisplacementMap().gather(displacementTrue);
    geometryContext.PreparePrimal(displacement, {.value = 1}, {.value = 1});

    ReducedGravityOperator gravityOperator(geometryContext);
    const mfem::Vector exact = gravity_prepared_test_utils::make_deterministic_vector(gravityOperator.Width(), 0.37);
    mfem::Vector rightHandSide(gravityOperator.Height());
    gravityOperator.Mult(exact, rightHandSide);

    const Clock::time_point setupStart = Clock::now();
    const auto block                   = preconditioning::GravityFieldBlock(
        backend::Diagonal{}, backend::HypreBoomerAMG{backend::FixedCycles{.cycles = 1}},
        preconditioning::GravityApproximateLDU{}
    );
    auto prepared          = preconditioning::prepare(finiteElements, geometryContext, block);
    const double setupTime = maximumRankSeconds(setupStart, communicator);

    mean_field::solver::InstrumentedOperator instrumentedGravity(gravityOperator);
    mean_field::solver::InstrumentedPreconditioner instrumentedPreconditioner(prepared);
    mean_field::solver::ResidualHistoryMonitor monitor;
    mfem::FGMRESSolver krylov(communicator);
    krylov.SetPreconditioner(instrumentedPreconditioner);
    krylov.SetOperator(instrumentedGravity);
    krylov.SetMonitor(monitor);
    krylov.SetRelTol(1.0e-8);
    krylov.SetAbsTol(1.0e-12);
    krylov.SetMaxIter(maximumIterations);
    krylov.SetKDim(restartDimension);
    krylov.SetPrintLevel(0);

    mfem::Vector solution(gravityOperator.Width());
    solution = 0.0;
    announce(communicator, "P4 follow-up: running 200-iteration approximate-LDU FGMRES");
    const Clock::time_point solveStart = Clock::now();
    krylov.Mult(rightHandSide, solution);
    const double solveSeconds = maximumRankSeconds(solveStart, communicator);

    mfem::Vector reconstructed(rightHandSide.Size());
    gravityOperator.Mult(solution, reconstructed);
    reconstructed -= rightHandSide;
    const double trueRelativeResidual =
        globalNorm(reconstructed, communicator) /
        std::max(globalNorm(rightHandSide, communicator), std::numeric_limits<double>::epsilon());
    const double initialNorm = std::abs(krylov.GetInitialNorm());
    const auto &history      = monitor.GetHistory();
    const int iteration1e4   = firstReportedThresholdIteration(history, initialNorm, 1.0e-4);
    const int iteration1e6   = firstReportedThresholdIteration(history, initialNorm, 1.0e-6);
    const int iteration1e8   = firstReportedThresholdIteration(history, initialNorm, 1.0e-8);

    REQUIRE(std::isfinite(trueRelativeResidual));
    REQUIRE_FALSE(history.empty());
    experiment::record_experiment_result(
        "gravity_preconditioning_p4_followup", "approximate_ldu_extended_linear_solve",
        commonParameters("approximate_ldu_extended", "linear_solve", gravityOperator.Width()),
        {{"maximum_iterations", static_cast<double>(maximumIterations)},
         {"restart_dimension", static_cast<double>(restartDimension)},
         {"setup_seconds_maximum_rank", setupTime},
         {"solver_converged", krylov.GetConverged() ? 1.0 : 0.0},
         {"outer_iterations", static_cast<double>(krylov.GetNumIterations())},
         {"reported_initial_residual_norm", initialNorm},
         {"reported_final_residual_norm", std::abs(krylov.GetFinalNorm())},
         {"reported_residual_reduction", initialNorm > 0.0 ? std::abs(krylov.GetFinalNorm()) / initialNorm : 0.0},
         {"reported_iteration_to_1e-4", static_cast<double>(iteration1e4)},
         {"reported_iteration_to_1e-6", static_cast<double>(iteration1e6)},
         {"reported_iteration_to_1e-8", static_cast<double>(iteration1e8)},
         {"true_relative_residual", trueRelativeResidual},
         {"solve_seconds_maximum_rank", solveSeconds},
         {"gravity_applications", static_cast<double>(instrumentedGravity.GetStatistics().applications)},
         {"gravity_application_seconds", instrumentedGravity.GetStatistics().totalSeconds},
         {"preconditioner_applications", static_cast<double>(instrumentedPreconditioner.GetStatistics().applications)},
         {"preconditioner_application_seconds", instrumentedPreconditioner.GetStatistics().totalSeconds}}
    );

    for (std::size_t index = 0; index < history.size(); ++index) {
        const auto &sample = history[index];
        experiment::record_experiment_result(
            "gravity_preconditioning_p4_followup", "approximate_ldu_history_" + std::to_string(index),
            commonParameters("approximate_ldu_extended", "fgmres_residual_history", gravityOperator.Width()),
            {{"history_sample", static_cast<double>(index)},
             {"iteration", static_cast<double>(sample.iteration)},
             {"reported_residual_norm", sample.reportedNorm},
             {"reported_relative_residual", initialNorm > 0.0 ? std::abs(sample.reportedNorm) / initialNorm : 0.0},
             {"final_measurement", sample.final ? 1.0 : 0.0}}
        );
    }
}

TEST_CASE(
    "Reduced Gravity P4 LDU Extended Arnoldi Convergence",
    "[preconditioning][gravity][diagnostics][experiment][spectrum][p4_followup][extended_arnoldi]"
) {
    constexpr int arnoldiDimension      = 96;

    const auto arguments                = test_utils::setup_args();
    mean_field::fem::FEM finiteElements = mean_field::fem::setup_fem(arguments.mesh_file, arguments, 0);
    const MPI_Comm communicator         = finiteElements.mesh->GetComm();
    using GeometryContext               = mean_field::operators::context::gravity_field::GravityFieldGeometryContext;
    GeometryContext geometryContext(finiteElements, *finiteElements.domainMapperStateless);

    mfem::Vector displacementTrue(finiteElements.displacementFes->GetTrueVSize());
    displacementTrue                = 0.0;
    const mfem::Vector displacement = geometryContext.GetDisplacementMap().gather(displacementTrue);
    geometryContext.PreparePrimal(displacement, {.value = 1}, {.value = 1});

    ReducedGravityOperator gravityOperator(geometryContext);
    const mfem::Vector arnoldiDirection =
        gravity_prepared_test_utils::make_deterministic_vector(gravityOperator.Width(), 0.83);

    const Clock::time_point setupStart = Clock::now();
    const auto block                   = preconditioning::GravityFieldBlock(
        backend::Diagonal{}, backend::HypreBoomerAMG{backend::FixedCycles{.cycles = 1}},
        preconditioning::GravityApproximateLDU{}
    );
    auto prepared          = preconditioning::prepare(finiteElements, geometryContext, block);
    const double setupTime = maximumRankSeconds(setupStart, communicator);

    mean_field::solver::InstrumentedOperator instrumentedGravity(gravityOperator);
    mean_field::solver::InstrumentedPreconditioner instrumentedPreconditioner(prepared);
    mean_field::solver::FixedRightPreconditionedOperator product(instrumentedGravity, instrumentedPreconditioner);
    announce(communicator, "P4 follow-up: measuring the 96-vector approximate-LDU Arnoldi spectrum");
    const auto spectrum = mean_field::solver::measureArnoldiSpectrum(
        product, arnoldiDirection, communicator,
        {.krylovDimension                  = arnoldiDimension,
         .breakdownRelativeTolerance       = 1.0e-13,
         .ritzConvergenceRelativeTolerance = 1.0e-7,
         .reorthogonalize                  = true}
    );

    REQUIRE(spectrum.achievedDimension > 32);
    recordSpectrum("approximate_ldu_arnoldi_96", spectrum, gravityOperator.Width(), setupTime);
}