module;

#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <type_traits>
#include <utility>

#include <mfem.hpp>
#include <mpi.h>

export module mean_field:solver.coupled_trial_correction;

export import :solver.linear_backend;

export namespace mean_field::solver {
    struct CoupledTrialCorrectionOptions final {
        LinearSolveControl control{.relativeTolerance = .03, .absoluteTolerance = 1.e-14, .maximumIterations = 160};
        linear::FGMRESOptions backend{.restartLength = 160, .printLevel = -1};
        double maximumCorrectionRatio{1.0};

        void Validate() const {
            control.Validate();
            backend.Validate();
            if (control.relativeTolerance <= 0.0 || !std::isfinite(maximumCorrectionRatio) ||
                maximumCorrectionRatio <= 0.0)
                throw std::invalid_argument("Coupled trial correction requires positive relative tolerance and norm cap.");
        }
    };

    enum class CoupledTrialCorrectionStatus : std::uint8_t { admissible, linear_failure, norm_limit, non_finite };

    struct CoupledTrialCorrectionReport final {
        CoupledTrialCorrectionStatus status{CoupledTrialCorrectionStatus::linear_failure};
        LinearSolveReport linearSolve{};
        double verifiedTrueResidualNorm{std::numeric_limits<double>::quiet_NaN()};
        double convergenceThreshold{std::numeric_limits<double>::quiet_NaN()};
        double correctionNorm{std::numeric_limits<double>::quiet_NaN()};
        double maximumCorrectionNorm{std::numeric_limits<double>::quiet_NaN()};
        double verificationSeconds{0.0};
        bool verifiedGate{false};
        bool verificationApplied{false};
        [[nodiscard]] bool Admissible() const noexcept { return status == CoupledTrialCorrectionStatus::admissible; }
    };

    enum class CoupledTrialAttemptStatus : std::uint8_t {
        not_attempted, original_accepted, budget_exhausted, kept, linear_failure, norm_limit,
        non_finite, inadmissible, not_improved
    };

    struct CoupledTrialDiagnostics final {
        std::uint64_t attempts{0}, kept{0}, fallbacks{0}, iterations{0}, preparations{0}, restorations{0};
        double solveSeconds{0.0}, verificationSeconds{0.0}, preparationSeconds{0.0}, refreshSeconds{0.0};
        // These fields describe the current trial; BeginTrial clears them even
        // when an earlier preparation rejection prevents another auxiliary solve.
        CoupledTrialAttemptStatus status{CoupledTrialAttemptStatus::not_attempted};
        std::optional<CoupledTrialCorrectionReport> currentSolve{};
        double originalMerit{std::numeric_limits<double>::quiet_NaN()};
        double correctedMerit{std::numeric_limits<double>::quiet_NaN()};
    };
}

namespace mean_field::solver::detail {
    template <typename Operation>
    concept CoupledGenerationAvailable = requires(const Operation &operation) {
        { operation.GetPreparationGeneration() } -> std::convertible_to<std::uint64_t>;
    } || requires(const Operation &operation) {
        { operation.GetProblem().GetPreparationGeneration() } -> std::convertible_to<std::uint64_t>;
    };

    template <CoupledGenerationAvailable Operation>
    std::uint64_t CoupledGeneration(const Operation &operation) {
        if constexpr (requires { operation.GetPreparationGeneration(); }) return operation.GetPreparationGeneration();
        else return operation.GetProblem().GetPreparationGeneration();
    }

    inline bool CoupledFinite(const mfem::Vector &values) {
        for (int index = 0; index < values.Size(); ++index)
            if (!std::isfinite(values(index))) return false;
        return true;
    }
    inline bool CoupledAll(bool condition, MPI_Comm communicator) {
        int local = condition, global = 0;
        if (MPI_Allreduce(&local, &global, 1, MPI_INT, MPI_MIN, communicator) != MPI_SUCCESS)
            throw std::runtime_error("A coupled trial correction could not reduce its validity check.");
        return global != 0;
    }
    inline double CoupledNorm(const mfem::Vector &values, MPI_Comm communicator) {
        const double local = values * values;
        double global = 0.0;
        if (MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_SUM, communicator) != MPI_SUCCESS)
            throw std::runtime_error("A coupled trial correction could not reduce its norm.");
        return std::sqrt(global);
    }
    inline void CoupledRequire(bool condition, MPI_Comm communicator, const char *message) {
        if (!CoupledAll(condition, communicator)) throw std::invalid_argument(message);
    }
    inline CoupledTrialCorrectionOptions PrepareCoupledOptions(CoupledTrialCorrectionOptions options, MPI_Comm communicator) {
        int initialized = 0, finalized = 0;
        if (MPI_Initialized(&initialized) != MPI_SUCCESS || !initialized ||
            MPI_Finalized(&finalized) != MPI_SUCCESS || finalized)
            throw std::logic_error("A coupled trial correction requires initialized MPI that has not been finalized.");
        if (communicator == MPI_COMM_NULL) throw std::invalid_argument("A coupled trial correction requires a communicator.");
        bool valid = true;
        try { options.Validate(); } catch (const std::invalid_argument &) { valid = false; }
        CoupledRequire(valid, communicator, "Coupled trial correction controls are invalid on at least one rank.");
        const std::array<int, 3> integers{options.control.maximumIterations, options.backend.restartLength, options.backend.printLevel};
        const std::array<double, 3> reals{options.control.relativeTolerance, options.control.absoluteTolerance, options.maximumCorrectionRatio};
        std::array<int, 3> imin{}, imax{};
        std::array<double, 3> rmin{}, rmax{};
        if (MPI_Allreduce(integers.data(), imin.data(), 3, MPI_INT, MPI_MIN, communicator) != MPI_SUCCESS ||
            MPI_Allreduce(integers.data(), imax.data(), 3, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS ||
            MPI_Allreduce(reals.data(), rmin.data(), 3, MPI_DOUBLE, MPI_MIN, communicator) != MPI_SUCCESS ||
            MPI_Allreduce(reals.data(), rmax.data(), 3, MPI_DOUBLE, MPI_MAX, communicator) != MPI_SUCCESS)
            throw std::runtime_error("Coupled trial correction controls could not be compared across ranks.");
        if (imin != imax || rmin != rmax)
            throw std::invalid_argument("Coupled trial correction controls must agree across ranks.");
        return options;
    }
    template <typename Operation, typename Inverse>
    CoupledTrialCorrectionOptions PrepareCoupledSystemOptions(CoupledTrialCorrectionOptions options,
        const Operation &operation, const Inverse &inverse, MPI_Comm communicator) {
        options = PrepareCoupledOptions(options, communicator);
        CoupledRequire(operation.Height() >= 0 && operation.Width() == operation.Height() &&
            inverse.Height() == operation.Width() && inverse.Width() == operation.Height(), communicator,
            "A coupled correction requires compatible square operator and inverse dimensions.");
        const long long localSize = operation.Height();
        long long globalSize = 0;
        if (MPI_Allreduce(&localSize, &globalSize, 1, MPI_LONG_LONG, MPI_SUM, communicator) != MPI_SUCCESS)
            throw std::runtime_error("A coupled correction could not determine its global size.");
        if (globalSize == 0) throw std::invalid_argument("A coupled correction requires a globally nonempty system.");
        return options;
    }
    inline linear::FGMRES PrepareCoupledBackend(const CoupledTrialCorrectionOptions &options, MPI_Comm communicator) {
        return linear::FGMRES(PrepareCoupledOptions(options, communicator).backend);
    }
}

export namespace mean_field::solver {
    template <typename Operation>
    concept CoupledTrialOperation = std::derived_from<Operation, mfem::Operator> &&
        detail::CoupledGenerationAvailable<Operation> && requires(const Operation &operation) {
            { operation.IsPrepared() } -> std::same_as<bool>;
        };

    template <typename Inverse>
    concept CoupledTrialInverse = std::derived_from<Inverse, mfem::Solver> && requires(const Inverse &inverse) {
        { inverse.IsCurrent() } -> std::same_as<bool>;
    };

    // Capture fixes the affine target at the accepted linearization. Later
    // SolveCorrection borrows the freshly prepared candidate Jacobian; neither
    // operation owns or mutates a model state or the outer Newton direction.
    template <CoupledTrialOperation Operation, CoupledTrialInverse Inverse,
              LinearBackendConfiguration Backend = linear::FGMRES>
        requires LinearBackendRuntimeAvailableFor<Backend, Operation, Inverse>
    class PreparedCoupledTrialCorrection final {
    public:
        PreparedCoupledTrialCorrection(const Operation &operation, Inverse &inverse, MPI_Comm communicator,
                                      CoupledTrialCorrectionOptions options = {})
            requires std::same_as<Backend, linear::FGMRES>
            : PreparedCoupledTrialCorrection(operation, inverse, communicator, options,
                                             detail::PrepareCoupledBackend(options, communicator)) { }

        PreparedCoupledTrialCorrection(const Operation &operation, Inverse &inverse, MPI_Comm communicator,
                                      CoupledTrialCorrectionOptions options, Backend backend)
            : m_operation(std::addressof(operation)), m_inverse(std::addressof(inverse)), m_communicator(communicator),
              m_options(detail::PrepareCoupledSystemOptions(options, operation, inverse, communicator)),
              m_backend(prepareLinearBackend(std::move(backend), operation, inverse, communicator)),
              m_baseResidual(operation.Height()), m_baseAction(operation.Height()), m_rightHandSide(operation.Height()),
              m_verifiedAction(operation.Height()), m_defect(operation.Height()) {
            detail::CoupledRequire(IsCurrent(), communicator, "Coupled correction preparation requires a current operator and inverse.");
        }
        PreparedCoupledTrialCorrection(const PreparedCoupledTrialCorrection &) = delete;
        PreparedCoupledTrialCorrection &operator=(const PreparedCoupledTrialCorrection &) = delete;
        PreparedCoupledTrialCorrection(PreparedCoupledTrialCorrection &&) = delete;
        PreparedCoupledTrialCorrection &operator=(PreparedCoupledTrialCorrection &&) = delete;

        [[nodiscard]] bool IsCurrent() const {
            return m_operation->IsPrepared() && m_inverse->IsCurrent() && m_backend.IsReady() &&
                   m_operation->Height() == m_operation->Width() &&
                   m_inverse->Height() == m_operation->Width() && m_inverse->Width() == m_operation->Height() &&
                   m_baseResidual.Size() == m_operation->Height() && m_baseAction.Size() == m_operation->Height();
        }
        [[nodiscard]] bool HasBase() const noexcept { return m_hasBase; }
        [[nodiscard]] std::uint64_t GetBaseGeneration() const noexcept { return m_baseGeneration; }
        [[nodiscard]] double GetBaseDirectionNorm() const noexcept { return m_directionNorm; }
        [[nodiscard]] MPI_Comm GetCommunicator() const noexcept { return m_communicator; }
        [[nodiscard]] const mfem::Vector &GetRightHandSide() const noexcept { return m_rightHandSide; }
        [[nodiscard]] const mfem::Vector &GetVerifiedAction() const noexcept { return m_verifiedAction; }
        [[nodiscard]] const mfem::Vector &GetVerifiedDefect() const noexcept { return m_defect; }
        [[nodiscard]] const auto &GetLastReport() const noexcept { return m_lastReport; }
        [[nodiscard]] const CoupledTrialCorrectionOptions &GetOptions() const noexcept { return m_options; }
        void Invalidate() noexcept { m_hasBase = false; }

        void CaptureBase(const mfem::Vector &residual, const mfem::Vector &direction) {
            m_hasBase = false;
            detail::CoupledRequire(IsCurrent() && residual.Size() == m_operation->Height() &&
                direction.Size() == m_operation->Width() && detail::CoupledFinite(residual) && detail::CoupledFinite(direction),
                m_communicator, "Coupled correction capture requires a current base and finite compatible vectors.");
            m_baseResidual = residual;
            m_operation->Mult(direction, m_baseAction);
            m_directionNorm = detail::CoupledNorm(direction, m_communicator);
            detail::CoupledRequire(detail::CoupledFinite(m_baseAction) && std::isfinite(m_directionNorm), m_communicator,
                                   "The captured coupled correction tangent is nonfinite.");
            m_baseGeneration = detail::CoupledGeneration(*m_operation);
            m_hasBase = true;
        }

        void BuildLinearPrediction(double alpha, mfem::Vector &prediction) const {
            detail::CoupledRequire(m_hasBase && prediction.Size() == m_baseResidual.Size(), m_communicator,
                                   "A coupled prediction requires a captured base and compatible output.");
            RequireAlpha(alpha);
            prediction = m_baseResidual;
            prediction.Add(alpha, m_baseAction);
            detail::CoupledRequire(detail::CoupledFinite(prediction), m_communicator, "The coupled affine target is nonfinite.");
        }

        [[nodiscard]] CoupledTrialCorrectionReport SolveCorrection(double alpha, const mfem::Vector &trialResidual,
                                                                   mfem::Vector &correction) {
            detail::CoupledRequire(m_hasBase && IsCurrent() && detail::CoupledGeneration(*m_operation) >= m_baseGeneration &&
                trialResidual.Size() == m_operation->Height() && correction.Size() == m_operation->Width(),
                m_communicator, "Coupled correction requires a captured base, current candidate and compatible vectors.");
            RequireAlpha(alpha);
            CoupledTrialCorrectionReport report;
            const auto finish = [&]() { m_lastReport = report; return report; };
            correction = 0.0;
            m_verifiedAction = std::numeric_limits<double>::quiet_NaN();
            m_defect = std::numeric_limits<double>::quiet_NaN();
            report.maximumCorrectionNorm = m_options.maximumCorrectionRatio * alpha * m_directionNorm;
            BuildLinearPrediction(alpha, m_rightHandSide);
            m_rightHandSide -= trialResidual;
            if (!detail::CoupledAll(detail::CoupledFinite(m_rightHandSide) && std::isfinite(report.maximumCorrectionNorm), m_communicator)) {
                report.status = CoupledTrialCorrectionStatus::non_finite;
                return finish();
            }
            const double rhsNorm = detail::CoupledNorm(m_rightHandSide, m_communicator);
            const double relativeThreshold = m_options.control.relativeTolerance * rhsNorm;
            if (!std::isfinite(rhsNorm) || !std::isfinite(relativeThreshold)) {
                report.status = CoupledTrialCorrectionStatus::non_finite;
                return finish();
            }
            report.convergenceThreshold = m_options.control.ConvergenceThreshold(rhsNorm);
            report.linearSolve = m_backend.Solve(m_rightHandSide, correction, m_options.control);
            if (!detail::CoupledAll(correction.Size() == m_operation->Width() && detail::CoupledFinite(correction), m_communicator)) {
                report.status = CoupledTrialCorrectionStatus::non_finite;
                return finish();
            }
            const auto start = std::chrono::steady_clock::now();
            m_operation->Mult(correction, m_verifiedAction);
            report.verificationApplied = true;
            m_defect = m_verifiedAction;
            m_defect -= m_rightHandSide;
            const bool finite = detail::CoupledAll(detail::CoupledFinite(m_defect), m_communicator);
            report.verifiedTrueResidualNorm = finite ? detail::CoupledNorm(m_defect, m_communicator) :
                std::numeric_limits<double>::quiet_NaN();
            report.correctionNorm = detail::CoupledNorm(correction, m_communicator);
            report.verificationSeconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
            report.verifiedGate = std::isfinite(report.convergenceThreshold) && report.convergenceThreshold >= 0.0 &&
                                  std::isfinite(report.verifiedTrueResidualNorm) &&
                                  report.verifiedTrueResidualNorm <= report.convergenceThreshold;
            const bool solved = detail::CoupledAll(report.linearSolve.Converged(), m_communicator);
            if (!finite || !std::isfinite(report.verifiedTrueResidualNorm) || !std::isfinite(report.correctionNorm))
                report.status = CoupledTrialCorrectionStatus::non_finite;
            else if (!solved || !report.verifiedGate) report.status = CoupledTrialCorrectionStatus::linear_failure;
            else if (report.correctionNorm > report.maximumCorrectionNorm) report.status = CoupledTrialCorrectionStatus::norm_limit;
            else report.status = CoupledTrialCorrectionStatus::admissible;
            return finish();
        }

    private:
        void RequireAlpha(double alpha) const {
            detail::CoupledRequire(std::isfinite(alpha) && alpha >= 0.0, m_communicator,
                                   "Coupled correction requires a finite nonnegative step length.");
            double minimum = 0.0, maximum = 0.0;
            if (MPI_Allreduce(&alpha, &minimum, 1, MPI_DOUBLE, MPI_MIN, m_communicator) != MPI_SUCCESS ||
                MPI_Allreduce(&alpha, &maximum, 1, MPI_DOUBLE, MPI_MAX, m_communicator) != MPI_SUCCESS)
                throw std::runtime_error("A coupled correction could not compare step lengths.");
            if (minimum != maximum) throw std::invalid_argument("Coupled correction step lengths must agree across ranks.");
        }
        const Operation *m_operation;
        Inverse *m_inverse;
        MPI_Comm m_communicator;
        CoupledTrialCorrectionOptions m_options;
        PreparedLinearBackendType<Backend, Operation, Inverse> m_backend;
        mfem::Vector m_baseResidual, m_baseAction, m_rightHandSide, m_verifiedAction, m_defect;
        double m_directionNorm{0.0};
        std::uint64_t m_baseGeneration{0};
        bool m_hasBase{false};
        std::optional<CoupledTrialCorrectionReport> m_lastReport{};
    };
}

// The state protocol is deliberately structural: production and focused probes
// execute the same transaction while their typed callbacks choose the metric
// and enthalpy policy. This helper never reads or writes accepted-state buffers.
export namespace mean_field::solver::detail {
    template <typename State> class PreparedCoupledTrialTransaction final {
        using Operation = std::remove_cvref_t<decltype(*std::declval<State>().normalizedOperator)>;
        using Inverse = std::remove_cvref_t<decltype(*std::declval<State>().normalizedInverse)>;
        using Clock = std::chrono::steady_clock;
    public:
        PreparedCoupledTrialTransaction(State &state, const CoupledTrialCorrectionOptions &options,
                                        int maximumAttemptsPerIteration = 1)
            : m_core(*state.normalizedOperator, *state.normalizedInverse, state.Problem().GetCommunicator(), options),
              m_maximumAttempts(maximumAttemptsPerIteration), m_originalState(state.normalizedOperator->Width()),
              m_originalResidual(state.normalizedOperator->Height()), m_originalPhysical(state.normalizedOperator->Width()),
              m_correction(state.normalizedOperator->Width()) {
            CoupledRequire(maximumAttemptsPerIteration > 0, m_core.GetCommunicator(), "A coupled trial attempt budget must be positive.");
            int minimum = 0, maximum = 0;
            if (MPI_Allreduce(&maximumAttemptsPerIteration, &minimum, 1, MPI_INT, MPI_MIN, m_core.GetCommunicator()) != MPI_SUCCESS ||
                MPI_Allreduce(&maximumAttemptsPerIteration, &maximum, 1, MPI_INT, MPI_MAX, m_core.GetCommunicator()) != MPI_SUCCESS)
                throw std::runtime_error("Coupled trial attempt budgets could not be compared.");
            if (minimum != maximum) throw std::invalid_argument("Coupled trial attempt budgets must agree across ranks.");
        }
        void CaptureBase(const mfem::Vector &residual, const mfem::Vector &direction) {
            m_core.CaptureBase(residual, direction);
            m_attemptsThisIteration = 0;
        }
        void BeginTrial() noexcept {
            m_diagnostics.status = CoupledTrialAttemptStatus::not_attempted;
            m_diagnostics.currentSolve.reset();
            m_diagnostics.originalMerit = m_diagnostics.correctedMerit = std::numeric_limits<double>::quiet_NaN();
        }
        [[nodiscard]] const CoupledTrialDiagnostics &GetDiagnostics() const noexcept { return m_diagnostics; }
        [[nodiscard]] const auto &GetCore() const noexcept { return m_core; }
        [[nodiscard]] const mfem::Vector &GetCorrection() const noexcept { return m_correction; }

        template <typename Preparation, typename Metric, typename Evaluate, typename Accepts, typename Better, typename CheckState>
        void TryImprove(State &state, double alpha, Preparation &preparation, Metric &metric,
                        Evaluate &&evaluate, Accepts &&accepts, Better &&better, CheckState &&checkState) {
            const auto communicator = m_core.GetCommunicator();
            CoupledRequire(preparation.has_value() && state.normalizedOperator->IsPrepared() &&
                CoupledFinite(state.trialNormalizedState) && CoupledFinite(state.trialNormalizedResidual) &&
                std::isfinite(metric.merit), communicator,
                "Coupled correction requires a successfully prepared original trial.");
            m_diagnostics.originalMerit = metric.merit;
            if (CoupledAll(accepts(metric), communicator)) {
                m_diagnostics.status = CoupledTrialAttemptStatus::original_accepted;
                return;
            }
            if (m_attemptsThisIteration >= m_maximumAttempts) {
                m_diagnostics.status = CoupledTrialAttemptStatus::budget_exhausted;
                return;
            }
            m_originalState = state.trialNormalizedState;
            m_originalResidual = state.trialNormalizedResidual;
            m_originalPhysical = state.candidatePhysicalState;
            Refresh(state);
            ++m_attemptsThisIteration;
            ++m_diagnostics.attempts;
            m_diagnostics.currentSolve = m_core.SolveCorrection(alpha, m_originalResidual, m_correction);
            const auto &report = *m_diagnostics.currentSolve;
            m_diagnostics.iterations += static_cast<std::uint64_t>(report.linearSolve.iterations);
            m_diagnostics.solveSeconds += report.linearSolve.solveSeconds;
            m_diagnostics.verificationSeconds += report.verificationSeconds;
            if (!report.Admissible()) {
                ++m_diagnostics.fallbacks;
                switch (report.status) {
                case CoupledTrialCorrectionStatus::linear_failure: m_diagnostics.status = CoupledTrialAttemptStatus::linear_failure; break;
                case CoupledTrialCorrectionStatus::norm_limit: m_diagnostics.status = CoupledTrialAttemptStatus::norm_limit; break;
                default: m_diagnostics.status = CoupledTrialAttemptStatus::non_finite; break;
                }
                // No candidate mutation has occurred; its prepared operator and
                // freshly refreshed inverse still describe the original trial.
                return;
            }
            const Metric originalMetric = metric;
            bool restorationStarted = false;
            auto restore = [&]() {
                restorationStarted = true;
                state.trialNormalizedState = m_originalState;
                state.candidatePhysicalState = m_originalPhysical;
                auto restored = Prepare(state);
                ++m_diagnostics.restorations;
                if (!CoupledAll(restored.has_value() && CoupledFinite(state.trialNormalizedResidual), communicator))
                    throw std::runtime_error("A coupled trial fallback could not reprepare its original candidate.");
                preparation = std::move(restored);
                Refresh(state);
                metric = evaluate(state.trialNormalizedResidual);
                ++m_diagnostics.fallbacks;
            };
            try {
                state.trialNormalizedState += m_correction;
                state.normalizedOperator->DenormalizeState(state.trialNormalizedState, state.candidatePhysicalState);
                if (!CoupledAll(CoupledFinite(state.trialNormalizedState) && CoupledFinite(state.candidatePhysicalState), communicator)) {
                    m_diagnostics.status = CoupledTrialAttemptStatus::non_finite;
                    restore();
                    return;
                }
                if (!CoupledAll(checkState(state.trialNormalizedState), communicator)) {
                    m_diagnostics.status = CoupledTrialAttemptStatus::inadmissible;
                    restore();
                    return;
                }
                auto corrected = Prepare(state);
                if (!CoupledAll(corrected.has_value(), communicator)) {
                    m_diagnostics.status = CoupledTrialAttemptStatus::inadmissible;
                    restore();
                    return;
                }
                if (!CoupledAll(CoupledFinite(state.trialNormalizedResidual), communicator)) {
                    m_diagnostics.status = CoupledTrialAttemptStatus::non_finite;
                    restore();
                    return;
                }
                const auto correctedMetric = evaluate(state.trialNormalizedResidual);
                m_diagnostics.correctedMerit = correctedMetric.merit;
                if (!CoupledAll(better(correctedMetric, originalMetric), communicator)) {
                    m_diagnostics.status = CoupledTrialAttemptStatus::not_improved;
                    restore();
                    return;
                }
                preparation = std::move(corrected);
                metric = correctedMetric;
                m_diagnostics.status = CoupledTrialAttemptStatus::kept;
                ++m_diagnostics.kept;
                // This is a nonlinear candidate, not an affine gravity completion.
                // Its fresh full residual is judged by the caller's normal metric.
                // As for an ordinary prepared trial, the caller refreshes the
                // inverse only if it commits or needs another auxiliary solve.
            } catch (...) {
                // Numerical rejection is handled above. For collectively
                // propagated exceptions, restore the original prepared trial
                // before rethrowing. A failed restoration itself propagates to
                // the owning Newton context, which restores the accepted base.
                if (!restorationStarted) restore();
                throw;
            }
        }

    private:
        void Refresh(State &state) {
            const auto start = Clock::now();
            state.physicalInverse->Refresh();
            m_diagnostics.refreshSeconds += std::chrono::duration<double>(Clock::now() - start).count();
        }
        auto Prepare(State &state) {
            const auto start = Clock::now();
            auto result = state.PrepareTrial();
            m_diagnostics.preparationSeconds += std::chrono::duration<double>(Clock::now() - start).count();
            ++m_diagnostics.preparations;
            return result;
        }
        PreparedCoupledTrialCorrection<Operation, Inverse> m_core;
        int m_maximumAttempts, m_attemptsThisIteration{0};
        mfem::Vector m_originalState, m_originalResidual, m_originalPhysical, m_correction;
        CoupledTrialDiagnostics m_diagnostics{};
    };
}