#pragma once

// Include after import mean_field. The caller must hold the diagnostic state lease.
#include <algorithm>
#include <array>
#include <cmath>
#include <exception>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <limits>
#include <stdexcept>
#include <string>
#include <type_traits>

namespace experiment::plateau_diagnostics {
struct NumericalStrategyProbeOptions final {
    mean_field::solver::LinearSolveControl inner{1.e-9, 1.e-13, 160};
    int restartLength{80};
    int maximumCompletionCorrections{2};
    double gravityCompletionTolerance{1.e-11};
    // Coarser differences resolve truncation before cancellation dominates;
    // the final small step independently checks derivative accuracy.
    std::array<double, 4> reducedDifferenceSteps{.1, .03, .01, .001};
};

namespace numerical_strategy_detail {
inline bool Finite(const mfem::Vector &values) {
    for (int index = 0; index < values.Size(); ++index)
        if (!std::isfinite(values(index))) return false;
    return true;
}
inline double Difference(const mfem::Vector &left, const mfem::Vector &right) {
    mfem::Vector difference(left);
    difference -= right;
    return difference.Norml2();
}
inline std::ofstream Open(const std::filesystem::path &path) {
    if (std::filesystem::exists(path))
        throw std::invalid_argument("Numerical strategy diagnostic output already exists: " + path.string());
    std::ofstream output(path);
    output.exceptions(std::ios::failbit | std::ios::badbit);
    output << std::setprecision(17);
    return output;
}
} // namespace numerical_strategy_detail

template <class State>
void ProbeNumericalStrategies(
    State &state, const mean_field::fem::FEM &finiteElements,
    const mfem::Vector &normalizedBase, const mfem::Vector &normalizedDirection,
    const std::filesystem::path &output, NumericalStrategyProbeOptions options = {}
) {
    using namespace mean_field;
    namespace d = numerical_strategy_detail;
    using Form = typename std::remove_cvref_t<decltype(state.Problem())>::FormType;
    options.inner.Validate();
    if (options.restartLength <= 0 || options.maximumCompletionCorrections <= 0 ||
        !std::isfinite(options.gravityCompletionTolerance) || options.gravityCompletionTolerance <= 0.)
        throw std::invalid_argument("Numerical strategy probes require valid inner and completion controls.");
    double previousDifferenceStep = std::numeric_limits<double>::infinity();
    for (const double step : options.reducedDifferenceSteps) {
        if (!std::isfinite(step) || step <= 0. || step >= previousDifferenceStep)
            throw std::invalid_argument("Reduced finite-difference steps must be finite, positive, and strictly decreasing.");
        previousDifferenceStep = step;
    }
    int ranks = 0;
    MPI_Comm_size(state.Problem().GetCommunicator(), &ranks);
    if (ranks != 1 || !finiteElements.mesh || normalizedBase.Size() != state.Problem().StateSize() ||
        normalizedDirection.Size() != normalizedBase.Size() || !d::Finite(normalizedBase) ||
        !d::Finite(normalizedDirection))
        throw std::invalid_argument("Numerical strategy probes require one rank and finite compatible root vectors.");

    const mfem::Vector base(normalizedBase), direction(normalizedDirection);
    const mfem::Vector savedAcceptedState(state.acceptedNormalizedState);
    const mfem::Vector savedAcceptedResidual(state.acceptedNormalizedResidual);
    const mfem::Vector savedAcceptedPhysical(state.AcceptedPhysicalState());
    const mfem::Vector savedCorrection(state.normalizedCorrection), savedPhysicalCorrection(state.physicalCorrection);
    const mfem::Vector savedVolumeDirection(state.volumeDisplacementDirection);
    std::filesystem::create_directories(output);
    auto summary = d::Open(output / "numerical_strategy_summary.csv");
    auto preparations = d::Open(output / "numerical_strategy_preparations.csv");
    auto inner = d::Open(output / "numerical_strategy_inner.csv");
    auto curvature = d::Open(output / "numerical_strategy_curvature.csv");
    auto completions = d::Open(output / "numerical_strategy_completions.csv");
    auto differences = d::Open(output / "numerical_strategy_reduced_derivative.csv");
    summary << "metric,value\n";
    preparations << "sequence,kind,alpha,prepared,rejection_reason,rejection_stage,residual_norm\n";
    inner << "kind,alpha,index,status,converged,iterations,right_hand_side_norm,true_residual_norm,"
             "relative_true_residual_norm,threshold,operator_applications,inverse_applications,seconds\n";
    curvature << "alpha,prepared,inner_converged,corrected_prepared,base_residual_norm,raw_residual_norm,"
                 "corrected_residual_norm,linear_prediction_norm,raw_gravity_remainder,corrected_gravity_remainder,"
                 "gravity_correction_norm,correction_over_alpha_squared,retained_state_change\n";
    completions << "kind,alpha,completed,corrections,initial_gravity_norm,final_gravity_norm,full_residual_norm,"
                   "retained_state_change\n";
    differences << "alpha,plus_completed,minus_completed,schur_action_norm,difference_quotient_norm,"
                   "difference_quotient_error,relative_error,plus_gravity_norm,minus_gravity_norm,"
                   "residual_subtraction_roundoff_scale\n";
    summary << "inner_relative_tolerance," << options.inner.relativeTolerance
            << "\ninner_absolute_tolerance," << options.inner.absoluteTolerance
            << "\ninner_maximum_iterations," << options.inner.maximumIterations
            << "\ninner_restart_length," << options.restartLength
            << "\ngravity_completion_tolerance," << options.gravityCompletionTolerance << '\n';
    for (std::size_t index = 0; index < options.reducedDifferenceSteps.size(); ++index)
        summary << "reduced_difference_step_" << index << ',' << options.reducedDifferenceSteps[index] << '\n';
    const double nan = std::numeric_limits<double>::quiet_NaN();
    int preparationCount = 0;
    double restorationError = 0.;
    mfem::Vector originalResidual(state.Problem().EquationSize());
    bool haveOriginalResidual = false;

    auto prepare = [&](const mfem::Vector &candidate, const std::string &kind, double alpha) {
        state.trialNormalizedState = candidate;
        const auto result = state.PrepareTrial();
        ++preparationCount;
        preparations << preparationCount << ',' << kind << ',' << alpha << ',' << bool(result) << ',';
        if (result)
            preparations << "-1,-1," << state.trialNormalizedResidual.Norml2() << '\n';
        else
            preparations << int(result.error().reason) << ',' << int(result.error().stage) << ',' << nan << '\n';
        preparations.flush();
        if (result && !d::Finite(state.trialNormalizedResidual))
            throw std::runtime_error("A numerical strategy probe produced a nonfinite prepared residual.");
        return bool(result);
    };
    auto restore = [&]() {
        if (!prepare(base, "restore_original", 0.))
            throw std::runtime_error("A numerical strategy probe could not restore its frozen base.");
        if (!haveOriginalResidual) {
            originalResidual = state.trialNormalizedResidual;
            haveOriginalResidual = true;
        }
        restorationError = std::max(restorationError, d::Difference(originalResidual, state.trialNormalizedResidual));
        state.physicalInverse->Refresh();
    };
    auto finish = [&]() {
        restore();
        state.normalizedCorrection = savedCorrection;
        state.physicalCorrection = savedPhysicalCorrection;
        state.volumeDisplacementDirection = savedVolumeDirection;
        const double acceptedStateChange = d::Difference(savedAcceptedState, state.acceptedNormalizedState);
        const double acceptedResidualChange = d::Difference(savedAcceptedResidual, state.acceptedNormalizedResidual);
        const double acceptedPhysicalChange = d::Difference(savedAcceptedPhysical, state.AcceptedPhysicalState());
        summary << "maximum_original_restoration_error," << restorationError
                << "\naccepted_state_change," << acceptedStateChange
                << "\naccepted_residual_change," << acceptedResidualChange
                << "\naccepted_physical_change," << acceptedPhysicalChange
                << "\npreparations," << preparationCount << '\n';
        summary.flush();
        if (acceptedStateChange != 0. || acceptedResidualChange != 0. || acceptedPhysicalChange != 0.)
            throw std::runtime_error("A numerical strategy probe changed an accepted state.");
    };
    auto recordInner = [&](const std::string &kind, double alpha, int index, const solver::LinearSolveReport &report) {
        inner << kind << ',' << alpha << ',' << index << ',' << int(report.status) << ',' << report.Converged()
              << ',' << report.iterations << ',' << report.rightHandSideNorm << ',' << report.trueResidualNorm
              << ',' << report.relativeTrueResidualNorm << ','
              << (std::isfinite(report.rightHandSideNorm) && report.rightHandSideNorm >= 0.
                      ? options.inner.ConvergenceThreshold(report.rightHandSideNorm) : nan)
              << ',' << report.operatorApplications
              << ',' << report.inversePreconditionerApplications << ',' << report.solveSeconds << '\n';
        inner.flush();
    };

    try {
        restore();
        auto partition = solver::MakeGravityBlockPartition<Form>(state.Problem().GetManifest().layout());
        solver::PreparedGravitySubsystem gravity(*state.normalizedOperator, *state.normalizedInverse, partition,
            state.Problem().GetCommunicator(), solver::linear::FGMRES({options.restartLength, -1}));
        solver::PreparedGravityCurvatureCorrection correction(gravity);
        solver::GravityReducedOperator reduced(gravity, options.inner);
        mfem::Vector baseAction(base.Size()), baseGravity(partition.GravityResidualSize());
        mfem::Vector gravityAction(partition.GravityResidualSize()), gravityResidual(partition.GravityResidualSize());
        mfem::Vector compactCorrection(partition.GravityValueSize());
        mfem::Vector retainedBefore(partition.RetainedValueSize()), retainedAfter(partition.RetainedValueSize());
        state.normalizedOperator->Mult(direction, baseAction);
        partition.GatherGravityResiduals(originalResidual, baseGravity);
        partition.GatherGravityResiduals(baseAction, gravityAction);
        summary << "original_residual_norm," << originalResidual.Norml2()
                << "\noriginal_gravity_norm," << baseGravity.Norml2() << '\n';
        for (const double alpha : {1., .1, .01, .001}) {
            restore();
            gravity.Refresh();
            correction.CaptureBase(originalResidual, direction);
            mfem::Vector candidate(base);
            candidate.Add(alpha, direction);
            const bool rawPrepared = prepare(candidate, "curvature_raw", alpha);
            if (!rawPrepared) {
                curvature << alpha << ",0,0,0," << originalResidual.Norml2()
                          << ",nan,nan,nan,nan,nan,nan,nan,nan\n";
                curvature.flush();
                continue;
            }
            const mfem::Vector rawResidual(state.trialNormalizedResidual);
            mfem::Vector prediction(originalResidual);
            prediction.Add(alpha, baseAction);
            partition.GatherGravityResiduals(rawResidual, gravityResidual);
            gravityResidual -= baseGravity;
            gravityResidual.Add(-alpha, gravityAction);
            const double rawGravityRemainder = gravityResidual.Norml2();
            state.physicalInverse->Refresh();
            gravity.Refresh();
            const auto report = correction.SolveCorrection(alpha, rawResidual, compactCorrection, options.inner);
            recordInner("curvature", alpha, 0, report);
            double correctedNorm = nan, correctedRemainder = nan, retainedChange = nan;
            bool correctedPrepared = false;
            if (report.Converged()) {
                partition.GatherRetainedValues(candidate, retainedBefore);
                partition.AddGravityValues(compactCorrection, candidate);
                partition.GatherRetainedValues(candidate, retainedAfter);
                retainedChange = d::Difference(retainedBefore, retainedAfter);
                correctedPrepared = prepare(candidate, "curvature_corrected", alpha);
                if (correctedPrepared) {
                    correctedNorm = state.trialNormalizedResidual.Norml2();
                    partition.GatherGravityResiduals(state.trialNormalizedResidual, gravityResidual);
                    gravityResidual -= baseGravity;
                    gravityResidual.Add(-alpha, gravityAction);
                    correctedRemainder = gravityResidual.Norml2();
                }
            }
            curvature << alpha << ",1," << report.Converged() << ',' << correctedPrepared << ','
                      << originalResidual.Norml2() << ',' << rawResidual.Norml2() << ',' << correctedNorm << ','
                      << prediction.Norml2() << ',' << rawGravityRemainder << ',' << correctedRemainder << ','
                      << compactCorrection.Norml2() << ',' << compactCorrection.Norml2() / (alpha * alpha) << ','
                      << retainedChange << '\n';
            curvature.flush();
        }

        // The residual used by a reduced nonlinear method is evaluated after
        // gravity completion, so finite differences must follow that same manifold.
        auto complete = [&](mfem::Vector &candidate, const std::string &kind, double alpha) {
            partition.GatherRetainedValues(candidate, retainedBefore);
            double initialNorm = nan, finalNorm = nan;
            int completedCorrections = 0;
            bool completed = false;
            if (prepare(candidate, kind + "_initial", alpha)) {
                for (int attempt = 0; attempt <= options.maximumCompletionCorrections; ++attempt) {
                    partition.GatherGravityResiduals(state.trialNormalizedResidual, gravityResidual);
                    finalNorm = gravityResidual.Norml2();
                    if (attempt == 0) initialNorm = finalNorm;
                    if (finalNorm <= options.gravityCompletionTolerance) {
                        completed = true;
                        break;
                    }
                    if (attempt == options.maximumCompletionCorrections) break;
                    state.physicalInverse->Refresh();
                    gravity.Refresh();
                    gravityResidual *= -1.;
                    compactCorrection = 0.;
                    const auto report = gravity.Solve(gravityResidual, compactCorrection, options.inner);
                    recordInner(kind, alpha, attempt, report);
                    if (!report.Converged()) break;
                    partition.AddGravityValues(compactCorrection, candidate);
                    ++completedCorrections;
                    if (!prepare(candidate, kind + "_corrected", alpha)) {
                        finalNorm = nan;
                        break;
                    }
                }
            }
            partition.GatherRetainedValues(candidate, retainedAfter);
            const double retainedChange = d::Difference(retainedBefore, retainedAfter);
            completions << kind << ',' << alpha << ',' << completed << ',' << completedCorrections << ','
                        << initialNorm << ',' << finalNorm << ','
                        << (completed ? state.trialNormalizedResidual.Norml2() : nan) << ',' << retainedChange << '\n';
            completions.flush();
            if (retainedChange != 0.)
                throw std::runtime_error("Gravity completion changed retained coordinates.");
            return completed;
        };

        restore();
        mfem::Vector completedBase(base);
        const bool baseCompleted = complete(completedBase, "reduced_base", 0.);
        summary << "reduced_base_completed," << baseCompleted << '\n';
        if (baseCompleted) {
            state.physicalInverse->Refresh();
            gravity.Refresh();
            const mfem::Vector completedBaseResidual(state.trialNormalizedResidual);
            mfem::Vector retainedDirection(partition.RetainedValueSize()), reducedAction(partition.RetainedResidualSize());
            mfem::Vector liftedDirection(base.Size()), liftedAction(base.Size());
            partition.GatherRetainedValues(direction, retainedDirection);
            const auto actionReport = reduced.TryMult(retainedDirection, reducedAction);
            recordInner("reduced_action", 0., 0, actionReport);
            const auto liftReport = reduced.TryLiftTangent(retainedDirection, liftedDirection);
            recordInner("reduced_lift", 0., 0, liftReport);
            if (actionReport.Converged() && liftReport.Converged()) {
                state.normalizedOperator->Mult(liftedDirection, liftedAction);
                partition.GatherGravityResiduals(liftedAction, gravityResidual);
                mfem::Vector retainedAction(partition.RetainedResidualSize());
                partition.GatherRetainedResiduals(liftedAction, retainedAction);
                summary << "reduced_action_norm," << reducedAction.Norml2()
                        << "\nlifted_gravity_action_norm," << gravityResidual.Norml2()
                        << "\nlifted_retained_action_difference," << d::Difference(retainedAction, reducedAction)
                        << "\ncompleted_base_residual_norm," << completedBaseResidual.Norml2() << '\n';
                for (const double alpha : options.reducedDifferenceSteps) {
                    std::array<mfem::Vector, 2> residuals{mfem::Vector(base.Size()), mfem::Vector(base.Size())};
                    std::array<bool, 2> completed{};
                    std::array<double, 2> gravityNorms{nan, nan};
                    for (int side = 0; side < 2; ++side) {
                        mfem::Vector candidate(completedBase);
                        candidate.Add(side == 0 ? alpha : -alpha, liftedDirection);
                        completed[side] = complete(candidate, side == 0 ? "reduced_plus" : "reduced_minus", alpha);
                        if (completed[side]) {
                            residuals[side] = state.trialNormalizedResidual;
                            partition.GatherGravityResiduals(residuals[side], gravityResidual);
                            gravityNorms[side] = gravityResidual.Norml2();
                        }
                    }
                    double quotientNorm = nan, error = nan, relativeError = nan, roundoffScale = nan;
                    if (completed[0] && completed[1]) {
                        mfem::Vector quotient(residuals[0]);
                        quotient -= residuals[1];
                        quotient /= 2. * alpha;
                        partition.GatherRetainedResiduals(quotient, retainedAction);
                        quotientNorm = retainedAction.Norml2();
                        error = d::Difference(retainedAction, reducedAction);
                        relativeError = error / std::max(reducedAction.Norml2(), std::numeric_limits<double>::min());
                        roundoffScale = 64. * std::numeric_limits<double>::epsilon() *
                            (1. + residuals[0].Norml2() + residuals[1].Norml2()) / (2. * alpha);
                    }
                    differences << alpha << ',' << completed[0] << ',' << completed[1] << ',' << reducedAction.Norml2()
                                << ',' << quotientNorm << ',' << error << ',' << relativeError << ',' << gravityNorms[0]
                                << ',' << gravityNorms[1] << ',' << roundoffScale << '\n';
                    differences.flush();
                }
            }
        }
        finish();
    } catch (...) {
        const auto failure = std::current_exception();
        finish();
        std::rethrow_exception(failure);
    }
}
} // namespace experiment::plateau_diagnostics