#include <algorithm>
#include <chrono>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <mfem.hpp>
#include <numbers>
#include <span>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
import mean_field;
#include "plateau_enthalpy_step_diagnostics.hpp"
#include "plateau_physical_diagnostics.hpp"
#include "plateau_preconditioner_diagnostics.hpp"
#include "plateau_preconditioner_benchmark.hpp"
#include "plateau_newton_preconditioner.hpp"
#include "plateau_pressure_obstruction_diagnostics.hpp"
#include "plateau_shape_curvature_diagnostics.hpp"
#include "plateau_numerical_strategy_diagnostics.hpp"
#include "plateau_gravity_application_diagnostics.hpp"
#include "plateau_positivity_strategy_diagnostics.hpp"
#include "plateau_polynomial_certificate.hpp"

using namespace mean_field;
namespace {
using Clock = std::chrono::steady_clock;
struct Options {
    double n = 1., angularMomentum = 0., atol = 0., rtol = 1e-8, linearTolerance = .03;
    int maxNewton = 12, maxLinear = 200, restart = 40, probeRestart = 40;
    bool probes = false, noCentral = false;
    bool strategyProbe = false, positivityProbe = false, positivityCertificate = false;
    bool preconditionerProbe = false;
    std::string preconditionerVariant = "independent";
    int schurCorrectionSweeps = 1;
    double schurRelaxation = 0.5;
    std::string strategy = "baseline";
    double gravityTolerance = 1e-10;
    int gravityIterations = 300, gravityRestart = 80;
    bool pressureProbe = false, shapeCurvature = false, surfaceTrace = false;
    std::vector<int> captureIterations;
    std::string mesh = "sandbox.smesh";
    std::filesystem::path output, probeState, probeDirection, replayCandidate;
};
Options Parse(int argc, char **argv) {
    Options o;
    for (int i = 1; i < argc; ++i) {
        std::string a = argv[i];
        auto v = [&]() {
            if (++i >= argc)
                throw std::invalid_argument("Missing value");
            return std::string(argv[i]);
        };
        if (a == "--index")
            o.n = std::stod(v());
        else if (a == "--angular-momentum")
            o.angularMomentum = std::stod(v());
        else if (a == "--output")
            o.output = v();
        else if (a == "--mesh")
            o.mesh = v();
        else if (a == "--max-newton")
            o.maxNewton = std::stoi(v());
        else if (a == "--max-linear")
            o.maxLinear = std::stoi(v());
        else if (a == "--restart")
            o.restart = std::stoi(v());
        else if (a == "--probe-restart")
            o.probeRestart = std::stoi(v());
        else if (a == "--linear-tol")
            o.linearTolerance = std::stod(v());
        else if (a == "--atol")
            o.atol = std::stod(v());
        else if (a == "--rtol")
            o.rtol = std::stod(v());
        else if (a == "--probe-state")
            o.probeState = v();
        else if (a == "--probe-direction")
            o.probeDirection = v();
        else if (a == "--replay-candidate")
            o.replayCandidate = v();
        else if (a == "--strategy")
            o.strategy = v();
        else if (a == "--gravity-tol")
            o.gravityTolerance = std::stod(v());
        else if (a == "--gravity-iterations")
            o.gravityIterations = std::stoi(v());
        else if (a == "--gravity-restart")
            o.gravityRestart = std::stoi(v());
        else if (a == "--strategy-probe")
            o.strategyProbe = true;
        else if (a == "--positivity-probe")
            o.positivityProbe = true;
        else if (a == "--positivity-certificate")
            o.positivityCertificate = true;
        else if (a == "--preconditioner-probe")
            o.preconditionerProbe = true;
        else if (a == "--preconditioner-variant")
            o.preconditionerVariant = v();
        else if (a == "--schur-sweeps")
            o.schurCorrectionSweeps = std::stoi(v());
        else if (a == "--schur-relaxation")
            o.schurRelaxation = std::stod(v());
        else if (a == "--pressure-probe")
            o.pressureProbe = true;
        else if (a == "--shape-curvature")
            o.shapeCurvature = true;
        else if (a == "--surface-trace")
            o.surfaceTrace = true;
        else if (a == "--probes")
            o.probes = true;
        else if (a == "--no-central")
            o.noCentral = true;
        else if (a == "--capture-iterations") {
            std::istringstream entries(v());
            std::string entry;
            while (std::getline(entries, entry, ',')) {
                size_t used = 0;
                const int iteration = std::stoi(entry, &used);
                if (used != entry.size() || iteration < 0)
                    throw std::invalid_argument("Invalid capture iteration");
                o.captureIterations.push_back(iteration);
            }
            if (o.captureIterations.empty())
                throw std::invalid_argument("Empty capture iteration list");
        }
        else
            throw std::invalid_argument("Unknown option: " + a);
    }
    if (o.output.empty() || o.n < 1 || o.n >= 5 || o.maxNewton < 0 || o.maxLinear < 1 || o.restart < 1)
        throw std::invalid_argument("Invalid options");
    if (o.probeRestart != 40 && o.probeRestart != 160)
        throw std::invalid_argument("--probe-restart must be 40 or 160");
    if ((o.pressureProbe || o.shapeCurvature || o.surfaceTrace || o.strategyProbe || o.positivityProbe) && (o.probeState.empty() || o.probeDirection.empty()))
        throw std::invalid_argument("Focused probes require --probe-state and --probe-direction");
    if (o.positivityCertificate && o.probeState.empty())
        throw std::invalid_argument("--positivity-certificate requires --probe-state");
    if (o.positivityCertificate && (o.probes || o.pressureProbe || o.shapeCurvature || o.surfaceTrace ||
        o.strategyProbe || o.positivityProbe || o.preconditionerProbe || !o.probeDirection.empty() || !o.replayCandidate.empty()))
        throw std::invalid_argument("--positivity-certificate is a certificate-only mode; omit other probe options");
    if (o.preconditionerProbe) {
        experiment::plateau_diagnostics::ParsePreconditionerBenchmarkVariant(o.preconditionerVariant);
        if (o.maxNewton != 0 || o.probes || o.pressureProbe || o.shapeCurvature || o.surfaceTrace ||
            o.strategyProbe || o.positivityProbe || o.positivityCertificate || !o.probeDirection.empty() ||
            !o.replayCandidate.empty() || !o.captureIterations.empty() || o.strategy != "baseline")
            throw std::invalid_argument("--preconditioner-probe requires --max-newton 0 and no nonlinear/probe strategies");
        if (!std::isfinite(o.linearTolerance) || o.linearTolerance <= 0.0)
            throw std::invalid_argument("A preconditioner benchmark requires a finite positive --linear-tol");
    }
    if (!o.preconditionerProbe) {
        if (o.preconditionerVariant != "independent" && o.preconditionerVariant != "gravity-schur" &&
            o.preconditionerVariant != "stellar-ldu" && o.preconditionerVariant != "stellar-ldu-gravity-schur")
            throw std::invalid_argument("Newton preconditioner variants are independent, gravity-schur, stellar-ldu, or stellar-ldu-gravity-schur");
        if (o.preconditionerVariant != "independent") {
            if (o.probes || o.pressureProbe || o.shapeCurvature || o.surfaceTrace || o.strategyProbe ||
                o.positivityProbe || o.positivityCertificate || !o.probeState.empty() || !o.probeDirection.empty() ||
                !o.replayCandidate.empty())
                throw std::invalid_argument("A custom Newton preconditioner cannot be combined with frozen or certificate probes");
            if (o.strategy != "baseline" && o.strategy != "positivity" && o.strategy != "curvature" &&
                o.strategy != "positivity-curvature")
                throw std::invalid_argument("Custom Newton preconditioners support baseline, positivity, curvature, or positivity-curvature");
            if (!std::isfinite(o.linearTolerance) || o.linearTolerance <= 0.0 || o.schurCorrectionSweeps < 0 ||
                !std::isfinite(o.schurRelaxation) || o.schurRelaxation <= 0.0)
                throw std::invalid_argument("Invalid custom Newton preconditioner controls");
        }
    }
    if ((!o.probeDirection.empty() || !o.replayCandidate.empty()) && o.probeState.empty())
        throw std::invalid_argument("A supplied direction or candidate requires --probe-state");
    if (o.strategy != "baseline" && o.strategy != "positivity" && o.strategy != "curvature" &&
        o.strategy != "reduced" && o.strategy != "positivity-curvature" && o.strategy != "positivity-reduced")
        throw std::invalid_argument("Unknown numerical strategy");
    if (!std::isfinite(o.gravityTolerance) || o.gravityTolerance <= 0.0 || o.gravityIterations < 1 || o.gravityRestart < 1)
        throw std::invalid_argument("Invalid gravity solve controls");
    if (std::filesystem::exists(o.output))
        throw std::runtime_error("Output directory already exists");
    std::filesystem::create_directories(o.output);
    return o;
}
std::ofstream File(const std::filesystem::path &p) {
    std::ofstream s(p);
    s.exceptions(std::ios::failbit | std::ios::badbit);
    s << std::setprecision(17);
    return s;
}
double Norm(std::span<const mfem::real_t> v, int offset, int size) {
    long double sum = 0;
    for (int i = offset; i < offset + size; ++i)
        sum += (long double)v[i] * v[i];
    return std::sqrt(sum);
}
double Norm(const mfem::Vector &v, int offset, int size) {
    return Norm(std::span<const mfem::real_t>{v.GetData(), (size_t)v.Size()}, offset, size);
}
struct Scales {
    double K, rho, xi, mu;
};
Scales Derive(double n) {
    if (n == 1.)
        return {2. * utils::G / std::numbers::pi, std::numbers::pi / 4., std::numbers::pi, std::numbers::pi};
    auto le = seed::integrateLaneEmden(n, 2000., 1e-3);
    if (!le.firstZeroCoordinate)
        throw std::runtime_error("No Lane Emden root");
    const double xi = *le.firstZeroCoordinate, mu = -xi * xi * le.thetaDerivative(le.thetaDerivative.Size() - 1);
    const double rho = xi * xi * xi / (4 * std::numbers::pi * mu);
    return {4 * std::numbers::pi * utils::G * std::pow(rho, 1. - 1. / n) / ((n + 1) * xi * xi), rho, xi, mu};
}
void SaveVector(const std::filesystem::path &p, const mfem::Vector &v) {
    auto s = File(p);
    s << v.Size() << '\n';
    v.Print(s, 1);
}
mfem::Vector LoadVector(const std::filesystem::path &path, int expectedSize) {
    std::ifstream input(path);
    int count = -1;
    input >> count;
    if (!input || count != expectedSize)
        throw std::runtime_error("Snapshot size mismatch: " + path.string());
    mfem::Vector result(count);
    for (int i = 0; i < count; ++i) {
        input >> result(i);
        if (!input || !std::isfinite(result(i)))
            throw std::runtime_error("Invalid snapshot: " + path.string());
    }
    std::string extra;
    if (input >> extra)
        throw std::runtime_error("Unexpected trailing snapshot data: " + path.string());
    return result;
}

template <class State>
void FocusedProbes(State &s, const fem::FEM &f, const Options &o, const mfem::Vector &physicalBase) {
    mfem::Vector base;
    s.normalizedOperator->NormalizeState(physicalBase, base);
    auto restore = [&]() {
        s.trialNormalizedState = base;
        if (!s.PrepareTrial())
            throw std::runtime_error("Focused probe base is inadmissible");
    };
    restore();
    const mfem::Vector residual(s.trialNormalizedResidual);
    SaveVector(o.output / "frozen_physical_base.txt", physicalBase);
    SaveVector(o.output / "frozen_normalized_base.txt", base);
    SaveVector(o.output / "frozen_normalized_residual.txt", residual);
    auto layout = File(o.output / "layout.csv");
    layout << "kind,block,offset,size,scale\n";
    for (const auto &b : s.Problem().GetManifest().residualBlocks())
        layout << "residual," << b.stableId << ',' << b.offset << ',' << b.size << ',' << b.scale << '\n';
    for (const auto &b : s.Problem().GetManifest().valueBlocks())
        layout << "value," << b.stableId << ',' << b.offset << ',' << b.size << ',' << b.scale << '\n';
    auto checks = File(o.output / "frozen_input_checks.csv");
    checks << "metric,value\nbase_residual_norm," << residual.Norml2() << '\n';
    const auto capturedBase = o.probeState.parent_path() / "normalized_base.txt";
    if (std::filesystem::is_regular_file(capturedBase)) {
        mfem::Vector difference = LoadVector(capturedBase, base.Size());
        difference -= base;
        checks << "captured_normalized_base_difference," << difference.Norml2() << '\n';
        if (difference.Norml2() > 64. * std::numeric_limits<double>::epsilon() * std::max(1., base.Norml2()))
            throw std::runtime_error("Frozen normalization disagrees with the captured base");
    }
    const auto capturedResidual = o.probeState.parent_path() / "normalized_residual.txt";
    if (std::filesystem::is_regular_file(capturedResidual)) {
        mfem::Vector difference = LoadVector(capturedResidual, residual.Size());
        difference -= residual;
        checks << "captured_residual_difference," << difference.Norml2() << '\n';
    }
    if (!o.replayCandidate.empty()) {
        s.trialNormalizedState = LoadVector(o.replayCandidate, base.Size());
        const auto prepared = s.PrepareTrial();
        auto replay = File(o.output / "replayed_candidate.csv");
        replay << "admissible,reason,stage,eos_code,minimum_jacobian,residual_norm\n";
        if (prepared)
            replay << "1,-1,-1,-1," << prepared->minimumJacobianDeterminant << ','
                   << s.trialNormalizedResidual.Norml2() << '\n';
        else
            replay << "0," << int(prepared.error().reason) << ',' << int(prepared.error().stage) << ','
                   << int(prepared.error().thermodynamicErrorCode) << ','
                   << prepared.error().minimumJacobianDeterminant << ",nan\n";
        restore();
    }
    if (!o.probeDirection.empty()) {
        const auto direction = LoadVector(o.probeDirection, base.Size());
        mfem::Vector action(base.Size());
        restore();
        s.normalizedOperator->Mult(direction, action);
        SaveVector(o.output / "probe_normalized_correction.txt", direction);
        checks << "direction_norm," << direction.Norml2() << "\naction_norm," << action.Norml2() << '\n';
        action += residual;
        checks << "linear_defect_norm," << action.Norml2() << "\nlinear_defect_relative,"
               << action.Norml2() / std::max(residual.Norml2(), 1e-300) << '\n';
        checks.flush();
        s.normalizedCorrection = direction;
        auto bounds = File(o.output / "frozen_enthalpy_bound.csv");
        bounds << "metric,value\n";
        for (const auto &[key, value] : experiment::plateau_diagnostics::MeasureEnthalpyStep(s, f, o.n, &physicalBase))
            bounds << key << ',' << value << '\n';
        bounds.flush();
        if (o.strategyProbe) {
            s.physicalInverse->Refresh();
            using Form = typename std::remove_cvref_t<decltype(s.Problem())>::FormType;
            solver::PreparedGravitySubsystem gravity(*s.normalizedOperator, *s.normalizedInverse,
                solver::MakeGravityBlockPartition<Form>(s.Problem().GetManifest().layout()),
                s.Problem().GetCommunicator());
            plateau_diagnostics::ProbeGravityApplications(gravity, o.output);
            experiment::plateau_diagnostics::ProbeNumericalStrategies(s, f, base, direction, o.output);
        }
        if (o.positivityProbe)
            experiment::plateau_diagnostics::ProbePolynomialEnthalpy(s, f, base, direction, o.output);
        if (o.surfaceTrace)
            experiment::plateau_diagnostics::ProbeSurfaceTrace(s, f, o.n, base, direction, o.output);
        if (o.pressureProbe)
            experiment::plateau_diagnostics::ProbePressureObstruction(s, f, o.n, base, direction, o.output);
        if (o.shapeCurvature)
            experiment::plateau_diagnostics::ProbeShapeCurvature(s, f, o.n, base, direction, o.output);
    }
    restore();
    mfem::Vector difference(s.trialNormalizedResidual);
    difference -= residual;
    checks << "final_restoration_residual_difference," << difference.Norml2() << '\n';
}

template <class State>
void FrozenProbes(State &s, const fem::FEM &f, const Options &o, const mfem::Vector *physicalBase = nullptr) {
    auto rows = s.Problem().GetManifest().residualBlocks();
    mfem::Vector frozenState(s.acceptedNormalizedState);
    if (physicalBase)
        s.normalizedOperator->NormalizeState(*physicalBase, frozenState);
    auto restoreFrozen = [&]() {
        s.trialNormalizedState = frozenState;
        auto p = s.PrepareTrial();
        if (!p)
            throw std::runtime_error("Frozen diagnostic state is inadmissible");
    };
    restoreFrozen();
    s.physicalInverse->Refresh();
    const mfem::Vector base(s.trialNormalizedResidual);
    if (base.Norml2() == 0.)
        return;
    auto results = File(o.output / "frozen_linear.csv");
    results << "restart,tolerance,iterations,converged,true_relative,seconds\n";
    auto histories = File(o.output / "frozen_history.csv");
    histories << "restart,tolerance,iteration,reported_norm,final\n";
    auto blocks = File(o.output / "frozen_blocks.csv");
    blocks << "restart,tolerance,block,F_norm,Jp_plus_F_norm,global_correction_norm\n";
    mfem::Vector selectedCorrection, baselineAction(base.Size());
    // These are two cold direct MFEM solves, not the configured production
    // backend or any recycled Krylov history.
    for (int restart : {40, 160}) {
        mfem::FGMRESSolver linear(MPI_COMM_WORLD);
        solver::ResidualHistoryMonitor monitor;
        linear.SetPreconditioner(*s.normalizedInverse);
        linear.SetOperator(*s.normalizedOperator);
        linear.SetKDim(restart);
        linear.SetRelTol(0.);
        linear.SetAbsTol(o.linearTolerance * base.Norml2());
        linear.SetMaxIter(o.maxLinear);
        linear.SetPrintLevel(-1);
        linear.SetMonitor(monitor);
        linear.iterative_mode = false;
        mfem::Vector b(base), p(base.Size()), action(base.Size());
        b *= -1.;
        p = 0.;
        const auto start = Clock::now();
        linear.Mult(b, p);
        double seconds = std::chrono::duration<double>(Clock::now() - start).count();
        s.normalizedOperator->Mult(p, action);
        action += base;
        const double rel = action.Norml2() / base.Norml2();
        results << restart << ',' << o.linearTolerance << ',' << linear.GetNumIterations() << ','
                << linear.GetConverged() << ',' << rel << ',' << seconds << '\n';
        results.flush();
        for (const auto &h : monitor.GetHistory())
            histories << restart << ',' << o.linearTolerance << ',' << h.iteration << ',' << h.reportedNorm << ','
                      << h.final << '\n';
        histories.flush();
        for (const auto &r : rows)
            blocks << restart << ',' << o.linearTolerance << ',' << r.stableId << ',' << Norm(base, r.offset, r.size)
                   << ',' << Norm(action, r.offset, r.size) << ',' << p.Norml2() << '\n';
        blocks.flush();
        std::cout << "Frozen restart=" << restart << " true=" << rel << " seconds=" << seconds << std::endl;
        if (restart == o.probeRestart)
            selectedCorrection = p;
    }
    std::cout << "Frozen probe direction restart=" << o.probeRestart << std::endl;
    s.normalizedCorrection = selectedCorrection;
    s.normalizedOperator->Mult(selectedCorrection, baselineAction);
    auto positivity = File(o.output / "enthalpy_step_bound.csv");
    positivity << "metric,value\n";
    for (const auto &[key, value] : experiment::plateau_diagnostics::MeasureEnthalpyStep(s, f, o.n, physicalBase))
        positivity << key << ',' << value << '\n';
    positivity.flush();
    SaveVector(o.output / "probe_normalized_correction.txt", selectedCorrection);
    auto fd = File(o.output / "finite_differences.csv");
    fd << "epsilon,admissible,block,absolute_error,relative_to_Jp,action_norm\n";
    const auto safe = s.EstimateLargestSafeStepSize(1., .9);
    if (!std::isfinite(safe.stepSize) || safe.stepSize <= 0.)
        return;
    auto trials = File(o.output / "frozen_trials.csv");
    trials << "alpha,admissible,rejection_reason,F_norm,merit_ratio\n";
    // Diagnostic trials are never committed. Each samples the same frozen Newton direction.
    for (double a : {1., .5, .1, .01, .001, .0001, .00001}) {
        const double epsilon = std::min(1., safe.stepSize) * a;
        s.trialNormalizedState = frozenState;
        s.trialNormalizedState.Add(epsilon, selectedCorrection);
        const auto prep = s.PrepareTrial();
        if (!prep) {
            trials << epsilon << ",0," << int(prep.error().reason) << ",nan,nan\n";
            fd << epsilon << ",0,all,nan,nan,nan\n";
        } else {
            trials << epsilon << ",1,-1," << s.trialNormalizedResidual.Norml2() << ','
                   << std::pow(s.trialNormalizedResidual.Norml2() / base.Norml2(), 2) << '\n';
            mfem::Vector error(s.trialNormalizedResidual);
            error -= base;
            error /= epsilon;
            error -= baselineAction;
            for (const auto &r : rows) {
                double e = Norm(error, r.offset, r.size), j = Norm(baselineAction, r.offset, r.size);
                fd << epsilon << ",1," << r.stableId << ',' << e << ',' << e / std::max(j, 1e-300) << ',' << j << '\n';
            }
        }
        trials.flush();
        fd.flush();
        restoreFrozen();
    }
}

template <bool FullDefaultStack, class Context>
int RunContext(Context &context, const Options &o, const Scales &scales) {
    std::vector<operators::RootBlockDescriptor> rows, values;
    mfem::Vector residualFactors, stateFactors;
    std::vector<int> centralRows;
    int centralOffset = -1;
    const bool contextUsesGravitySchur = !o.preconditionerProbe &&
        (o.preconditionerVariant == "gravity-schur" || o.preconditionerVariant == "stellar-ldu-gravity-schur");
    auto metadata = File(o.output / "metadata.txt");
    metadata << "n=" << o.n << "\nK=" << scales.K << "\nrho_c=" << scales.rho << "\nJ=" << o.angularMomentum
             << "\nM=1\nR_reference=1\nxi1=" << scales.xi << "\nmu1=" << scales.mu << "\ncentral_phase=" << !o.noCentral
             << "\nrestart=" << o.restart << "\nprobe_restart=" << o.probeRestart
             << "\nfrozen_probe_backend=" << (o.preconditionerProbe ? "direct_mfem_fgmres_preconditioner_benchmark" :
                 o.positivityCertificate ? "polynomial_certificate_no_linear_solve" :
                 o.probeDirection.empty() ? "direct_mfem_fgmres_cold" : "supplied_direction_no_linear_solve")
             << "\nlinear_tol=" << o.linearTolerance << "\nmax_linear=" << o.maxLinear
             << "\nmax_newton=" << o.maxNewton << "\natol=" << o.atol << "\nrtol=" << o.rtol
             << "\nstrategy=" << o.strategy << "\ngravity_tolerance=" << o.gravityTolerance
             << "\ngravity_iterations=" << o.gravityIterations << "\ngravity_restart=" << o.gravityRestart
             << "\npreconditioner_probe=" << o.preconditionerProbe
             << "\npreconditioner_variant=" << o.preconditionerVariant
             << "\ncontext_preconditioner_variant=" << (o.preconditionerProbe ? "independent" : o.preconditionerVariant)
             << "\ncontext_gravity_schur_sweeps="
             << (contextUsesGravitySchur ? o.schurCorrectionSweeps : 0)
             << "\ncontext_gravity_schur_relaxation=" << (contextUsesGravitySchur ? o.schurRelaxation : 0.0)
             << "\nschur_correction_sweeps=" << o.schurCorrectionSweeps << "\nschur_relaxation=" << o.schurRelaxation
             << "\ncompiler=" << __VERSION__ << "\ncompiled=" << __DATE__ << ' ' << __TIME__ << '\n';
    metadata.flush();
    // Custom Newton recipes do not instantiate the nine-variant frozen suite.
    if constexpr (FullDefaultStack) if (o.preconditionerProbe) {
        metadata << "benchmark_input=" << (o.probeState.empty() ? "lane_emden_seed" : "physical_snapshot") << '\n';
        if (!o.probeState.empty()) metadata << "probe_state=" << o.probeState << '\n';
        solver::detail::StellarEquilibriumContextDiagnostics::WithState(context, [&](auto &s, const fem::FEM &) {
            if (o.probeState.empty()) {
                s.trialNormalizedState = s.acceptedNormalizedState;
            } else {
                const auto snapshot = LoadVector(o.probeState, s.AcceptedPhysicalState().Size());
                s.normalizedOperator->NormalizeState(snapshot, s.trialNormalizedState);
            }
            const auto preparationStart = Clock::now();
            if (!s.PrepareTrial())
                throw std::runtime_error("Preconditioner benchmark state is inadmissible to the physical operator");
            metadata << "benchmark_frozen_preparation_seconds="
                     << std::chrono::duration<double>(Clock::now() - preparationStart).count() << '\n';
            metadata.flush();
            experiment::plateau_diagnostics::BenchmarkPreconditioners(s, o.output,
                {.variant = o.preconditionerVariant, .restart = o.restart,
                 .maximumIterations = o.maxLinear, .relativeTolerance = o.linearTolerance,
                 .schurCorrectionSweeps = o.schurCorrectionSweeps, .schurRelaxation = o.schurRelaxation});
        });
        // A missed linear gate is a measured outcome, not failure to execute
        // the experiment. The CSV retains both recurrence and true gate results.
        return 0;
    }
    if constexpr (FullDefaultStack) if (!o.probeState.empty()) {
        if (o.maxNewton != 0)
            throw std::invalid_argument("--probe-state requires --max-newton 0");
        metadata << "probe_state=" << o.probeState << '\n';
        if (!o.probeDirection.empty())
            metadata << "probe_direction=" << o.probeDirection << '\n';
        if (!o.replayCandidate.empty())
            metadata << "replay_candidate=" << o.replayCandidate << '\n';
        metadata << "pressure_probe=" << o.pressureProbe << "\nshape_curvature=" << o.shapeCurvature
                 << "\nsurface_trace=" << o.surfaceTrace << "\npositivity_certificate=" << o.positivityCertificate << '\n';
        metadata.flush();
        solver::detail::StellarEquilibriumContextDiagnostics::WithState(context, [&](auto &s, const fem::FEM &f) {
            const auto snapshot = LoadVector(o.probeState, s.AcceptedPhysicalState().Size());
            if (o.positivityCertificate) {
                // A legacy endpoint may already violate admissibility. Its
                // polynomial must be inspectable without preparing that state.
                mfem::Vector normalizedSnapshot;
                s.normalizedOperator->NormalizeState(snapshot,normalizedSnapshot);
                experiment::plateau_diagnostics::CertifyPolynomialEnthalpy(s,normalizedSnapshot,o.output);
                return;
            }
            s.normalizedOperator->NormalizeState(snapshot, s.trialNormalizedState);
            auto prep = s.PrepareTrial();
            if (!prep)
                throw std::runtime_error("Snapshot is inadmissible");
            std::cout << "Snapshot F=" << s.trialNormalizedResidual.Norml2() << std::endl;
            if (!o.probeDirection.empty() || !o.replayCandidate.empty()) {
                FocusedProbes(s, f, o, snapshot);
                return;
            }
            experiment::plateau_diagnostics::InspectPreconditioner(s, o.output);
            FrozenProbes(s, f, o, &snapshot);
        });
        return 0;
    }
    auto physicalMetrics = [&](auto &s, const fem::FEM &f, const std::string &label) {
        std::cout << label << " physical diagnostics" << std::endl;
        experiment::polytrope_validation::PhysicalState physical(s, f);
        auto out = File(o.output / (label + "_physical.csv"));
        out << "quadrature,metric,value\n";
        for (int q : {14, 18}) {
            auto metrics = experiment::plateau_diagnostics::MeasurePhysical(physical, o.n, scales.K, q);
            metrics["normalized_bordered_residual"] = physical.normalizedBorderedResidualNorm;
            metrics["normalized_unbordered_residual"] = physical.normalizedUnborderedResidualNorm;
            metrics["central_border"] = physical.centralBorder;
            metrics["normalized_border_action"] = physical.normalizedCentralBorderActionNorm;
            metrics["minimum_sampled_jacobian"] = s.acceptedMinimumJacobianDeterminant;
            metrics["mass_absolute_error"] = std::abs(metrics["mass"] - 1.);
            metrics["angular_momentum_absolute_error"] = std::abs(metrics["angular_momentum_z"] - o.angularMomentum);
            std::cout << "  q=" << q << " virial_h=" << metrics["enthalpy_virial_error"] << " M=" << metrics["mass"]
                      << std::endl;
            for (auto &[k, v] : metrics)
                out << q << ',' << k << ',' << v << '\n';
            out.flush();
        }
    };
    solver::detail::StellarEquilibriumContextDiagnostics::WithState(context, [&](auto &s, const fem::FEM &f) {
        auto r = s.Problem().GetManifest().residualBlocks(), v = s.Problem().GetManifest().valueBlocks();
        rows.assign(r.begin(), r.end());
        values.assign(v.begin(), v.end());
        residualFactors = s.normalizedOperator->GetNormalization().ResidualFactors();
        stateFactors = s.normalizedOperator->GetNormalization().StateFactors();
        for (const auto &b : values)
            if (b.stableId == "fixed_central_density.border")
                centralOffset = b.offset;
        experiment::polytrope_validation::PhysicalState physical(s, f);
        centralRows = physical.centralHydrostaticRows;
        auto layout = File(o.output / "layout.csv");
        layout << "kind,block,offset,size,scale\n";
        for (auto &b : rows)
            layout << "residual," << b.stableId << ',' << b.offset << ',' << b.size << ',' << b.scale << '\n';
        for (auto &b : values)
            layout << "value," << b.stableId << ',' << b.offset << ',' << b.size << ',' << b.scale << '\n';
        physicalMetrics(s, f, "initial");
    });
    auto history = File(o.output / "newton.csv");
    history << "iteration,accepted,F,alpha,trials,linear_status,linear_iterations,linear_true,seconds,linear_seconds,"
               "prepare_seconds,refresh_seconds,geometry_limited,geometry_safe_step,central_border,unbordered_F,"
               "linear_initial_residual,linear_rhs_norm,linear_initial_over_rhs\n";
    auto blocklog = File(o.output / "blocks.csv");
    blocklog << "iteration,stage,block,normalized_norm,physical_norm,max_abs_normalized\n";
    auto triallog = File(o.output / "trials.csv");
    triallog << "iteration,trial,alpha,disposition,source,F,min_jacobian\n";
    auto stateConsistency = File(o.output / "trial_state_consistency.csv");
    stateConsistency << "iteration,trial,normalized_physical_roundtrip_error,relative_error\n";
    mfem::Vector captureBase;
    auto captureDirectory = [&](int iteration) {
        return o.output / ("capture_" + std::to_string(iteration));
    };
    auto shouldCapture = [&](int iteration) {
        return std::find(o.captureIterations.begin(), o.captureIterations.end(), iteration) != o.captureIterations.end();
    };
    auto copySpan = [](std::span<const mfem::real_t> values) {
        mfem::Vector result(static_cast<int>(values.size()));
        for (int i = 0; i < result.Size(); ++i)
            result(i) = values[i];
        return result;
    };
    auto record = [&](int it, const char *stage, std::span<const mfem::real_t> residual) {
        for (const auto &b : rows) {
            long double ps = 0;
            double mx = 0;
            for (int i = b.offset; i < b.offset + b.size; ++i) {
                const double p = residual[i] / residualFactors(i);
                ps += (long double)p * p;
                mx = std::max(mx, std::abs(residual[i]));
            }
            blocklog << it << ',' << stage << ',' << b.stableId << ',' << Norm(residual, b.offset, b.size) << ','
                     << std::sqrt(ps) << ',' << mx << '\n';
        }
        blocklog.flush();
    };
    auto observer = solver::nonlinear::makeObserver(
        [&](const solver::nonlinear::BeforeIteration &e) {
            record(e.iteration, "before", e.normalizedResidual);
            if (shouldCapture(e.iteration)) {
                const auto directory = captureDirectory(e.iteration);
                std::filesystem::create_directory(directory);
                captureBase = copySpan(e.normalizedState);
                SaveVector(directory / "normalized_base.txt", captureBase);
                SaveVector(directory / "physical_base.txt", copySpan(e.physicalState));
                SaveVector(directory / "normalized_residual.txt", copySpan(e.normalizedResidual));
            }
            std::cout << "Newton " << e.iteration << " F=" << e.residualNorm << std::endl;
        },
        [&](const solver::nonlinear::AfterLineSearchTrial &e) {
            if (e.candidatePhysicalState.size() != e.candidateNormalizedState.size() ||
                e.candidateNormalizedState.size() != static_cast<size_t>(stateFactors.Size()))
                throw std::runtime_error("Trial observer exposed incompatible state representations");
            long double stateError = 0.0L, stateNorm = 0.0L;
            for (int i = 0; i < stateFactors.Size(); ++i) {
                const long double difference = e.candidatePhysicalState[i] * stateFactors(i) - e.candidateNormalizedState[i];
                stateError += difference*difference;
                stateNorm += static_cast<long double>(e.candidateNormalizedState[i])*e.candidateNormalizedState[i];
            }
            const double relativeStateError = std::sqrt(stateError)/std::max(1.0L,std::sqrt(stateNorm));
            stateConsistency << e.iteration << ',' << e.trial << ',' << std::sqrt(stateError) << ',' << relativeStateError << '\n';
            stateConsistency.flush();
            if (std::isfinite(relativeStateError) && relativeStateError > 64.0*std::numeric_limits<double>::epsilon())
                throw std::runtime_error("Trial observer physical state disagrees with its normalized state");
            triallog << e.iteration << ',' << e.trial << ',' << e.stepLength << ',' << int(e.disposition) << ','
                     << e.rejectionSource << ',' << (e.metric ? e.metric->residualNorm : NAN) << ','
                     << e.minimumJacobianDeterminant.value_or(NAN) << '\n';
            triallog.flush();
            if (e.trial == 0 && shouldCapture(e.iteration)) {
                if (!(e.stepLength > 0.) || !std::isfinite(e.stepLength) ||
                    e.candidateNormalizedState.size() != static_cast<size_t>(captureBase.Size()))
                    throw std::runtime_error("Invalid candidate for diagnostic capture");
                const auto directory = captureDirectory(e.iteration);
                const auto candidate = copySpan(e.candidateNormalizedState);
                SaveVector(directory / "normalized_candidate.txt", candidate);
                SaveVector(directory / "physical_candidate.txt", copySpan(e.candidatePhysicalState));
                mfem::Vector direction(candidate), reconstructed(captureBase);
                direction -= captureBase;
                direction /= e.stepLength;
                reconstructed.Add(e.stepLength, direction);
                reconstructed -= candidate;
                SaveVector(directory / "normalized_effective_trial_direction.txt", direction);
                SaveVector(directory / "normalized_direction.txt", copySpan(e.normalizedDirection));
                auto info = File(directory / "capture.txt");
                info << "iteration=" << e.iteration << "\nfirst_trial_alpha=" << e.stepLength
                     << "\ndisposition=" << int(e.disposition) << "\nsource=" << e.rejectionSource
                     << "\nbase_norm=" << captureBase.Norml2() << "\ndirection_norm=" << direction.Norml2()
                     << "\ncandidate_reconstruction_error=" << reconstructed.Norml2()
                     << "\ndirection_roundoff_scale_estimate="
                     << 8. * std::numeric_limits<double>::epsilon() *
                            (captureBase.Norml2() + candidate.Norml2()) / e.stepLength
                     << "\ndirection_source=exact_observer_tangent\n"
                     << "effective_trial_direction_source=candidate_difference_divided_by_alpha\n"
                     << "roundoff_scale_applies_to_effective_trial_direction_only=1\n"
                     << "roundoff_scale_is_estimate_not_certificate=1\n";
            }
        },
        [&](const solver::nonlinear::AfterIteration &e) {
            record(e.iteration, "after", e.normalizedResidual);
            double border = centralOffset < 0 ? 0 : e.physicalState[centralOffset];
            mfem::Vector unbordered((int)e.normalizedResidual.size());
            for (int i = 0; i < unbordered.Size(); ++i)
                unbordered(i) = e.normalizedResidual[i];
            for (int i : centralRows)
                unbordered(i) -= residualFactors(i) * border;
            const double linearInitial = e.linearSolve ? e.linearSolve->initialResidualNorm : NAN;
            const double linearRhs = e.linearSolve ? e.linearSolve->rightHandSideNorm : NAN;
            double linearInitialRelative = NAN;
            if (linearRhs > 0.)
                linearInitialRelative = linearInitial / linearRhs;
            else if (linearRhs == 0. && std::isfinite(linearInitial))
                linearInitialRelative = linearInitial == 0. ? 0. : std::numeric_limits<double>::infinity();
            history << e.iteration << ',' << e.stepAccepted << ',' << e.residualNorm << ',' << e.acceptedStepLength
                    << ',' << e.lineSearchTrials << ',' << (e.linearSolve ? int(e.linearSolve->status) : -1) << ','
                    << (e.linearSolve ? e.linearSolve->iterations : 0) << ','
                    << (e.linearSolve ? e.linearSolve->relativeTrueResidualNorm : NAN) << ',' << e.iterationSeconds
                    << ',' << (e.linearSolve ? e.linearSolve->solveSeconds : 0) << ',' << e.trialPreparationSeconds
                    << ',' << e.preconditionerRefreshSeconds << ','
                    << (e.geometryPreflight ? e.geometryPreflight->limitedByGeometry : false) << ','
                    << (e.geometryPreflight ? e.geometryPreflight->stepSize : NAN) << ',' << border << ','
                    << unbordered.Norml2() << ',' << linearInitial << ',' << linearRhs << ','
                    << linearInitialRelative << '\n';
            history.flush();
            std::cout << "  alpha=" << e.acceptedStepLength << " F=" << e.residualNorm
                      << " linear=" << (e.linearSolve ? e.linearSolve->relativeTrueResidualNorm : NAN)
                      << " lambda=" << border << " unbordered=" << unbordered.Norml2()
                      << " seconds=" << e.iterationSeconds << std::endl;
        });
    bool converged = false;
    if (o.maxNewton > 0) {
        namespace nl = solver::nonlinear;
        const nl::NewtonOptions options{
            .relativeTolerance = o.rtol, .absoluteTolerance = o.atol, .maximumIterations = o.maxNewton,
            .linearSolve = {.relativeTolerance = o.linearTolerance, .absoluteTolerance = 0., .maximumIterations = o.maxLinear}};
        const nl::GravitySolveOptions gravity{
            .control = {.relativeTolerance = o.gravityTolerance, .absoluteTolerance = 1e-14,
                        .maximumIterations = o.gravityIterations},
            .backend = {.restartLength = o.gravityRestart, .printLevel = -1}};
        auto run = [&](auto numerics) {
            auto solve = solver::make(context, nl::Newton(options, nl::NormalizedResidualMetric{}, numerics), observer);
            return solve.evaluate();
        };
        auto report = [&] {
            if (o.strategy == "positivity") return run(nl::NewtonNumerics{nl::PolynomialEnthalpy{}, nl::MonolithicGravity{}});
            if (o.strategy == "curvature") return run(nl::NewtonNumerics{nl::SampledEnthalpy{}, nl::GravityCurvatureCorrection{gravity}});
            if constexpr (FullDefaultStack)
                if (o.strategy == "reduced") return run(nl::NewtonNumerics{nl::SampledEnthalpy{}, nl::EliminatedGravity{gravity}});
            if (o.strategy == "positivity-curvature") return run(nl::NewtonNumerics{nl::PolynomialEnthalpy{}, nl::GravityCurvatureCorrection{gravity}});
            if constexpr (FullDefaultStack)
                if (o.strategy == "positivity-reduced") return run(nl::NewtonNumerics{nl::PolynomialEnthalpy{}, nl::EliminatedGravity{gravity}});
            return run(nl::NewtonNumerics{});
        }();
        const auto &strategyDiagnostics = report.diagnostics().numericalStrategies;
        auto strategyMetrics = File(o.output / "numerical_strategies.csv");
        strategyMetrics << "metric,value\n";
        const auto metric = [&](const char *name, auto value) { strategyMetrics << name << ',' << value << '\n'; };
        metric("polynomial_checks", strategyDiagnostics.polynomialChecks);
        metric("direction_repairs", strategyDiagnostics.directionRepairs);
        metric("last_repair_fraction", strategyDiagnostics.lastRepairFraction);
        metric("last_merit_directional_derivative", strategyDiagnostics.lastMeritDirectionalDerivative);
        metric("gravity_solves", strategyDiagnostics.gravitySolves);
        metric("gravity_iterations", strategyDiagnostics.gravityIterations);
        metric("gravity_failures", strategyDiagnostics.gravityFailures);
        metric("gravity_prepared_failures", strategyDiagnostics.gravityPreparedFailures);
        metric("compact_gravity_operator", strategyDiagnostics.compactGravityOperator);
        metric("compact_gravity_inverse", strategyDiagnostics.compactGravityInverse);
        metric("linear_report_space", static_cast<int>(strategyDiagnostics.lastLinearReportSpace));
        metric("gravity_solve_seconds", strategyDiagnostics.gravitySolveSeconds);
        metric("last_gravity_true_residual", strategyDiagnostics.lastGravityTrueResidual);
        metric("last_gravity_correction_norm", strategyDiagnostics.lastGravityCorrectionNorm);
        metric("last_gravity_prepared_defect", strategyDiagnostics.lastGravityPreparedDefect);
        metric("last_gravity_affine_disagreement", strategyDiagnostics.lastGravityAffineDisagreement);
        metric("last_gravity_arithmetic_allowance", strategyDiagnostics.lastGravityArithmeticAllowance);
        metric("last_gravity_prepared_target_threshold", strategyDiagnostics.lastGravityPreparedTargetThreshold);
        metric("initial_gravity_residual", strategyDiagnostics.initialGravityResidual);
        metric("completed_initial_gravity_residual", strategyDiagnostics.completedInitialGravityResidual);
        metric("last_lifted_relative_true_residual", strategyDiagnostics.lastLiftedRelativeTrueResidual);
        converged = report.converged();
        metadata << "converged=" << converged << "\naccepted_steps=" << report.completedNonlinearIterations() << '\n';
        if (!converged)
            metadata << "failure=" << report.failure().message << '\n';
        metadata.flush();
    }
    solver::detail::StellarEquilibriumContextDiagnostics::WithState(context, [&](auto &s, const fem::FEM &f) {
        SaveVector(o.output / "accepted_state.txt", s.AcceptedPhysicalState());
        physicalMetrics(s, f, "final");
        const auto &schur = s.physicalInverse->GetGroupedPreconditioner().GetFactorization().GetSchurComplement();
        auto sh = File(o.output / "border_schur.csv");
        for (int i = 0; i < schur.Height(); ++i) {
            for (int j = 0; j < schur.Width(); ++j)
                sh << (j ? "," : "") << schur(i, j);
            sh << '\n';
        }
        if constexpr (FullDefaultStack) if (o.probes) {
            experiment::plateau_diagnostics::InspectPreconditioner(s, o.output);
            FrozenProbes(s, f, o);
        }
    });
    return converged ? 0 : 3;
}
template <class Model> int Setup(Model model, fem::FEM f, const Options &o, const Scales &scales) {
    auto d = equilibrium::makeStellarDiscretization(
        std::move(f), normalization::PhysicalRieszDiagonal{dimensions::LengthValue{1.}, utils::G});
    const auto withRecipe = [&](auto recipe) {
        using Recipe = decltype(recipe);
        using Problem = equilibrium::StellarEquilibriumProblem<Model, decltype(d)>;
        static_assert(preconditioning::StellarPreconditionerRuntimeAvailableFor<Recipe, Problem>);
        using Inverse = preconditioning::PreparedStellarInverseType<Recipe, Problem>;
        static_assert(solver::LinearPreconditionerApplicationContractAvailable<Inverse>);
        static_assert(solver::linearPreconditionerApplicationContract<Inverse> ==
                      preconditioning::ApplicationContract::stationary_linear);
        auto context = solver::makeContext(std::move(model), std::move(d), std::move(recipe),
                                           solver::linear::FGMRES({.restartLength = o.restart, .printLevel = -1}),
                                           seed::LaneEmden({.centralDensity = dimensions::DensityValue{scales.rho}}));
        constexpr bool fullDefaultStack = std::same_as<Recipe, preconditioning::DefaultStellarPreconditioner>;
        return RunContext<fullDefaultStack>(context, o, scales);
    };
    if (o.preconditionerProbe || o.preconditionerVariant == "independent")
        return withRecipe(preconditioning::makePreconditioner());
    if (o.preconditionerVariant == "gravity-schur")
        return withRecipe(experiment::plateau_diagnostics::GravitySchurNewtonPreconditioner{
            {.correctionSweeps = o.schurCorrectionSweeps, .relaxation = o.schurRelaxation}});
    if (o.preconditionerVariant == "stellar-ldu")
        return withRecipe(experiment::plateau_diagnostics::StellarLDUNewtonPreconditioner{
            {.correctionSweeps = 0, .relaxation = 1.0}});
    if (o.preconditionerVariant == "stellar-ldu-gravity-schur")
        return withRecipe(experiment::plateau_diagnostics::StellarLDUGravitySchurNewtonPreconditioner{
            {.correctionSweeps = o.schurCorrectionSweeps, .relaxation = o.schurRelaxation}});
    throw std::logic_error("An unsupported Newton preconditioner bypassed option validation.");
}
int Run(const Options &o) {
    const auto scales = Derive(o.n);
    std::cout << "n=" << o.n << " K=" << scales.K << " rho_c=" << scales.rho << " J=" << o.angularMomentum << std::endl;
    utils::Args args;
    args.mesh_file = o.mesh;
    args.p.rtol = 1e-12;
    args.p.atol = 1e-12;
    auto f = fem::setup_fem(o.mesh, args, 0);
    if (!f.okay())
        throw std::runtime_error("FEM setup failed");
    auto eos = eos::Polytrope({.n = o.n, .K = scales.K});
    auto surf = surface::Isobaric({.Psurf = dimensions::PressureValue{0.}});
    auto mass = integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.}});
    auto angular = integral::FixedAngularMomentum(
        {.Jtotal = dimensions::AngularMomentumValue{o.angularMomentum}, .axis = {0., 0., 1.}, .center = {0., 0., 0.}});
    if (o.noCentral)
        return Setup(model::StellarModel(eos, surf, mass, angular), std::move(f), o, scales);
    return Setup(model::StellarModel(eos, surf, mass, angular,
                                     constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{scales.rho}})),
                 std::move(f), o, scales);
}
} // namespace
int main(int argc, char **argv) {
    mfem::Mpi::Init(argc, argv);
    int result = 0;
    try {
        int count;
        MPI_Comm_size(MPI_COMM_WORLD, &count);
        if (count != 1)
            throw std::runtime_error("Single rank only");
        mfem::Device device("cpu");
        std::cout << std::setprecision(12);
        result = Run(Parse(argc, argv));
    } catch (const std::exception &e) {
        std::cerr << "Diagnostic failure: " << e.what() << std::endl;
        result = 2;
    }
    mfem::Mpi::Finalize();
    return result;
}