#pragma once

// Experiment-only: one frozen Jacobian and fixed normalization for every split.
#include "plateau_nonlinear_step_diagnostics.hpp"
#include "shape_field_physical_diagnostics.hpp"
#include <vector>

namespace experiment::plateau_diagnostics {
template <class State>
void ProbeWeakShapeField(State &state, const mean_field::fem::FEM &fem,
                        const mfem::Vector &normalizedBase, const mfem::Vector *suppliedDirection,
                        const std::filesystem::path &output, NonlinearStepProbeOptions options = {}) {
    using namespace mean_field;
    namespace d = nonlinear_step_detail;
    options.linear.Validate();
    int ranks = 0;
    MPI_Comm_size(state.Problem().GetCommunicator(), &ranks);
    const int size = state.Problem().StateSize();
    if (ranks != 1 || size != state.Problem().EquationSize() || normalizedBase.Size() != size ||
        !d::Finite(normalizedBase) || options.restartLength < 1 ||
        (suppliedDirection && (suppliedDirection->Size() != size || !d::Finite(*suppliedDirection))))
        throw std::invalid_argument("Weak shape probes require one rank, finite compatible vectors and valid controls.");
    const auto begin = d::Clock::now();
    const mfem::Vector base(normalizedBase);
    const auto &normalization = state.normalizedOperator->GetNormalization();
    const mfem::Vector stateFactors(normalization.StateFactors()), residualFactors(normalization.ResidualFactors());
    const mfem::Vector savedAccepted(state.acceptedNormalizedState), savedResidual(state.acceptedNormalizedResidual),
        savedPhysical(state.AcceptedPhysicalState()), savedCorrection(state.normalizedCorrection),
        savedPhysicalCorrection(state.physicalCorrection), savedVolumeDirection(state.volumeDisplacementDirection);
    std::filesystem::create_directories(output);
    const auto file = [&](const std::string &name) { return output / ("weak_shape_" + name); };
    auto summary = d::Open(file("summary.csv"));
    summary << "metric,value\n";
    const auto metric = [&](const char *key, auto value) { summary << key << ',' << value << '\n'; };
    auto manifest = d::Open(file("layout.csv"));
    manifest << "kind,block,offset,size\n";
    const auto values = state.Problem().GetManifest().valueBlocks();
    const auto rows = state.Problem().GetManifest().residualBlocks();
    for (const auto &b : values) manifest << "value," << b.stableId << ',' << b.offset << ',' << b.size << '\n';
    for (const auto &b : rows) manifest << "residual," << b.stableId << ',' << b.offset << ',' << b.size << '\n';
    manifest.flush();
    mfem::Vector residual(size), direction(size), action(size), physicalBase, physicalDirection;
    bool haveBase = false, finalized = false;
    int preparations = 0, applications = 0;
    double maximumRestoreError = 0.;
    const auto prepare = [&](const mfem::Vector &candidate) {
        state.trialNormalizedState = candidate;
        ++preparations;
        return state.PrepareTrial();
    };
    const auto restore = [&]() {
        if (!prepare(base)) throw std::runtime_error("Weak shape probe cannot restore its base.");
        if (!haveBase) { residual = state.trialNormalizedResidual; haveBase = true; }
        if (!d::Finite(state.trialNormalizedResidual)) throw std::runtime_error("Nonfinite frozen residual.");
        maximumRestoreError = std::max(maximumRestoreError, d::Difference(residual, state.trialNormalizedResidual));
    };
    const auto apply = [&](const mfem::Vector &v, mfem::Vector &a) {
        state.normalizedOperator->Mult(v, a);
        ++applications;
        if (!d::Finite(a)) throw std::runtime_error("Nonfinite weak shape Jacobian action.");
    };
    const auto finish = [&]() {
        // Finalization is attempted once, including on exceptional paths; an
        // invariant failure must not be obscured by overwriting its artifacts.
        finalized = true;
        restore();
        state.physicalInverse->Refresh();
        state.normalizedCorrection = savedCorrection;
        state.physicalCorrection = savedPhysicalCorrection;
        state.volumeDisplacementDirection = savedVolumeDirection;
        const bool acceptedSame = d::Identical(savedAccepted, state.acceptedNormalizedState) &&
            d::Identical(savedResidual, state.acceptedNormalizedResidual) &&
            d::Identical(savedPhysical, state.AcceptedPhysicalState());
        const bool factorsSame = d::Identical(stateFactors, normalization.StateFactors()) &&
            d::Identical(residualFactors, normalization.ResidualFactors());
        metric("accepted_buffers_bitwise_identical", acceptedSame);
        metric("normalization_factors_bitwise_identical", factorsSame);
        metric("direction_buffers_bitwise_identical", d::Identical(savedCorrection, state.normalizedCorrection) &&
            d::Identical(savedPhysicalCorrection, state.physicalCorrection) &&
            d::Identical(savedVolumeDirection, state.volumeDisplacementDirection));
        metric("maximum_restore_error", maximumRestoreError);
        d::Save(file("restored_residual.txt"), state.trialNormalizedResidual);
        metric("preparations", preparations);
        metric("direct_jacobian_applications", applications);
        metric("total_probe_seconds", d::Seconds(begin));
        summary.flush();
        if (!acceptedSame || !factorsSame || maximumRestoreError != 0.)
            throw std::runtime_error("Weak shape probe changed frozen invariants.");
    };
    try {
        restore();
        state.physicalInverse->Refresh();
        state.normalizedOperator->DenormalizeState(base, physicalBase);
        metric("schema_version", 1);
        metric("state_size", size);
        metric("supplied_direction", suppliedDirection != nullptr);
        metric("accepted_commits", 0);
        metric("outer_tolerance", options.linear.relativeTolerance);
        metric("outer_absolute_tolerance", options.linear.absoluteTolerance);
        metric("maximum_iterations", options.linear.maximumIterations);
        metric("restart", options.restartLength);
        if (suppliedDirection) direction = *suppliedDirection;
        else {
            mfem::Vector rhs(residual);
            rhs *= -1.;
            direction = 0.;
            auto backend = solver::linear::prepareLinearBackend(
                solver::linear::FGMRES({.restartLength = options.restartLength, .printLevel = -1}),
                *state.normalizedOperator, *state.normalizedInverse, state.Problem().GetCommunicator());
            const auto solve = backend.Solve(rhs, direction, options.linear);
            metric("linear_status", int(solve.status));
            metric("linear_iterations", solve.iterations);
            metric("linear_seconds", solve.solveSeconds);
            metric("linear_reported_true_norm", solve.trueResidualNorm);
            restore();
        }
        if (!d::Finite(direction) || d::Norm(direction) == 0.)
            throw std::runtime_error("Weak shape probe requires a nonzero finite direction.");
        apply(direction, action);
        mfem::Vector defect(action);
        defect += residual;
        const double threshold = options.linear.ConvergenceThreshold(d::Norm(residual));
        for (const double norm : {d::Norm(residual), d::Norm(direction), d::Norm(action), d::Norm(defect), threshold})
            if (!std::isfinite(norm)) throw std::runtime_error("Overflowing weak shape norm or linear threshold.");
        metric("base_norm", d::Norm(residual));
        metric("direction_norm", d::Norm(direction));
        metric("action_norm", d::Norm(action));
        metric("verified_defect_norm", d::Norm(defect));
        metric("verified_relative_defect", d::Ratio(d::Norm(defect), d::Norm(residual)));
        metric("verified_threshold", threshold);
        metric("verified_gate", d::Norm(defect) <= threshold);
        // A finite failed solve is still a directional observation, never an
        // accepted Newton step or an equal-accuracy order comparison.
        d::Save(file("base.txt"), base);
        d::Save(file("physical_base.txt"), physicalBase);
        d::Save(file("residual.txt"), residual);
        d::Save(file("direction.txt"), direction);
        d::Save(file("action.txt"), action);
        d::Save(file("defect.txt"), defect);
        d::Save(file("state_factors.txt"), stateFactors);
        d::Save(file("residual_factors.txt"), residualFactors);
        state.normalizedOperator->DenormalizeState(direction, physicalDirection);
        d::Save(file("physical_direction.txt"), physicalDirection);

        auto split = d::Open(file("blocks.csv"));
        split << "column,value_block,value_offset,value_size,residual_block,residual_offset,residual_size,direction_norm,action_norm\n";
        mfem::Vector sum(size), blockDirection(size), blockAction(size);
        sum = 0.;
        int column = 0;
        for (const auto &b : values) {
            blockDirection = 0.;
            for (int i = b.offset; i < b.offset + b.size; ++i) blockDirection(i) = direction(i);
            apply(blockDirection, blockAction);
            sum += blockAction;
            d::Save(file("action_b" + std::to_string(column) + ".txt"), blockAction);
            for (const auto &r : rows)
                split << column << ',' << b.stableId << ',' << b.offset << ',' << b.size << ',' << r.stableId << ','
                    << r.offset << ',' << r.size << ',' << d::Norm(blockDirection) << ','
                    << d::Norm(blockAction, r.offset, r.size) << '\n';
            ++column;
        }
        metric("split_columns", column);
        metric("split_action_sum_error", d::Difference(sum, action));
        split.flush();
        MeasureShapeFieldPhysical(state, fem, physicalBase, physicalDirection, output);

        auto trials = d::Open(file("trials.csv"));
        trials << "trial,alpha,prepared,reason,stage,residual_norm,remainder_norm,remainder_over_alpha,linear_prediction_norm\n";
        int trial = 0;
        for (const double alpha : {1./128., 1./512., 1./2048.}) {
            mfem::Vector candidate(base), prediction(residual);
            candidate.Add(alpha, direction);
            prediction.Add(alpha, action);
            const auto prepared = prepare(candidate);
            d::Save(file("candidate_a" + std::to_string(trial) + ".txt"), candidate);
            const double nan = std::numeric_limits<double>::quiet_NaN();
            double residualNorm = nan, remainderNorm = nan;
            if (prepared) {
                if (!d::Finite(state.trialNormalizedResidual)) throw std::runtime_error("Nonfinite trial residual.");
                mfem::Vector remainder(state.trialNormalizedResidual);
                remainder -= prediction;
                residualNorm = d::Norm(state.trialNormalizedResidual);
                remainderNorm = d::Norm(remainder);
                d::Save(file("residual_a" + std::to_string(trial) + ".txt"), state.trialNormalizedResidual);
            }
            trials << trial++ << ',' << alpha << ',' << bool(prepared) << ','
                << (prepared ? -1 : int(prepared.error().reason)) << ','
                << (prepared ? -1 : int(prepared.error().stage)) << ',' << residualNorm << ',' << remainderNorm << ','
                << remainderNorm / alpha << ',' << d::Norm(prediction) << '\n';
            trials.flush();
        }
        restore();
        mfem::Vector restoredAction(size);
        apply(direction, restoredAction);
        d::Save(file("restored_action.txt"), restoredAction);
        metric("restored_action_error", d::Difference(action, restoredAction));
        if (!d::Identical(action, restoredAction)) throw std::runtime_error("Frozen Jacobian action changed.");
        finish();
        metric("probe_exception", 0);
        summary.flush();
    } catch (...) {
        const auto error = std::current_exception();
        metric("probe_exception", 1);
        if (!finalized) finish();
        summary.flush();
        std::rethrow_exception(error);
    }
}
} // namespace experiment::plateau_diagnostics