#include <algorithm>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <expected>
#include <limits>
#include <memory>
#include <stdexcept>
#include <utility>
#include <vector>

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

import mean_field;

namespace coupled_trial_correction_test {
    using namespace mean_field;
    using namespace mean_field::solver;

    class DenseOperation final : public mfem::Operator {
    public:
        explicit DenseOperation(mfem::DenseMatrix matrix)
            : mfem::Operator(matrix.Height(), matrix.Width()), m_matrix(std::move(matrix)) { }

        void Mult(const mfem::Vector &input, mfem::Vector &output) const override {
            if (!prepared) throw std::logic_error("The test operator is not prepared.");
            output.SetSize(Height());
            m_matrix.Mult(input, output);
            ++applications;
        }
        void DenormalizeState(const mfem::Vector &normalized, mfem::Vector &physical) const { physical = normalized; }
        void Prepare(mfem::DenseMatrix matrix) {
            m_matrix = std::move(matrix);
            prepared = true;
            ++generation;
        }
        [[nodiscard]] bool IsPrepared() const noexcept { return prepared; }
        [[nodiscard]] std::uint64_t GetPreparationGeneration() const noexcept { return generation; }
        [[nodiscard]] const DenseOperation &GetProblem() const noexcept { return *this; }
        [[nodiscard]] const mfem::DenseMatrix &Matrix() const noexcept { return m_matrix; }

        bool prepared{true};
        mutable int applications{0};
        std::uint64_t generation{1};
    private:
        mfem::DenseMatrix m_matrix;
    };

    class ExactInverse final : public mfem::Solver {
    public:
        static constexpr auto applicationContract = preconditioning::ApplicationContract::stationary_linear;
        explicit ExactInverse(const DenseOperation &operation)
            : mfem::Solver(operation.Width()), m_operation(operation) { Refresh(); }
        void SetOperator(const mfem::Operator &) override { }
        void Mult(const mfem::Vector &input, mfem::Vector &output) const override {
            if (!IsCurrent()) throw std::logic_error("The test inverse is stale.");
            output.SetSize(Height());
            if (Height() > 0) {
                mfem::DenseMatrixInverse inverse(m_operation.Matrix());
                inverse.Mult(input, output);
                if (Height() > 1) output(1) += extraSecondComponent * input(0);
            }
            ++applications;
        }
        void Refresh() { m_generation = m_operation.GetPreparationGeneration(); }
        [[nodiscard]] bool IsCurrent() const noexcept {
            return m_operation.IsPrepared() && m_generation == m_operation.GetPreparationGeneration();
        }
        mutable int applications{0};
        double extraSecondComponent{0.0};
    private:
        const DenseOperation &m_operation;
        std::uint64_t m_generation{};
    };

    [[nodiscard]] mfem::DenseMatrix identity(const int size) {
        mfem::DenseMatrix matrix(size);
        matrix = 0.0;
        for (int i = 0; i < size; ++i) matrix(i, i) = 1.0;
        return matrix;
    }

    [[nodiscard]] mfem::Vector coupledResidual(const mfem::Vector &state) {
        mfem::Vector result(3);
        result(0) = 1.0 + state(0) + 10.0 * state(1) * state(2);
        result(1) = 1.0 + state(1);
        result(2) = 1.0 + state(2);
        return result;
    }

    [[nodiscard]] mfem::DenseMatrix coupledJacobian(const mfem::Vector &state) {
        auto matrix = identity(3);
        matrix(0, 1) = 10.0 * state(2);
        matrix(0, 2) = 10.0 * state(1);
        return matrix;
    }

    [[nodiscard]] CoupledTrialCorrectionOptions exactOptions() {
        return {.control = {.relativeTolerance = 1.e-12, .absoluteTolerance = 1.e-14, .maximumIterations = 6},
                .backend = {.restartLength = 4, .printLevel = -1}, .maximumCorrectionRatio = 4.0};
    }

    enum class BackendBehavior { exact, inaccurate, failed_status, non_finite };
    struct BackendProbe {
        BackendBehavior behavior{BackendBehavior::exact};
        int calls{0};
        bool allInitialGuessesZero{true};
    };
    struct TestBackend final : LinearBackendConfigurationTag {
        static constexpr auto supportedPreconditionerContract = preconditioning::ApplicationContract::flexible;
        std::shared_ptr<BackendProbe> probe;
        explicit TestBackend(std::shared_ptr<BackendProbe> value) : probe(std::move(value)) { }
    };

    template <typename Operation, typename Inverse>
    class PreparedTestBackend final {
    public:
        PreparedTestBackend(const Operation &operation, Inverse &inverse, const MPI_Comm communicator,
                            std::shared_ptr<BackendProbe> probe)
            : m_operation(operation), m_inverse(inverse), m_communicator(communicator), m_probe(std::move(probe)) { }
        [[nodiscard]] const Operation &GetOperator() const noexcept { return m_operation; }
        [[nodiscard]] const Inverse &GetPreconditioner() const noexcept { return m_inverse; }
        [[nodiscard]] MPI_Comm GetCommunicator() const noexcept { return m_communicator; }
        [[nodiscard]] bool IsReady() const noexcept { return m_communicator != MPI_COMM_NULL; }
        [[nodiscard]] int RightHandSideSize() const noexcept { return m_operation.Height(); }
        [[nodiscard]] int CorrectionSize() const noexcept { return m_operation.Width(); }
        [[nodiscard]] LinearSolveReport Solve(const mfem::Vector &rhs, mfem::Vector &correction,
                                               const LinearSolveControl &control) {
            ++m_probe->calls;
            for (int i = 0; i < correction.Size(); ++i)
                m_probe->allInitialGuessesZero &= correction(i) == 0.0;
            m_inverse.Mult(rhs, correction);
            if (m_probe->behavior == BackendBehavior::inaccurate) correction *= 0.5;
            if (m_probe->behavior == BackendBehavior::non_finite && correction.Size() > 0)
                correction(0) = std::numeric_limits<double>::quiet_NaN();
            double local = rhs * rhs, global = 0.0;
            if (MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_SUM, m_communicator) != MPI_SUCCESS)
                throw std::runtime_error("The test backend could not reduce its norm.");
            // A backend claiming exact convergence is deliberately insufficient
            // evidence: the correction owner must verify the actual J action.
            return {.status = m_probe->behavior == BackendBehavior::failed_status ?
                                  LinearSolveStatus::maximum_iterations : LinearSolveStatus::converged,
                    .control = control, .iterations = 1, .rightHandSideNorm = std::sqrt(global),
                    .trueResidualNorm = 0.0, .relativeTrueResidualNorm = 0.0};
        }
    private:
        const Operation &m_operation;
        Inverse &m_inverse;
        MPI_Comm m_communicator;
        std::shared_ptr<BackendProbe> m_probe;
    };

    template <typename Operation, typename Inverse>
    [[nodiscard]] auto prepareLinearBackend(TestBackend recipe, const Operation &operation, Inverse &inverse,
                                            const MPI_Comm communicator) {
        return PreparedTestBackend<Operation, Inverse>(operation, inverse, communicator, std::move(recipe.probe));
    }

    static_assert(CoupledTrialOperation<DenseOperation>);
    static_assert(CoupledTrialInverse<ExactInverse>);
    static_assert(!CoupledTrialOperation<mfem::Operator>);
    static_assert(!CoupledTrialInverse<mfem::Solver>);
    static_assert(CoupledCorrectionSearch<FullCoupledStep>);
    static_assert(CoupledCorrectionSearch<BacktrackedCoupledStep>);
    static_assert(!CoupledCorrectionSearch<int>);
    static_assert(CoupledAttemptCriterion<EveryRejectedCoupledAttempt>);
    static_assert(CoupledAttemptCriterion<NearAcceptanceCoupledAttempt>);
    static_assert(!CoupledAttemptCriterion<int>);
    static_assert(LinearBackendRuntimeAvailableFor<TestBackend, DenseOperation, ExactInverse>);
    using ScriptedCorrection = PreparedCoupledTrialCorrection<DenseOperation, ExactInverse, TestBackend>;
} // namespace coupled_trial_correction_test

TEST_CASE("A coupled trial correction preserves the tangent and removes mixed nonlinear curvature",
          "[solver][coupled_trial_correction]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    mfem::Vector base(3), direction(3), correction(3), prediction(3);
    base = 0.0;
    direction = -1.0;
    const auto baseResidual = coupledResidual(base);
    DenseOperation operation(coupledJacobian(base));
    ExactInverse inverse(operation);
    PreparedCoupledTrialCorrection prepared(operation, inverse, MPI_COMM_SELF, exactOptions());
    prepared.CaptureBase(baseResidual, direction);
    const auto capturedGeneration = prepared.GetBaseGeneration();
    CHECK(prepared.GetBaseDirectionNorm() == Approx(std::sqrt(3.0)));
    double previousScaledDisplacement = 0.0;
    for (const double alpha : {.5, .25, .125, .0625}) {
        CAPTURE(alpha);
        mfem::Vector trial(direction);
        trial *= alpha;
        const auto rawResidual = coupledResidual(trial);
        operation.Prepare(coupledJacobian(trial));
        inverse.Refresh();
        correction = 17.0;
        const auto report = prepared.SolveCorrection(alpha, rawResidual, correction);
        REQUIRE(report.Admissible());
        REQUIRE(report.verifiedGate);
        CHECK(report.linearSolve.Converged());
        CHECK(report.linearSolve.iterations <= exactOptions().control.maximumIterations);
        CHECK(report.verifiedTrueResidualNorm <= report.convergenceThreshold);
        CHECK(correction(0) == Approx(-10.0 * alpha * alpha).margin(1.e-13));
        CHECK(std::abs(correction(1)) < 1.e-13);
        CHECK(std::abs(correction(2)) < 1.e-13);
        trial += correction;
        const auto correctedResidual = coupledResidual(trial);
        prepared.BuildLinearPrediction(alpha, prediction);
        mfem::Vector remainder(correctedResidual);
        remainder -= prediction;
        CHECK(remainder.Norml2() < 1.e-12);
        CHECK(correctedResidual.Norml2() < rawResidual.Norml2());
        if (alpha == .5) {
            CHECK(rawResidual.Norml2() > baseResidual.Norml2());
            CHECK(correctedResidual.Norml2() == Approx(.5 * baseResidual.Norml2()));
        }
        trial /= alpha;
        trial -= direction;
        const double scaledDisplacement = trial.Norml2();
        CHECK(scaledDisplacement == Approx(10.0 * alpha));
        if (previousScaledDisplacement > 0.0)
            CHECK(scaledDisplacement == Approx(previousScaledDisplacement / 2.0));
        previousScaledDisplacement = scaledDisplacement;
        CHECK(prepared.GetBaseGeneration() == capturedGeneration);
        for (int i = 0; i < 3; ++i) CHECK(direction(i) == -1.0);
    }
}

TEST_CASE("Coupled correction captures the base action while solving with the current trial Jacobian",
          "[solver][coupled_trial_correction][lifetime]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    DenseOperation operation(identity(2));
    ExactInverse inverse(operation);
    auto probe = std::make_shared<BackendProbe>();
    ScriptedCorrection prepared(operation, inverse, MPI_COMM_SELF, exactOptions(), TestBackend(probe));
    mfem::Vector f(2), p(2), prediction(2), correction(2), trialResidual(2);
    f = 1.0;
    p = -1.0;
    prepared.CaptureBase(f, p);
    const auto generation = prepared.GetBaseGeneration();
    auto matrix = identity(2);
    matrix(0, 0) = 2.0;
    matrix(1, 1) = 4.0;
    operation.Prepare(matrix);
    const int applications = operation.applications;
    prepared.BuildLinearPrediction(.5, prediction);
    CHECK(operation.applications == applications);
    CHECK(prediction(0) == .5);
    CHECK(prediction(1) == .5);
    trialResidual = prediction;
    trialResidual(0) += .2;
    trialResidual(1) += .4;
    REQUIRE_THROWS_AS(prepared.SolveCorrection(.5, trialResidual, correction), std::invalid_argument);
    inverse.Refresh();
    const auto report = prepared.SolveCorrection(.5, trialResidual, correction);
    REQUIRE(report.Admissible());
    CHECK(correction(0) == Approx(-.1));
    CHECK(correction(1) == Approx(-.1));
    CHECK(prepared.GetBaseGeneration() == generation);
    CHECK(probe->allInitialGuessesZero);
    prepared.Invalidate();
    CHECK_FALSE(prepared.HasBase());
    CHECK_THROWS_AS(prepared.BuildLinearPrediction(.5, prediction), std::invalid_argument);
}

TEST_CASE("Coupled correction independently rejects a false auxiliary convergence report",
          "[solver][coupled_trial_correction][failure]") {
    using namespace coupled_trial_correction_test;
    DenseOperation operation(identity(2));
    ExactInverse inverse(operation);
    auto probe = std::make_shared<BackendProbe>();
    ScriptedCorrection prepared(operation, inverse, MPI_COMM_SELF, exactOptions(), TestBackend(probe));
    mfem::Vector f(2), p(2), trial(2), correction(2);
    f = 1.0;
    p = -1.0;
    trial = .75;
    prepared.CaptureBase(f, p);
    SECTION("The returned correction fails the independently computed true gate") {
        probe->behavior = BackendBehavior::inaccurate;
        const auto report = prepared.SolveCorrection(.5, trial, correction);
        CHECK_FALSE(report.Admissible());
        CHECK(report.status == CoupledTrialCorrectionStatus::linear_failure);
        CHECK(report.linearSolve.Converged());
        CHECK_FALSE(report.verifiedGate);
        CHECK(report.verifiedTrueResidualNorm > report.convergenceThreshold);
    }
    SECTION("An accurate correction does not erase backend nonconvergence") {
        probe->behavior = BackendBehavior::failed_status;
        const auto report = prepared.SolveCorrection(.5, trial, correction);
        CHECK_FALSE(report.Admissible());
        CHECK(report.status == CoupledTrialCorrectionStatus::linear_failure);
        CHECK_FALSE(report.linearSolve.Converged());
    }
    SECTION("A nonfinite returned correction is never admissible") {
        probe->behavior = BackendBehavior::non_finite;
        const auto report = prepared.SolveCorrection(.5, trial, correction);
        CHECK_FALSE(report.Admissible());
        CHECK(report.status == CoupledTrialCorrectionStatus::non_finite);
        CHECK_FALSE(report.verifiedGate);
    }
    CHECK(probe->calls == 1);
    CHECK(probe->allInitialGuessesZero);
    for (int i = 0; i < 2; ++i) {
        CHECK(f(i) == 1.0);
        CHECK(p(i) == -1.0);
        CHECK(trial(i) == .75);
    }
}

TEST_CASE("Coupled correction bounds work and correction size without changing the Newton direction",
          "[solver][coupled_trial_correction][safeguard]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    DenseOperation operation(identity(2));
    ExactInverse inverse(operation);
    auto options = exactOptions();
    options.maximumCorrectionRatio = .1;
    auto probe = std::make_shared<BackendProbe>();
    ScriptedCorrection prepared(operation, inverse, MPI_COMM_SELF, options, TestBackend(probe));
    mfem::Vector f(2), p(2), trial(2), correction(2);
    f = 1.0;
    p = -1.0;
    trial = .9;
    prepared.CaptureBase(f, p);
    const auto report = prepared.SolveCorrection(.5, trial, correction);
    CHECK(report.status == CoupledTrialCorrectionStatus::norm_limit);
    CHECK_FALSE(report.Admissible());
    CHECK(report.verifiedGate);
    CHECK(report.maximumCorrectionNorm == Approx(.1 * .5 * std::sqrt(2.0)));
    CHECK(report.correctionNorm > report.maximumCorrectionNorm);
    CHECK(report.linearSolve.iterations <= options.control.maximumIterations);
    CHECK(probe->calls == 1);
    CHECK(probe->allInitialGuessesZero);
    CHECK(p(0) == -1.0);
    CHECK(p(1) == -1.0);
}

TEST_CASE("Zero coupled remainder does not manufacture an auxiliary displacement",
          "[solver][coupled_trial_correction][safeguard]") {
    using namespace coupled_trial_correction_test;
    DenseOperation operation(identity(2));
    ExactInverse inverse(operation);
    auto probe = std::make_shared<BackendProbe>();
    ScriptedCorrection prepared(operation, inverse, MPI_COMM_SELF, exactOptions(), TestBackend(probe));
    mfem::Vector f(2), p(2), trial(2), correction(2);
    f = 1.0;
    p = -1.0;
    trial = .5;
    correction = 900.0;
    prepared.CaptureBase(f, p);
    const auto report = prepared.SolveCorrection(.5, trial, correction);
    CHECK(report.Admissible());
    CHECK(report.verifiedGate);
    CHECK(correction.Norml2() == 0.0);
    CHECK(report.correctionNorm == 0.0);
    CHECK(report.verifiedTrueResidualNorm == 0.0);
    CHECK(probe->calls <= 1);
    CHECK(probe->allInitialGuessesZero);
}

TEST_CASE("Coupled correction rejects invalid controls before numerical work",
          "[solver][coupled_trial_correction][configuration]") {
    using namespace coupled_trial_correction_test;
    auto options = exactOptions();
    SECTION("A nonfinite cap cannot define a bounded correction") {
        options.maximumCorrectionRatio = std::numeric_limits<double>::infinity();
    }
    SECTION("A negative cap has no feasible correction size") { options.maximumCorrectionRatio = -1.0; }
    SECTION("A zero cap cannot enable the correction strategy") { options.maximumCorrectionRatio = 0.0; }
    SECTION("The auxiliary iteration budget must be finite and positive") { options.control.maximumIterations = 0; }
    SECTION("The auxiliary tolerance cannot be NaN") {
        options.control.relativeTolerance = std::numeric_limits<double>::quiet_NaN();
    }
    CHECK_THROWS_AS(options.Validate(), std::invalid_argument);
}

TEST_CASE("Coupled correction decisions remain collective with an empty peer rank",
          "[solver][coupled_trial_correction][mpi]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    int rank = 0;
    REQUIRE(MPI_Comm_rank(MPI_COMM_WORLD, &rank) == MPI_SUCCESS);
    // Only one rank owns this tiny synthetic problem. A distributed stop must
    // not rely on every rank having a local entry to inspect.
    const int localSize = rank == 0 ? 2 : 0;
    DenseOperation operation(identity(localSize));
    ExactInverse inverse(operation);
    auto probe = std::make_shared<BackendProbe>();
    ScriptedCorrection prepared(operation, inverse, MPI_COMM_WORLD, exactOptions(), TestBackend(probe));
    mfem::Vector f(localSize), p(localSize), trial(localSize), correction(localSize);
    f = 1.0;
    p = -1.0;
    trial = .75;
    prepared.CaptureBase(f, p);
    SECTION("Finite global norms include the owning rank exactly once") {
        const auto report = prepared.SolveCorrection(.5, trial, correction);
        CHECK(report.Admissible());
        CHECK(report.correctionNorm == Approx(std::sqrt(2.0) * .25));
        CHECK(probe->calls == 1);
    }
    SECTION("One nonfinite trial prevents the backend call on every rank") {
        if (rank == 0) trial(0) = std::numeric_limits<double>::quiet_NaN();
        const auto report = prepared.SolveCorrection(.5, trial, correction);
        CHECK(report.status == CoupledTrialCorrectionStatus::non_finite);
        CHECK_FALSE(report.Admissible());
        CHECK(probe->calls == 0);
    }
    SECTION("One nonfinite correction fails the independently checked global gate") {
        probe->behavior = BackendBehavior::non_finite;
        const auto report = prepared.SolveCorrection(.5, trial, correction);
        CHECK(report.status == CoupledTrialCorrectionStatus::non_finite);
        CHECK_FALSE(report.Admissible());
        CHECK_FALSE(report.verifiedGate);
        CHECK(probe->calls == 1);
    }
}

namespace coupled_trial_correction_test {
    struct ToyProblem final {
        MPI_Comm communicator{MPI_COMM_SELF};
        [[nodiscard]] MPI_Comm GetCommunicator() const noexcept { return communicator; }
    };

    struct ToyState final {
        DenseOperation operation{identity(3)};
        ExactInverse inverse{operation};
        DenseOperation *normalizedOperator{&operation};
        ExactInverse *normalizedInverse{&inverse};
        ExactInverse *physicalInverse{&inverse};
        mfem::Vector trialNormalizedState = mfem::Vector(3);
        mfem::Vector trialNormalizedResidual = mfem::Vector(3);
        mfem::Vector candidatePhysicalState = mfem::Vector(3);
        mfem::Vector acceptedState = mfem::Vector(3);
        mfem::Vector acceptedResidual = mfem::Vector(3);
        mfem::Vector normalizedCorrection = mfem::Vector(3);
        LinearSolveReport outerReport{.status = LinearSolveStatus::converged, .iterations = 11,
                                     .trueResidualNorm = .0123};
        ToyProblem problem{};
        int preparations{0};
        bool quarticBarrier{false}, rejectCorrectedPreparation{false}, nonfiniteCorrectedResidual{false};
        bool throwCorrectedPreparation{false};
        bool cubicSearchModel{false};
        bool rejectPositivePreparation{false};
        double barrierMultiplier{1.0};
        int throwAtPreparation{-1};
        std::vector<double> preparedCoordinates;

        ToyState() {
            trialNormalizedState = 0.0;
            candidatePhysicalState = 0.0;
            trialNormalizedResidual = 1.0;
            acceptedState = 0.0;
            acceptedResidual = 1.0;
            normalizedCorrection = -1.0;
        }
        [[nodiscard]] ToyProblem &Problem() noexcept { return problem; }
        [[nodiscard]] const ToyProblem &Problem() const noexcept { return problem; }
        [[nodiscard]] std::expected<void, int> PrepareTrial() {
            ++preparations;
            candidatePhysicalState = trialNormalizedState;
            preparedCoordinates.push_back(trialNormalizedState(0));
            if (preparations == throwAtPreparation) {
                operation.prepared = false;
                throw std::runtime_error("Synthetic delayed preparation failure.");
            }
            if (throwCorrectedPreparation && trialNormalizedState(0) < -2.0) {
                operation.prepared = false;
                throw std::runtime_error("Synthetic corrected preparation failure.");
            }
            if ((rejectCorrectedPreparation && trialNormalizedState(0) < -2.0) ||
                (rejectPositivePreparation && trialNormalizedState(0) > 0.0)) {
                operation.prepared = false;
                return std::unexpected(7);
            }
            trialNormalizedResidual = coupledResidual(trialNormalizedState);
            auto matrix = coupledJacobian(trialNormalizedState);
            if (cubicSearchModel) {
                // At alpha=.5, c=(1.5,0,0) overshoots badly at full scale.
                // Its quarter step improves the ordinary residual merit.
                const double x = trialNormalizedState(0);
                trialNormalizedResidual(0) = 1.0 + x + 10.0*x*x + 11.0*x*x*x;
                matrix = identity(3);
                matrix(0, 0) = 1.0 + 20.0*x + 33.0*x*x;
            }
            if (quarticBarrier) {
                // The extra term and its derivative vanish at the base and at
                // alpha=.5. It makes the solved correction harmful without
                // supplying the transaction an inconsistent trial Jacobian.
                const double x = trialNormalizedState(0);
                trialNormalizedResidual(0) += barrierMultiplier*x*x*(x+.5)*(x+.5);
                matrix(0, 0) += barrierMultiplier*2.0*x*(x+.5)*(2.0*x+.5);
            }
            operation.Prepare(matrix);
            if (nonfiniteCorrectedResidual && trialNormalizedState(0) < -2.0)
                trialNormalizedResidual(0) = std::numeric_limits<double>::quiet_NaN();
            return {};
        }
        [[nodiscard]] std::expected<void, int> PrepareAlpha(const double alpha) {
            trialNormalizedState = normalizedCorrection;
            trialNormalizedState *= alpha;
            return PrepareTrial();
        }
    };

    [[nodiscard]] nonlinear::MetricEvaluation toyMetric(const mfem::Vector &residual) {
        const double norm = residual.Norml2();
        return {.residualNorm = norm, .merit = .5*norm*norm};
    }
    [[nodiscard]] bool toyAccepts(const nonlinear::MetricEvaluation &metric) {
        return std::isfinite(metric.merit) && metric.merit < 1.5;
    }
    [[nodiscard]] bool toyImproves(const nonlinear::MetricEvaluation &candidate,
                                   const nonlinear::MetricEvaluation &original) {
        return std::isfinite(candidate.merit) && candidate.merit < original.merit;
    }
    [[nodiscard]] bool admissibleToy(const mfem::Vector &) { return true; }

    void checkOuterData(const ToyState &state) {
        for (int i = 0; i < 3; ++i) {
            CHECK(state.acceptedState(i) == 0.0);
            CHECK(state.acceptedResidual(i) == 1.0);
            CHECK(state.normalizedCorrection(i) == -1.0);
        }
        CHECK(state.outerReport.status == LinearSolveStatus::converged);
        CHECK(state.outerReport.iterations == 11);
        CHECK(state.outerReport.trueResidualNorm == .0123);
    }
    void checkRestoredTrial(const ToyState &state, const mfem::Vector &originalState,
                            const mfem::Vector &originalResidual, const mfem::Vector &originalAction) {
        for (int i = 0; i < 3; ++i) {
            CHECK(state.trialNormalizedState(i) == originalState(i));
            CHECK(state.candidatePhysicalState(i) == originalState(i));
            CHECK(state.trialNormalizedResidual(i) == originalResidual(i));
        }
        CHECK(state.operation.IsPrepared());
        CHECK(state.inverse.IsCurrent());
        mfem::Vector action(3);
        state.operation.Mult(state.normalizedCorrection, action);
        for (int i = 0; i < 3; ++i) CHECK(action(i) == originalAction(i));
        checkOuterData(state);
    }
}

TEST_CASE("The production coupled transaction keeps an improved full residual and leaves outer data alone",
          "[solver][coupled_trial_correction][transaction]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    ToyState state;
    solver::detail::PreparedCoupledTrialTransaction transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    REQUIRE(preparation.has_value());
    auto metric = toyMetric(state.trialNormalizedResidual);
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric, toyAccepts, toyImproves, admissibleToy);
    const auto &diagnostics = transaction.GetDiagnostics();
    CHECK(diagnostics.status == CoupledTrialAttemptStatus::kept);
    REQUIRE(diagnostics.currentSolve.has_value());
    CHECK(diagnostics.currentSolve->Admissible());
    CHECK(diagnostics.currentSolve->verifiedGate);
    CHECK(diagnostics.attempts == 1);
    CHECK(diagnostics.kept == 1);
    CHECK(diagnostics.fallbacks == 0);
    CHECK(diagnostics.preparations == 1);
    CHECK(diagnostics.restorations == 0);
    CHECK(metric.merit == Approx(.375));
    CHECK(metric.merit < diagnostics.originalMerit);
    REQUIRE(preparation.has_value());
    for (int i = 0; i < 3; ++i) CHECK(state.trialNormalizedResidual(i) == Approx(.5));
    checkOuterData(state);
    transaction.BeginTrial();
    CHECK(transaction.GetDiagnostics().status == CoupledTrialAttemptStatus::not_attempted);
    CHECK_FALSE(transaction.GetDiagnostics().currentSolve.has_value());
    CHECK(std::isnan(transaction.GetDiagnostics().originalMerit));
    CHECK(transaction.GetDiagnostics().attempts == 1);
}

TEST_CASE("The production coupled transaction restores rejected corrections and their prepared operators",
          "[solver][coupled_trial_correction][transaction][rollback]") {
    using namespace coupled_trial_correction_test;
    ToyState state;
    bool rejectAdmissibility = false;
    CoupledTrialAttemptStatus expected = CoupledTrialAttemptStatus::inadmissible;
    SECTION("Corrected physical preparation rejects the candidate") { state.rejectCorrectedPreparation = true; }
    SECTION("Corrected polynomial admissibility rejects before preparation") { rejectAdmissibility = true; }
    SECTION("A consistent nonlinear map makes the solved correction worse") {
        state.quarticBarrier = true;
        expected = CoupledTrialAttemptStatus::not_improved;
    }
    SECTION("A successfully prepared nonfinite residual is unusable") {
        state.nonfiniteCorrectedResidual = true;
        expected = CoupledTrialAttemptStatus::non_finite;
    }
    solver::detail::PreparedCoupledTrialTransaction transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    REQUIRE(preparation.has_value());
    const mfem::Vector originalState(state.trialNormalizedState), originalResidual(state.trialNormalizedResidual);
    mfem::Vector originalAction(3);
    state.operation.Mult(state.normalizedCorrection, originalAction);
    auto metric = toyMetric(originalResidual);
    const double originalMerit = metric.merit;
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric, toyAccepts, toyImproves,
        [rejectAdmissibility](const mfem::Vector &candidate) { return !rejectAdmissibility || candidate(0) >= -2.0; });
    const auto &diagnostics = transaction.GetDiagnostics();
    CHECK(diagnostics.status == expected);
    REQUIRE(diagnostics.currentSolve.has_value());
    CHECK(diagnostics.currentSolve->Admissible());
    CHECK(diagnostics.attempts == 1);
    CHECK(diagnostics.kept == 0);
    CHECK(diagnostics.fallbacks == 1);
    CHECK(diagnostics.restorations == 1);
    CHECK(diagnostics.preparations == (rejectAdmissibility ? 1 : 2));
    CHECK(metric.merit == originalMerit);
    if (expected == CoupledTrialAttemptStatus::not_improved)
        CHECK(diagnostics.correctedMerit > diagnostics.originalMerit);
    REQUIRE(preparation.has_value());
    checkRestoredTrial(state, originalState, originalResidual, originalAction);
}

TEST_CASE("The coupled transaction skips acceptable trials and bounds attempts per captured direction",
          "[solver][coupled_trial_correction][transaction][budget]") {
    using namespace coupled_trial_correction_test;
    ToyState state;
    state.quarticBarrier = true;
    solver::detail::PreparedCoupledTrialTransaction transaction(state, exactOptions(), 1);
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.05);
    auto metric = toyMetric(state.trialNormalizedResidual);
    REQUIRE(toyAccepts(metric));
    const int before = state.preparations;
    transaction.BeginTrial();
    transaction.TryImprove(state, .05, preparation, metric, toyMetric, toyAccepts, toyImproves, admissibleToy);
    CHECK(transaction.GetDiagnostics().status == CoupledTrialAttemptStatus::original_accepted);
    CHECK(transaction.GetDiagnostics().attempts == 0);
    CHECK(state.preparations == before);
    CHECK_FALSE(transaction.GetDiagnostics().currentSolve.has_value());

    preparation = state.PrepareAlpha(.5);
    metric = toyMetric(state.trialNormalizedResidual);
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric, toyAccepts, toyImproves, admissibleToy);
    CHECK(transaction.GetDiagnostics().status == CoupledTrialAttemptStatus::not_improved);
    CHECK(transaction.GetDiagnostics().attempts == 1);
    const auto iterations = transaction.GetDiagnostics().iterations;
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric, toyAccepts, toyImproves, admissibleToy);
    CHECK(transaction.GetDiagnostics().status == CoupledTrialAttemptStatus::budget_exhausted);
    CHECK(transaction.GetDiagnostics().attempts == 1);
    CHECK(transaction.GetDiagnostics().iterations == iterations);
    CHECK_FALSE(transaction.GetDiagnostics().currentSolve.has_value());
    checkOuterData(state);
}

TEST_CASE("A rank-local corrected admissibility rejection restores the original trial everywhere",
          "[solver][coupled_trial_correction][transaction][rollback][mpi]") {
    using namespace coupled_trial_correction_test;
    int rank = 0;
    REQUIRE(MPI_Comm_rank(MPI_COMM_WORLD, &rank) == MPI_SUCCESS);
    ToyState state;
    state.problem.communicator = MPI_COMM_WORLD;
    solver::detail::PreparedCoupledTrialTransaction transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    const mfem::Vector originalState(state.trialNormalizedState), originalResidual(state.trialNormalizedResidual);
    mfem::Vector originalAction(3);
    state.operation.Mult(state.normalizedCorrection, originalAction);
    auto metric = toyMetric(originalResidual);
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric, toyAccepts, toyImproves,
        [rank](const mfem::Vector &) { return rank != 0; });
    CHECK(transaction.GetDiagnostics().status == CoupledTrialAttemptStatus::inadmissible);
    CHECK(transaction.GetDiagnostics().restorations == 1);
    checkRestoredTrial(state, originalState, originalResidual, originalAction);
}

TEST_CASE("A coupled auxiliary gate or cap failure preserves the already prepared original candidate",
          "[solver][coupled_trial_correction][transaction][failure]") {
    using namespace coupled_trial_correction_test;
    ToyState state;
    auto options = exactOptions();
    auto expected = CoupledTrialAttemptStatus::norm_limit;
    SECTION("A correction outside the cap is rejected without mutation") {
        options.maximumCorrectionRatio = .1;
    }
    SECTION("A bounded unsuccessful Krylov solve is rejected without mutation") {
        // This stationary perturbation prevents one Krylov vector from spanning
        // the exact correction. The failure comes from a real verified solve.
        state.inverse.extraSecondComponent = .1;
        options.control.maximumIterations = 1;
        options.backend.restartLength = 1;
        expected = CoupledTrialAttemptStatus::linear_failure;
    }
    solver::detail::PreparedCoupledTrialTransaction transaction(state, options);
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    const mfem::Vector originalState(state.trialNormalizedState), originalResidual(state.trialNormalizedResidual);
    mfem::Vector originalAction(3);
    state.operation.Mult(state.normalizedCorrection, originalAction);
    auto metric = toyMetric(originalResidual);
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric, toyAccepts, toyImproves, admissibleToy);
    const auto &diagnostics = transaction.GetDiagnostics();
    CHECK(diagnostics.status == expected);
    REQUIRE(diagnostics.currentSolve.has_value());
    CHECK_FALSE(diagnostics.currentSolve->Admissible());
    CHECK(diagnostics.attempts == 1);
    CHECK(diagnostics.fallbacks == 1);
    CHECK(diagnostics.preparations == 0);
    CHECK(diagnostics.restorations == 0);
    CHECK(diagnostics.iterations <= static_cast<std::uint64_t>(options.control.maximumIterations));
    if (expected == CoupledTrialAttemptStatus::linear_failure)
        CHECK_FALSE(diagnostics.currentSolve->verifiedGate);
    checkRestoredTrial(state, originalState, originalResidual, originalAction);
}

TEST_CASE("Coupled correction options and step lengths must agree across ranks",
          "[solver][coupled_trial_correction][configuration][mpi]") {
    using namespace coupled_trial_correction_test;
    int rank = 0, ranks = 0;
    REQUIRE(MPI_Comm_rank(MPI_COMM_WORLD, &rank) == MPI_SUCCESS);
    REQUIRE(MPI_Comm_size(MPI_COMM_WORLD, &ranks) == MPI_SUCCESS);
    DenseOperation operation(identity(2));
    ExactInverse inverse(operation);
    auto options = exactOptions();
    SECTION("A rank-local invalid cap fails collectively") {
        if (rank == 0) options.maximumCorrectionRatio = std::numeric_limits<double>::quiet_NaN();
        const auto create = [&] { return PreparedCoupledTrialCorrection(operation, inverse, MPI_COMM_WORLD, options); };
        CHECK_THROWS_AS(create(), std::invalid_argument);
    }
    SECTION("Different valid caps fail collectively") {
        if (ranks < 2) return;
        options.maximumCorrectionRatio = rank == 0 ? 1.0 : 2.0;
        const auto create = [&] { return PreparedCoupledTrialCorrection(operation, inverse, MPI_COMM_WORLD, options); };
        CHECK_THROWS_AS(create(), std::invalid_argument);
    }
    SECTION("Different step lengths cannot form one global trial") {
        if (ranks < 2) return;
        PreparedCoupledTrialCorrection prepared(operation, inverse, MPI_COMM_WORLD, options);
        mfem::Vector f(2), p(2), prediction(2);
        f = 1.0;
        p = -1.0;
        prepared.CaptureBase(f, p);
        CHECK_THROWS_AS(prepared.BuildLinearPrediction(rank == 0 ? .5 : .25, prediction), std::invalid_argument);
    }
}

TEST_CASE("Finite entries with an overflowing auxiliary norm never authorize a correction",
          "[solver][coupled_trial_correction][failure]") {
    using namespace coupled_trial_correction_test;
    DenseOperation operation(identity(2));
    ExactInverse inverse(operation);
    auto probe = std::make_shared<BackendProbe>();
    ScriptedCorrection prepared(operation, inverse, MPI_COMM_SELF, exactOptions(), TestBackend(probe));
    mfem::Vector f(2), p(2), trial(2), correction(2);
    f = 1.0;
    p = -1.0;
    trial = 1.e200;
    prepared.CaptureBase(f, p);
    const auto report = prepared.SolveCorrection(.5, trial, correction);
    CHECK(report.status == CoupledTrialCorrectionStatus::non_finite);
    CHECK_FALSE(report.Admissible());
    CHECK_FALSE(report.verifiedGate);
    CHECK(probe->calls == 0);
}

TEST_CASE("Exceptions after coupled mutation restore the original prepared candidate before escaping",
          "[solver][coupled_trial_correction][transaction][rollback][exception]") {
    using namespace coupled_trial_correction_test;
    ToyState state;
    enum class FailureSite { admissibility, preparation, metric, comparison };
    FailureSite site = FailureSite::admissibility;
    SECTION("The admissibility callback throws") { site = FailureSite::admissibility; }
    SECTION("Physical preparation invalidates its operator and throws") {
        site = FailureSite::preparation;
        state.throwCorrectedPreparation = true;
    }
    SECTION("Evaluation of the corrected metric throws") { site = FailureSite::metric; }
    SECTION("Comparison of the corrected full merit throws") { site = FailureSite::comparison; }
    solver::detail::PreparedCoupledTrialTransaction transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    const mfem::Vector originalState(state.trialNormalizedState), originalResidual(state.trialNormalizedResidual);
    mfem::Vector originalAction(3);
    state.operation.Mult(state.normalizedCorrection, originalAction);
    auto metric = toyMetric(originalResidual);
    const double originalMerit = metric.merit;
    int evaluations = 0;
    transaction.BeginTrial();
    const auto attempt = [&] {
        transaction.TryImprove(state, .5, preparation, metric,
            [&](const mfem::Vector &residual) {
                if (++evaluations == 1 && site == FailureSite::metric)
                    throw std::runtime_error("Synthetic corrected metric failure.");
                return toyMetric(residual);
            }, toyAccepts,
            [&](const auto &candidate, const auto &original) {
                if (site == FailureSite::comparison)
                    throw std::runtime_error("Synthetic corrected comparison failure.");
                return toyImproves(candidate, original);
            }, [&](const mfem::Vector &) {
                if (site == FailureSite::admissibility)
                    throw std::runtime_error("Synthetic corrected admissibility failure.");
                return true;
            });
    };
    CHECK_THROWS_AS(attempt(), std::runtime_error);
    CHECK(transaction.GetDiagnostics().kept == 0);
    CHECK(transaction.GetDiagnostics().restorations == 1);
    CHECK(metric.merit == originalMerit);
    REQUIRE(preparation.has_value());
    checkRestoredTrial(state, originalState, originalResidual, originalAction);
}

TEST_CASE("Damping one verified coupled correction accepts a nonlinear step that its full scale misses",
          "[solver][coupled_trial_correction][backtracked]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    ToyState state;
    state.cubicSearchModel = true;
    SECTION("Finite full and half corrections overshoot the nonlinear residual") { }
    SECTION("Rejected full and half preparations are followed by a valid smaller correction") {
        state.rejectPositivePreparation = true;
    }
    solver::detail::PreparedCoupledTrialTransaction<ToyState, BacktrackedCoupledStep> transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    auto metric = toyMetric(state.trialNormalizedResidual);
    REQUIRE(metric.merit == Approx(1.5703125));
    const auto accepts = [](const auto &candidate) { return candidate.merit <= 1.5*(1.0-1.e-4*.5); };
    REQUIRE_FALSE(accepts(metric));
    const int initialActions = state.operation.applications, initialInverses = state.inverse.applications;
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric, accepts, toyImproves, admissibleToy);
    const auto &diagnostics = transaction.GetDiagnostics();
    REQUIRE(diagnostics.currentSolve.has_value());
    REQUIRE(diagnostics.currentSolve->Admissible());
    REQUIRE(diagnostics.currentSolve->verifiedGate);
    CHECK(diagnostics.status == CoupledTrialAttemptStatus::kept);
    CHECK(diagnostics.selectedBeta == .25);
    CHECK(diagnostics.attempts == 1);
    CHECK(diagnostics.kept == 1);
    CHECK(diagnostics.fallbacks == 0);
    CHECK(diagnostics.candidateCount == 3);
    REQUIRE(diagnostics.candidates.size() == 3);
    CHECK(diagnostics.candidates[0].beta == 1.0);
    CHECK_FALSE(diagnostics.candidates[0].accepted);
    if (state.rejectPositivePreparation) {
        CHECK(diagnostics.candidates[0].status == CoupledTrialCandidateStatus::inadmissible);
        CHECK(diagnostics.candidates[1].status == CoupledTrialCandidateStatus::inadmissible);
    } else CHECK(diagnostics.candidates[0].merit > diagnostics.originalMerit);
    CHECK(diagnostics.candidates[1].beta == .5);
    CHECK_FALSE(diagnostics.candidates[1].accepted);
    CHECK(diagnostics.candidates[2].beta == .25);
    CHECK(diagnostics.candidates[2].accepted);
    CHECK(diagnostics.preparations == 3);
    CHECK(diagnostics.restorations == 0);
    CHECK(transaction.GetCorrection()(0) == Approx(1.5));
    CHECK(transaction.GetSelectedCorrection()(0) == Approx(.375));
    CHECK(state.trialNormalizedState(0) == Approx(-.125));
    CHECK(state.trialNormalizedResidual(0) == Approx(1.009765625));
    CHECK(accepts(metric));
    // Candidate preparations change J but must not launch another solve or
    // verify a scaled correction against a different candidate Jacobian.
    CHECK(state.operation.applications-initialActions == diagnostics.currentSolve->linearSolve.operatorApplications+1);
    CHECK(state.inverse.applications-initialInverses == diagnostics.currentSolve->linearSolve.inversePreconditionerApplications);
    CHECK(diagnostics.iterations == static_cast<std::uint64_t>(diagnostics.currentSolve->linearSolve.iterations));
    CHECK(state.operation.IsPrepared());
    CHECK_FALSE(state.inverse.IsCurrent());
    checkOuterData(state);
    transaction.BeginTrial();
    CHECK(std::isnan(transaction.GetDiagnostics().selectedBeta));
    CHECK(transaction.GetDiagnostics().candidateCount == 0);
    CHECK(transaction.GetDiagnostics().candidates.empty());
    CHECK_FALSE(transaction.GetDiagnostics().currentSolve.has_value());
    CHECK(transaction.GetSelectedCorrection().Norml2() == 0.0);
}

TEST_CASE("A bounded coupled search freshly restores its best earlier beta when no beta passes Armijo",
          "[solver][coupled_trial_correction][backtracked][rollback]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    ToyState state;
    state.cubicSearchModel = true;
    solver::detail::PreparedCoupledTrialTransaction<ToyState, BacktrackedCoupledStep> transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    auto metric = toyMetric(state.trialNormalizedResidual);
    // This stringent but valid Armijo parameter excludes the quarter step;
    // it remains the best full-merit improvement among all four proposals.
    const auto accepts = [](const auto &candidate) { return candidate.merit <= 1.5*(1.0-.99*.5); };
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric, accepts, toyImproves, admissibleToy);
    const auto &diagnostics = transaction.GetDiagnostics();
    CHECK(diagnostics.status == CoupledTrialAttemptStatus::kept);
    CHECK(diagnostics.selectedBeta == .25);
    CHECK(diagnostics.candidateCount == 4);
    REQUIRE(diagnostics.candidates.size() == 4);
    for (const auto &candidate : diagnostics.candidates) CHECK_FALSE(candidate.accepted);
    CHECK(diagnostics.candidates[2].merit < diagnostics.candidates[3].merit);
    CHECK_FALSE(accepts(metric));
    CHECK(metric.merit < diagnostics.originalMerit);
    CHECK(diagnostics.preparations == 5);
    CHECK(diagnostics.restorations == 1);
    CHECK(diagnostics.fallbacks == 0);
    REQUIRE(state.preparedCoordinates.size() == 6);
    CHECK(state.preparedCoordinates[4] == Approx(-.3125));
    CHECK(state.preparedCoordinates[5] == Approx(-.125));
    CHECK(state.trialNormalizedState(0) == Approx(-.125));
    CHECK(state.candidatePhysicalState(0) == Approx(-.125));
    CHECK(state.trialNormalizedResidual(0) == Approx(1.009765625));
    CHECK(state.operation.Matrix()(0, 0) == Approx(1.0+20.0*(-.125)+33.0*.125*.125));
    CHECK(state.operation.IsPrepared());
    checkOuterData(state);
}

TEST_CASE("Every rejected damping proposal falls back to the original prepared trial",
          "[solver][coupled_trial_correction][backtracked][rollback]") {
    using namespace coupled_trial_correction_test;
    ToyState state;
    bool rejectAdmissibility = false, invalidMetric = false;
    auto expected = CoupledTrialAttemptStatus::not_improved;
    SECTION("Every nonlinear correction has worse full merit") {
        state.quarticBarrier = true;
        state.barrierMultiplier = 100.0;
    }
    SECTION("Every beta fails admissibility before preparation") {
        rejectAdmissibility = true;
        expected = CoupledTrialAttemptStatus::inadmissible;
    }
    SECTION("A finite merit cannot hide a nonfinite metric norm") {
        invalidMetric = true;
        expected = CoupledTrialAttemptStatus::non_finite;
    }
    solver::detail::PreparedCoupledTrialTransaction<ToyState, BacktrackedCoupledStep> transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    const mfem::Vector originalState(state.trialNormalizedState), originalResidual(state.trialNormalizedResidual);
    mfem::Vector originalAction(3);
    state.operation.Mult(state.normalizedCorrection, originalAction);
    auto metric = toyMetric(originalResidual);
    const double originalMerit = metric.merit;
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric,
        [&](const mfem::Vector &residual) {
            if (invalidMetric && state.trialNormalizedState(0) != originalState(0))
                return nonlinear::MetricEvaluation{.residualNorm = std::numeric_limits<double>::quiet_NaN(), .merit = .1};
            return toyMetric(residual);
        }, toyAccepts, toyImproves, [rejectAdmissibility](const mfem::Vector &) { return !rejectAdmissibility; });
    const auto &diagnostics = transaction.GetDiagnostics();
    CHECK(diagnostics.status == expected);
    REQUIRE(diagnostics.currentSolve.has_value());
    CHECK(diagnostics.currentSolve->Admissible());
    CHECK(diagnostics.attempts == 1);
    CHECK(diagnostics.candidateCount == 4);
    CHECK(diagnostics.kept == 0);
    CHECK(diagnostics.fallbacks == 1);
    CHECK(diagnostics.restorations == 1);
    CHECK(diagnostics.preparations == (rejectAdmissibility ? 1 : 5));
    CHECK(std::isnan(diagnostics.selectedBeta));
    CHECK(transaction.GetSelectedCorrection().Norml2() == 0.0);
    CHECK(metric.merit == originalMerit);
    REQUIRE(preparation.has_value());
    checkRestoredTrial(state, originalState, originalResidual, originalAction);
}

TEST_CASE("Exceptions during later damping or best-candidate replay restore the original trial",
          "[solver][coupled_trial_correction][backtracked][rollback][exception]") {
    using namespace coupled_trial_correction_test;
    ToyState state;
    state.cubicSearchModel = true;
    enum class FailureSite { later_admissibility, later_metric, replay_preparation, replay_metric };
    FailureSite site = FailureSite::later_admissibility;
    SECTION("The second admissibility callback throws") { site = FailureSite::later_admissibility; }
    SECTION("The second candidate metric throws") { site = FailureSite::later_metric; }
    SECTION("Repreparing the earlier best invalidates J and throws") { site = FailureSite::replay_preparation; }
    SECTION("The freshly replayed best metric throws") { site = FailureSite::replay_metric; }
    solver::detail::PreparedCoupledTrialTransaction<ToyState, BacktrackedCoupledStep> transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    const mfem::Vector originalState(state.trialNormalizedState), originalResidual(state.trialNormalizedResidual);
    mfem::Vector originalAction(3);
    state.operation.Mult(state.normalizedCorrection, originalAction);
    auto metric = toyMetric(originalResidual);
    const double originalMerit = metric.merit;
    if (site == FailureSite::replay_preparation) state.throwAtPreparation = state.preparations+5;
    int metricCalls = 0, admissibilityCalls = 0;
    transaction.BeginTrial();
    const auto attempt = [&] {
        transaction.TryImprove(state, .5, preparation, metric,
            [&](const mfem::Vector &residual) {
                ++metricCalls;
                if ((site == FailureSite::later_metric && metricCalls == 2) ||
                    (site == FailureSite::replay_metric && metricCalls == 5))
                    throw std::runtime_error("Synthetic damping metric failure.");
                return toyMetric(residual);
            }, [](const auto &candidate) { return candidate.merit <= 1.5*(1.0-.99*.5); }, toyImproves,
            [&](const mfem::Vector &) {
                if (++admissibilityCalls == 2 && site == FailureSite::later_admissibility)
                    throw std::runtime_error("Synthetic later admissibility failure.");
                return true;
            });
    };
    CHECK_THROWS_AS(attempt(), std::runtime_error);
    CHECK(transaction.GetDiagnostics().kept == 0);
    CHECK(std::isnan(transaction.GetDiagnostics().selectedBeta));
    CHECK(transaction.GetSelectedCorrection().Norml2() == 0.0);
    CHECK(metric.merit == originalMerit);
    REQUIRE(preparation.has_value());
    checkRestoredTrial(state, originalState, originalResidual, originalAction);
}

TEST_CASE("Near-acceptance selection defers work without spending the captured direction budget",
          "[solver][coupled_trial_correction][backtracked][budget]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    ToyState state;
    solver::detail::PreparedCoupledTrialTransaction<ToyState, BacktrackedCoupledStep,
        NearAcceptanceCoupledAttempt> transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    auto metric = toyMetric(state.trialNormalizedResidual);
    const int actions = state.operation.applications, inverses = state.inverse.applications, preparations = state.preparations;
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric,
        [](const auto &candidate) { return candidate.merit <= 1.5*(1.0-1.e-4*.5); }, toyImproves, admissibleToy);
    CHECK(transaction.GetDiagnostics().status == CoupledTrialAttemptStatus::prediction_deferred);
    CHECK(transaction.GetDiagnostics().baseMerit == Approx(1.5));
    CHECK(transaction.GetDiagnostics().linearPredictionMerit == Approx(.375));
    CHECK(transaction.GetDiagnostics().predictedDecrease == Approx(1.125));
    CHECK(transaction.GetDiagnostics().predictionExcessRatio == Approx(26.0/9.0));
    CHECK(transaction.GetDiagnostics().deferredAttempts == 1);
    CHECK(transaction.GetDiagnostics().attempts == 0);
    CHECK_FALSE(transaction.GetDiagnostics().currentSolve.has_value());
    CHECK(state.operation.applications == actions);
    CHECK(state.inverse.applications == inverses);
    CHECK(state.preparations == preparations);

    preparation = state.PrepareAlpha(.25);
    metric = toyMetric(state.trialNormalizedResidual);
    REQUIRE_FALSE(toyAccepts(metric));
    transaction.BeginTrial();
    CHECK(std::isnan(transaction.GetDiagnostics().predictionExcessRatio));
    transaction.TryImprove(state, .25, preparation, metric, toyMetric, toyAccepts, toyImproves, admissibleToy);
    CHECK(transaction.GetDiagnostics().predictionExcessRatio == Approx(1.0/84.0));
    CHECK(transaction.GetDiagnostics().status == CoupledTrialAttemptStatus::kept);
    CHECK(transaction.GetDiagnostics().selectedBeta == 1.0);
    CHECK(transaction.GetDiagnostics().attempts == 1);
    CHECK(transaction.GetDiagnostics().deferredAttempts == 1);
    CHECK(metric.merit == Approx(.84375));
    CHECK(toyAccepts(metric));

    preparation = state.PrepareAlpha(.25);
    metric = toyMetric(state.trialNormalizedResidual);
    transaction.BeginTrial();
    transaction.TryImprove(state, .25, preparation, metric, toyMetric, toyAccepts, toyImproves, admissibleToy);
    CHECK(transaction.GetDiagnostics().status == CoupledTrialAttemptStatus::budget_exhausted);
    CHECK(transaction.GetDiagnostics().attempts == 1);
    CHECK_FALSE(transaction.GetDiagnostics().currentSolve.has_value());
    CHECK(transaction.GetDiagnostics().candidateCount == 0);
    CHECK(transaction.GetDiagnostics().candidates.empty());
    CHECK(std::isnan(transaction.GetDiagnostics().selectedBeta));
    CHECK(transaction.GetSelectedCorrection().Norml2() == 0.0);
    checkOuterData(state);
}

TEST_CASE("Near-acceptance selection skips acceptable trials and requires a predicted decrease",
          "[solver][coupled_trial_correction][backtracked][budget]") {
    using namespace coupled_trial_correction_test;
    ToyState state;
    solver::detail::PreparedCoupledTrialTransaction<ToyState, BacktrackedCoupledStep,
        NearAcceptanceCoupledAttempt> transaction(state, exactOptions());
    mfem::Vector direction(state.normalizedCorrection);
    bool accepted = false;
    SECTION("An already acceptable original never starts auxiliary work") { accepted = true; }
    SECTION("A zero predicted decrease cannot define a near-acceptance allowance") { direction = 0.0; }
    transaction.CaptureBase(state.acceptedResidual, direction);
    auto preparation = state.PrepareAlpha(accepted ? .05 : 0.0);
    auto metric = toyMetric(state.trialNormalizedResidual);
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric, toyMetric,
        [](const auto &candidate) { return candidate.merit <= 1.5*(1.0-1.e-4*.5); }, toyImproves, admissibleToy);
    CHECK(transaction.GetDiagnostics().status == (accepted ? CoupledTrialAttemptStatus::original_accepted :
        CoupledTrialAttemptStatus::prediction_deferred));
    CHECK(transaction.GetDiagnostics().attempts == 0);
    CHECK_FALSE(transaction.GetDiagnostics().currentSolve.has_value());
    CHECK(transaction.GetDiagnostics().candidateCount == 0);
    CHECK(transaction.GetSelectedCorrection().Norml2() == 0.0);
    if (!accepted) CHECK(transaction.GetDiagnostics().predictedDecrease == 0.0);
    checkOuterData(state);
}

TEST_CASE("Damped coupled candidates make collective admissibility and metric decisions",
          "[solver][coupled_trial_correction][backtracked][mpi]") {
    using namespace coupled_trial_correction_test;
    using Catch::Approx;
    int rank = 0;
    REQUIRE(MPI_Comm_rank(MPI_COMM_WORLD, &rank) == MPI_SUCCESS);
    ToyState state;
    state.problem.communicator = MPI_COMM_WORLD;
    state.cubicSearchModel = true;
    bool invalidMetric = false, strictAcceptance = false;
    SECTION("One rank rejects the first two admissibility checks") { }
    SECTION("One rank reports a nonfinite quarter-step metric") { invalidMetric = true; }
    SECTION("One rank requires the full bounded search before keeping its best") { strictAcceptance = true; }
    solver::detail::PreparedCoupledTrialTransaction<ToyState, BacktrackedCoupledStep> transaction(state, exactOptions());
    transaction.CaptureBase(state.acceptedResidual, state.normalizedCorrection);
    auto preparation = state.PrepareAlpha(.5);
    auto metric = toyMetric(state.trialNormalizedResidual);
    transaction.BeginTrial();
    transaction.TryImprove(state, .5, preparation, metric,
        [&](const mfem::Vector &residual) {
            if (invalidMetric && rank == 0 && std::abs(state.trialNormalizedState(0)+.125) < 1.e-12)
                return nonlinear::MetricEvaluation{.residualNorm = std::numeric_limits<double>::quiet_NaN(), .merit = .1};
            return toyMetric(residual);
        }, [&](const auto &candidate) {
            return candidate.merit <= 1.5*(1.0-(strictAcceptance && rank == 0 ? .99 : 1.e-4)*.5);
        }, toyImproves, [&](const mfem::Vector &candidate) {
            return invalidMetric || strictAcceptance || rank != 0 || candidate(0) <= 0.0;
        });
    const auto &diagnostics = transaction.GetDiagnostics();
    REQUIRE(diagnostics.currentSolve.has_value());
    REQUIRE(diagnostics.currentSolve->Admissible());
    CHECK(diagnostics.status == CoupledTrialAttemptStatus::kept);
    CHECK(diagnostics.attempts == 1);
    CHECK(diagnostics.selectedBeta == (invalidMetric ? .125 : .25));
    CHECK(diagnostics.candidateCount == (invalidMetric || strictAcceptance ? 4 : 3));
    REQUIRE(diagnostics.candidates.size() == diagnostics.candidateCount);
    if (invalidMetric) CHECK(diagnostics.candidates[2].status == CoupledTrialCandidateStatus::non_finite);
    else if (strictAcceptance) CHECK_FALSE(diagnostics.candidates[2].accepted);
    else {
        CHECK(diagnostics.candidates[0].status == CoupledTrialCandidateStatus::inadmissible);
        CHECK(diagnostics.candidates[1].status == CoupledTrialCandidateStatus::inadmissible);
        CHECK(diagnostics.preparations == 1);
    }
    CHECK(state.trialNormalizedState(0) == Approx(invalidMetric ? -.3125 : -.125));
    checkOuterData(state);
}

TEST_CASE("Damped coupled search controls reject invalid or inconsistent work budgets collectively",
          "[solver][coupled_trial_correction][backtracked][configuration][mpi]") {
    using namespace coupled_trial_correction_test;
    int rank = 0, ranks = 0;
    REQUIRE(MPI_Comm_rank(MPI_COMM_WORLD, &rank) == MPI_SUCCESS);
    REQUIRE(MPI_Comm_size(MPI_COMM_WORLD, &ranks) == MPI_SUCCESS);
    ToyState state;
    state.problem.communicator = MPI_COMM_WORLD;
    BacktrackedCoupledStep search;
    NearAcceptanceCoupledAttempt selection;
    SECTION("A rank-local zero candidate budget cannot bound a useful search") { if (rank == 0) search.maximumSteps = 0; }
    SECTION("A rank-local nondecreasing scale is invalid") { if (rank == 0) search.reduction = 1.0; }
    SECTION("A rank-local zero scale is invalid") { if (rank == 0) search.reduction = 0.0; }
    SECTION("A rank-local nonfinite scale is invalid") { if (rank == 0) search.reduction = std::numeric_limits<double>::quiet_NaN(); }
    SECTION("A nominally finite budget whose last scale underflows is invalid") {
        if (rank == 0) search.maximumSteps = 4096;
    }
    SECTION("A rank-local zero acceptance allowance is invalid") { if (rank == 0) selection.maximumExcessRatio = 0.0; }
    SECTION("A rank-local nonfinite allowance is invalid") { if (rank == 0) selection.maximumExcessRatio = std::numeric_limits<double>::infinity(); }
    SECTION("Different valid candidate budgets cannot share a collective search") {
        if (ranks < 2) return;
        search.maximumSteps = rank == 0 ? 4 : 3;
    }
    SECTION("Different valid reductions cannot share a collective search") {
        if (ranks < 2) return;
        search.reduction = rank == 0 ? .5 : .25;
    }
    SECTION("Different valid allowances cannot select one global attempt") {
        if (ranks < 2) return;
        selection.maximumExcessRatio = rank == 0 ? 1.0 : 2.0;
    }
    const auto create = [&] {
        return solver::detail::PreparedCoupledTrialTransaction(state, exactOptions(), 1, search, selection);
    };
    CHECK_THROWS_AS(create(), std::invalid_argument);
}