#pragma once

#include "plateau_coupled_trial_diagnostics.hpp"

namespace experiment::plateau_diagnostics {
struct DampedCoupledProbeOptions final {
    CoupledTrialProbeOptions common{};
    // The uncorrected control still captures the same tangent so its measured
    // search cost is comparable without hiding the optional feature's setup.
    std::string mode{"near"};
    mean_field::solver::BacktrackedCoupledStep search{};
    mean_field::solver::NearAcceptanceCoupledAttempt attempt{};
    int maximumTrials{13};
};

template <class State, class Search, class Attempt>
void ProbeDampedCoupledPath(State &state, const mfem::Vector &base, const mfem::Vector &direction,
                           const std::filesystem::path &output, DampedCoupledProbeOptions options,
                           Search search, Attempt attempt) {
    using namespace mean_field;
    namespace d = nonlinear_step_detail;
    using Gravity = coupled_trial_detail::GravityWorkspace<State>;
    using GravityCorrection = solver::PreparedGravityCurvatureCorrection<typename Gravity::Subsystem>;
    using Transaction = solver::detail::PreparedCoupledTrialTransaction<State, Search, Attempt>;
    const auto comm = state.Problem().GetCommunicator();
    const int size = state.Problem().StateSize();
    int ranks = 0;
    MPI_Comm_size(comm, &ranks);
    if (ranks != 1 || size != state.Problem().EquationSize() || base.Size() != size || direction.Size() != size ||
        !d::Finite(base) || !d::Finite(direction) || options.maximumTrials < 1 || options.maximumTrials > 30 ||
        !std::isfinite(options.common.outerLinearTolerance) || options.common.outerLinearTolerance <= 0. ||
        options.common.outerLinearTolerance >= 1. || !std::isfinite(options.common.sufficientDecrease) ||
        options.common.sufficientDecrease <= 0. || options.common.sufficientDecrease >= 1.)
        throw std::invalid_argument("Sequential coupled probes require one rank and compatible finite vectors.");
    const auto file = [&](const std::string &name) { return output / ("damped_coupled_" + name); };
    auto summary = d::Open(file("summary.csv"));
    auto trials = d::Open(file("trials.csv"));
    auto candidates = d::Open(file("candidates.csv"));
    auto evaluations = d::Open(file("evaluations.csv"));
    auto layout = d::Open(file("layout.csv"));
    layout << "kind,block,offset,size\n";
    for (const auto &v : state.Problem().GetManifest().valueBlocks())
        layout << "value," << v.stableId << ',' << v.offset << ',' << v.size << '\n';
    for (const auto &v : state.Problem().GetManifest().residualBlocks())
        layout << "residual," << v.stableId << ',' << v.offset << ',' << v.size << '\n';
    layout.flush();
    summary << "metric,value\n";
    trials << "trial,alpha,raw_prepared,raw_reason,raw_stage,gravity_gate,usable,before_norm,before_merit,"
              "before_accepted,status,attempted,attempts_total,aux_iterations,aux_gate,aux_threshold,aux_defect,"
              "correction_norm,correction_cap,selected_beta,candidate_count,predicted_decrease,excess_ratio,"
              "selected_prepared,selected_norm,selected_merit,selected_accepted,selected_original,"
              "raw_seconds,gravity_seconds,correction_seconds,preparations,restorations,solve_seconds,"
              "refresh_seconds,verification_seconds,preparation_seconds,aux_core_status,aux_status,aux_reported_gate,"
              "aux_rhs_norm,aux_reported_true_norm,aux_verification_applied,rhs_saved,correction_saved,verified_vectors_saved,"
              "gravity_target_defect,gravity_target_threshold,gravity_affine_disagreement\n";
    candidates << "trial,candidate,beta,status,merit,accepted,preparation_seconds\n";
    evaluations << "trial,evaluation,beta,residual_norm,merit,accepted\n";
    const auto metric = [&](const char *key, auto value) { summary << key << ',' << value << '\n'; };
    const double nan = std::numeric_limits<double>::quiet_NaN();
    const auto &normalization = state.normalizedOperator->GetNormalization();
    const mfem::Vector accepted(state.acceptedNormalizedState), acceptedF(state.acceptedNormalizedResidual),
        acceptedPhysical(state.AcceptedPhysicalState()), savedDirection(state.normalizedCorrection),
        savedPhysicalDirection(state.physicalCorrection), savedVolumeDirection(state.volumeDisplacementDirection),
        stateFactors(normalization.StateFactors()), residualFactors(normalization.ResidualFactors());
    const double acceptedJacobian = state.acceptedMinimumJacobianDeterminant;
    const auto evaluate = [&](const mfem::Vector &v) {
        return solver::nonlinear::getMetric(solver::nonlinear::NormalizedResidualMetric{}, v, comm);
    };
    const auto prepare = [&](const mfem::Vector &x) {
        state.trialNormalizedState = x;
        state.normalizedOperator->DenormalizeState(x, state.candidatePhysicalState);
        return state.PrepareTrial();
    };
    mfem::Vector baseF(size), action(size);
    bool captured = false, finished = false;
    const auto finish = [&]() {
        finished = true;
        if (!prepare(base)) throw std::runtime_error("Sequential coupled probe could not restore its base.");
        state.physicalInverse->Refresh();
        mfem::Vector restoredAction(size);
        state.normalizedOperator->Mult(direction, restoredAction);
        const bool acceptedSame = d::Identical(accepted, state.acceptedNormalizedState) &&
            d::Identical(acceptedF, state.acceptedNormalizedResidual) &&
            d::Identical(acceptedPhysical, state.AcceptedPhysicalState()) &&
            acceptedJacobian == state.acceptedMinimumJacobianDeterminant;
        const bool directionsSame = d::Identical(savedDirection, state.normalizedCorrection) &&
            d::Identical(savedPhysicalDirection, state.physicalCorrection) &&
            d::Identical(savedVolumeDirection, state.volumeDisplacementDirection);
        const bool factorsSame = d::Identical(stateFactors, normalization.StateFactors()) &&
            d::Identical(residualFactors, normalization.ResidualFactors());
        metric("accepted_buffers_bitwise_identical", acceptedSame);
        metric("direction_buffers_bitwise_identical", directionsSame);
        metric("normalization_factors_bitwise_identical", factorsSame);
        metric("restored_residual_error", captured ? d::Difference(baseF, state.trialNormalizedResidual) : nan);
        metric("restored_action_error", captured ? d::Difference(action, restoredAction) : nan);
        metric("fresh_operator_and_inverse_current", state.normalizedOperator->IsPrepared() &&
            state.normalizedInverse->IsCurrent() && state.physicalInverse->IsCurrent());
        d::Save(file("restored_residual.txt"), state.trialNormalizedResidual);
        d::Save(file("restored_action.txt"), restoredAction);
        summary.flush();
        if (!acceptedSame || !directionsSame || !factorsSame)
            throw std::runtime_error("Sequential coupled probe mutated an accepted buffer or normalization.");
    };
    metric("candidate_status_evaluated", int(solver::CoupledTrialCandidateStatus::evaluated));
    metric("candidate_status_inadmissible", int(solver::CoupledTrialCandidateStatus::inadmissible));
    metric("candidate_status_non_finite", int(solver::CoupledTrialCandidateStatus::non_finite));
    metric("schema_version", 1);
    metric("mode", options.mode);
    metric("gravity_path", options.common.gravityPreconditionPath);
    metric("attempt_budget", 1);
    metric("maximum_trials", options.maximumTrials);
    metric("maximum_iterations", options.common.coupled.control.maximumIterations);
    metric("restart", options.common.coupled.backend.restartLength);
    metric("maximum_beta_steps", std::same_as<Search, solver::FullCoupledStep> ? 1 : options.search.maximumSteps);
    metric("beta_reduction", options.search.reduction);
    metric("maximum_excess_ratio", options.attempt.maximumExcessRatio);
    metric("sufficient_decrease", options.common.sufficientDecrease);
    metric("relative_tolerance", options.common.coupled.control.relativeTolerance);
    metric("absolute_tolerance", options.common.coupled.control.absoluteTolerance);
    metric("maximum_correction_ratio", options.common.coupled.maximumCorrectionRatio);
    metric("accepted_commits", 0);
    metric("shared_iteration_budget", 1);
    metric("status_prediction_deferred", int(solver::CoupledTrialAttemptStatus::prediction_deferred));
    metric("status_original_accepted", int(solver::CoupledTrialAttemptStatus::original_accepted));
    metric("status_budget_exhausted", int(solver::CoupledTrialAttemptStatus::budget_exhausted));
    metric("status_kept", int(solver::CoupledTrialAttemptStatus::kept));
    try {
        if (!prepare(base)) throw std::runtime_error("Sequential coupled probe base does not prepare.");
        state.physicalInverse->Refresh();
        baseF = state.trialNormalizedResidual;
        state.normalizedOperator->Mult(direction, action);
        captured = true;
        const auto baseMetric = evaluate(baseF);
        mfem::Vector defect(action);
        defect += baseF;
        const double defectNorm = d::Norm(defect);
        const bool outerGate = d::Finite(baseF) && d::Finite(action) && d::Finite(defect) &&
            std::isfinite(baseMetric.residualNorm) && baseMetric.residualNorm > 0. &&
            std::isfinite(baseMetric.merit) && baseMetric.merit >= 0. && std::isfinite(d::Norm(direction)) &&
            std::isfinite(d::Norm(action)) &&
            std::isfinite(defectNorm) && defectNorm <= options.common.outerLinearTolerance * baseMetric.residualNorm;
        metric("base_norm", baseMetric.residualNorm);
        metric("base_merit", baseMetric.merit);
        metric("outer_relative_defect", d::Norm(defect) / baseMetric.residualNorm);
        metric("outer_linear_tolerance", options.common.outerLinearTolerance);
        metric("outer_verified_gate", outerGate);
        d::Save(file("base.txt"), base);
        d::Save(file("base_physical.txt"), state.candidatePhysicalState);
        d::Save(file("base_residual.txt"), baseF);
        d::Save(file("direction.txt"), direction);
        d::Save(file("base_action.txt"), action);
        d::Save(file("state_factors.txt"), stateFactors);
        d::Save(file("residual_factors.txt"), residualFactors);
        if (!outerGate) throw std::runtime_error("Sequential probe supplied direction fails the full-J gate.");
        Transaction transaction(state, options.common.coupled, 1, search, attempt);
        const auto captureStart = d::Clock::now();
        transaction.CaptureBase(baseF, direction);
        std::optional<Gravity> gravity;
        std::optional<GravityCorrection> gravityCorrection;
        if (options.common.gravityPreconditionPath) {
            gravity.emplace(state, options.common.gravity);
            gravityCorrection.emplace(gravity->subsystem);
            gravity->subsystem.Refresh();
            gravityCorrection->CaptureBase(baseF, direction);
        }
        metric("capture_seconds", d::Seconds(captureStart));
        const auto searchStart = d::Clock::now();
        double operationSeconds = 0.;
        int acceptedTrial = -1;
        for (int trial = 0; trial < options.maximumTrials; ++trial) {
            const double alpha = std::ldexp(1., -trial);
            const std::string suffix = "_a" + std::to_string(trial) + ".txt";
            transaction.BeginTrial();
            if (gravity) gravity->BeginGravityAttempt();
            mfem::Vector raw(base), prediction(baseF);
            raw.Add(alpha, direction);
            prediction.Add(alpha, action);
            d::Save(file("raw" + suffix), raw);
            d::Save(file("prediction" + suffix), prediction);
            const auto rawStart = d::Clock::now();
            auto preparation = prepare(raw);
            const double rawSeconds = d::Seconds(rawStart);
            const bool rawPrepared = bool(preparation);
            if (rawPrepared) d::Save(file("raw_residual" + suffix), state.trialNormalizedResidual);
            const int rawReason = rawPrepared ? -1 : int(preparation.error().reason);
            const int rawStage = rawPrepared ? -1 : int(preparation.error().stage);
            bool gravityGate = !gravity;
            double gravitySeconds = 0.;
            if (rawPrepared && gravity) {
                const auto started = d::Clock::now();
                gravity->Refresh(state);
                gravityCorrection->BuildLinearPrediction(alpha, gravity->desiredResidual);
                gravity->lastSolve = gravityCorrection->SolveCorrection(alpha, state.trialNormalizedResidual,
                    gravity->gravityCorrection, options.common.gravity.control);
                gravity->trialSolveSucceeded = gravity->lastSolve.Converged();
                if (gravity->trialSolveSucceeded) {
                    preparation = gravity->ApplyCorrection(state);
                    gravityGate = bool(preparation) && gravity->TrialSolveSucceeded();

                }
                gravitySeconds = d::Seconds(started);
                if (gravityGate) d::Save(file("gravity_correction" + suffix), gravity->fullCorrection);
            }
            // Production rejects an overflowing prepared residual before the
            // correction transaction, whose precondition requires a valid metric.
            const bool preparedFinite = preparation.has_value() && gravityGate && d::Finite(state.trialNormalizedResidual);
            const auto preparedMetric = preparedFinite ? evaluate(state.trialNormalizedResidual) :
                solver::nonlinear::MetricEvaluation{nan, nan};
            const bool usable = preparedFinite && std::isfinite(preparedMetric.residualNorm) &&
                preparedMetric.residualNorm >= 0. && std::isfinite(preparedMetric.merit) && preparedMetric.merit >= 0.;
            const mfem::Vector before(state.trialNormalizedState);
            auto beforeMetric = solver::nonlinear::MetricEvaluation{nan, nan};
            if (usable) {
                beforeMetric = preparedMetric;
                d::Save(file("before" + suffix), before);
                d::Save(file("before_residual" + suffix), state.trialNormalizedResidual);
            }
            const auto accepts = [&](const auto &m) {
                return std::isfinite(m.merit) && m.merit >= 0. &&
                    m.merit <= (1. - options.common.sufficientDecrease * alpha) * baseMetric.merit;
            };
            const auto better = [](const auto &a, const auto &b) {
                return std::isfinite(a.merit) && a.merit >= 0. && std::isfinite(a.residualNorm) &&
                    a.residualNorm >= 0. && a.merit < b.merit;
            };
            auto selected = beforeMetric;
            int observed = 0;
            struct Observation { mfem::Vector x, f; double beta; solver::nonlinear::MetricEvaluation metric; };
            std::vector<Observation> observations;
            const auto observe = [&](const mfem::Vector &v) {
                const auto m = evaluate(v);
                // Prediction metrics use other vectors. Copy observations in
                // memory here; writing large artifacts must not inflate the
                // timed correction's preparation/selection cost.
                if (&v == &state.trialNormalizedResidual) {
                    mfem::Vector delta(state.trialNormalizedState);
                    delta -= before;
                    const auto &c = transaction.GetCorrection();
                    const double c2 = c * c;
                    observations.push_back({state.trialNormalizedState, v, c2 > 0. ? (delta * c) / c2 : 0., m});
                }
                return m;
            };
            const auto previous = transaction.GetDiagnostics();
            double correctionSeconds = 0.;
            if (usable && options.mode != "none") {
                const auto started = d::Clock::now();
                transaction.TryImprove(state, alpha, preparation, selected, observe, accepts, better,
                    [](const mfem::Vector &) { return true; });
                correctionSeconds = d::Seconds(started);
            }
            operationSeconds += rawSeconds + gravitySeconds + correctionSeconds;
            const auto &diag = transaction.GetDiagnostics();
            const auto &report = diag.currentSolve;
            const bool rhsSaved = report && d::Finite(transaction.GetCore().GetRightHandSide());
            const bool correctionSaved = report && d::Finite(transaction.GetCorrection());
            const bool verifiedSaved = report && report->verificationApplied &&
                d::Finite(transaction.GetCore().GetVerifiedAction()) && d::Finite(transaction.GetCore().GetVerifiedDefect());
            if (report) {
                if (rhsSaved) d::Save(file("rhs" + suffix), transaction.GetCore().GetRightHandSide());
                if (correctionSaved) d::Save(file("correction" + suffix), transaction.GetCorrection());
                if (verifiedSaved) {
                    d::Save(file("action" + suffix), transaction.GetCore().GetVerifiedAction());
                    d::Save(file("defect" + suffix), transaction.GetCore().GetVerifiedDefect());
                }
            }
            for (const auto &o : observations) {
                const auto tag = "_a" + std::to_string(trial) + "_e" + std::to_string(observed) + ".txt";
                d::Save(file("evaluated_state" + tag), o.x);
                d::Save(file("evaluated_residual" + tag), o.f);
                evaluations << trial << ',' << observed++ << ',' << o.beta << ',' << o.metric.residualNorm
                            << ',' << o.metric.merit << ',' << accepts(o.metric) << '\n';
            }
            int candidate = 0;
            for (const auto &c : diag.candidates)
                candidates << trial << ',' << candidate++ << ',' << c.beta << ',' << int(c.status) << ','
                           << c.merit << ',' << c.accepted << ',' << c.preparationSeconds << '\n';
            if (usable && preparation) {
                selected = evaluate(state.trialNormalizedResidual);
                d::Save(file("selected" + suffix), state.trialNormalizedState);
                d::Save(file("selected_residual" + suffix), state.trialNormalizedResidual);
            }
            const bool acceptedNow = usable && preparation && accepts(selected);
            trials << trial << ',' << alpha << ',' << rawPrepared << ',' << rawReason << ',' << rawStage << ','
                   << gravityGate << ',' << usable << ',' << beforeMetric.residualNorm << ',' << beforeMetric.merit << ','
                   << (usable && accepts(beforeMetric)) << ',' << int(diag.status) << ',' << bool(report) << ','
                   << diag.attempts << ',' << (report ? report->linearSolve.iterations : 0) << ','
                   << (report && report->verifiedGate) << ',' << (report ? report->convergenceThreshold : nan) << ','
                   << (report ? report->verifiedTrueResidualNorm : nan) << ',' << (report ? report->correctionNorm : nan) << ','
                   << (report ? report->maximumCorrectionNorm : nan) << ',' << diag.selectedBeta << ',' << diag.candidateCount << ','
                   << diag.predictedDecrease << ',' << diag.predictionExcessRatio << ',' << (usable && bool(preparation)) << ','
                   << selected.residualNorm << ',' << selected.merit << ',' << acceptedNow << ','
                   << (usable && d::Identical(before, state.trialNormalizedState)) << ',' << rawSeconds << ','
                   << gravitySeconds << ',' << correctionSeconds << ',' << diag.preparations - previous.preparations << ','
                   << diag.restorations - previous.restorations << ',' << diag.solveSeconds - previous.solveSeconds << ','
                   << diag.refreshSeconds - previous.refreshSeconds << ',' << diag.verificationSeconds - previous.verificationSeconds << ','
                   << diag.preparationSeconds - previous.preparationSeconds << ','
                   << (report ? int(report->status) : -1) << ',' << (report ? int(report->linearSolve.status) : -1) << ','
                   << (report && report->linearSolve.Converged()) << ',' << (report ? report->linearSolve.rightHandSideNorm : nan) << ','
                   << (report ? report->linearSolve.trueResidualNorm : nan) << ',' << (report && report->verificationApplied) << ','
                   << rhsSaved << ',' << correctionSaved << ',' << verifiedSaved << ','
                   << (gravity && rawPrepared ? gravity->lastPreparedDefect : nan) << ','
                   << (gravity && rawPrepared ? gravity->lastPreparedTargetThreshold : nan) << ','
                   << (gravity && rawPrepared ? gravity->lastAffineDisagreement : nan) << '\n';
            trials.flush(); candidates.flush(); evaluations.flush();
            if (acceptedNow) {
                acceptedTrial = trial;
                metric("accepted_alpha", alpha);
                metric("accepted_norm", selected.residualNorm);
                metric("accepted_beta", diag.selectedBeta);
                d::Save(output / ("coupled_trial_selected_physical" + suffix), state.candidatePhysicalState);
                break;
            }
        }
        metric("search_wall_seconds", d::Seconds(searchStart));
        metric("search_operation_seconds", operationSeconds);
        metric("accepted_trial", acceptedTrial);
        const auto &diag = transaction.GetDiagnostics();
        metric("attempts", diag.attempts);
        metric("aux_iterations", diag.iterations);
        metric("deferred_attempts", diag.deferredAttempts);
        metric("probe_exception", 0);
    } catch (...) {
        const auto error = std::current_exception();
        metric("probe_exception", 1);
        if (!finished) finish();
        std::rethrow_exception(error);
    }
    finish();
}

template <class State>
void ProbeDampedCoupled(State &state, const mfem::Vector &base, const mfem::Vector &direction,
                       const std::filesystem::path &output, DampedCoupledProbeOptions options) {
    using namespace mean_field::solver;
    if (options.mode == "none" || options.mode == "full")
        ProbeDampedCoupledPath(state, base, direction, output, options, FullCoupledStep{}, EveryRejectedCoupledAttempt{});
    else if (options.mode == "damped")
        ProbeDampedCoupledPath(state, base, direction, output, options, options.search, EveryRejectedCoupledAttempt{});
    else if (options.mode == "near")
        ProbeDampedCoupledPath(state, base, direction, output, options, options.search, options.attempt);
    else throw std::invalid_argument("Unknown sequential coupled probe mode.");
}
} // namespace experiment::plateau_diagnostics