#pragma once

// A frozen, affine surface restriction tests approximation-space compatibility
// without changing the production problem or committing a filtered model.
#include "plateau_nonlinear_step_diagnostics.hpp"
#include "plateau_coupled_trial_physical.hpp"
#include "shape_field_physical_diagnostics.hpp"
#include "reduced_surface_basis.hpp"
#include "reduced_surface_operators.hpp"
#include <memory>

namespace experiment::plateau_diagnostics {
namespace reduced_surface_probe_detail {
inline void SaveMatrix(const std::filesystem::path &path, const mfem::DenseMatrix &matrix) {
    auto file = nonlinear_step_detail::Open(path);
    file << matrix.Height() << ' ' << matrix.Width() << '\n';
    for (int i = 0; i < matrix.Height(); ++i)
        for (int j = 0; j < matrix.Width(); ++j) {
            if (!std::isfinite(matrix(i, j))) throw std::runtime_error("Nonfinite reduced surface matrix.");
            file << matrix(i, j) << '\n';
        }
}
} // namespace reduced_surface_probe_detail

template <class State>
void ProbeReducedSurface(State &state, const mean_field::fem::FEM &fem,
                         const mfem::Vector &normalizedBase, const std::filesystem::path &output,
                         int maximumDegree, double n, double K, NonlinearStepProbeOptions options = {}) {
    using namespace mean_field;
    namespace d = nonlinear_step_detail;
    namespace pd = reduced_surface_probe_detail;
    const auto start = d::Clock::now();
    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)
        throw std::invalid_argument("Reduced surface probes require one rank and finite square root data.");
    const auto algebra = RunReducedSurfaceAlgebraSelfCheck(state.Problem().GetCommunicator());
    const bool identity = maximumDegree == 0;
    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);
    const double savedMinimumJacobian = state.acceptedMinimumJacobianDeterminant;
    const auto values = state.Problem().GetManifest().valueBlocks();
    const auto rows = state.Problem().GetManifest().residualBlocks();
    const auto shapeValue = shape_field_physical_detail::RequireBlock(values, "surface_deformation", size);
    const auto shapeRow = shape_field_physical_detail::RequireBlock(rows, "surface_shape_balance", size);
    if (shapeValue.size != shapeRow.size) throw std::invalid_argument("Shape dimensions disagree.");
    std::filesystem::create_directories(output);
    const auto file = [&](const std::string &name) { return output / ("reduced_surface_" + 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 layout = d::Open(file("layout.csv"));
    layout << "kind,block,offset,size\n";
    for (const auto &b : values) layout << "value," << b.stableId << ',' << b.offset << ',' << b.size << '\n';
    for (const auto &b : rows) layout << "residual," << b.stableId << ',' << b.offset << ',' << b.size << '\n';
    layout.flush();
    mfem::Vector residual(size), direction(size), action(size), physicalBase, physicalDirection;
    bool haveBase = false, finalized = false;
    int preparations = 0;
    double restoreError = 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("Reduced surface probe cannot restore base.");
        if (!haveBase) { residual = state.trialNormalizedResidual; haveBase = true; }
        if (!d::Finite(state.trialNormalizedResidual)) throw std::runtime_error("Nonfinite base residual.");
        restoreError = std::max(restoreError, d::Difference(residual, state.trialNormalizedResidual));
    };
    const auto finish = [&]() {
        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()) &&
            savedMinimumJacobian == state.acceptedMinimumJacobianDeterminant;
        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("restored_residual_error", restoreError);
        d::Save(file("restored_residual.txt"), state.trialNormalizedResidual);
        metric("preparations", preparations);
        metric("total_probe_seconds", d::Seconds(start));
        summary.flush();
        if (!acceptedSame || !factorsSame || restoreError != 0.)
            throw std::runtime_error("Reduced surface probe changed frozen invariants.");
    };
    try {
        restore();
        state.physicalInverse->Refresh();
        physicalBase = state.normalizedOperator->GetPhysicalState();
        metric("schema_version", 1);
        metric("maximum_degree", maximumDegree);
        metric("identity_control", identity);
        metric("full_size", size);
        metric("shape_size", shapeValue.size);
        metric("shape_value_offset", shapeValue.offset);
        metric("shape_residual_offset", shapeRow.offset);
        metric("linear_relative_tolerance", options.linear.relativeTolerance);
        metric("linear_absolute_tolerance", options.linear.absoluteTolerance);
        metric("linear_maximum_iterations", options.linear.maximumIterations);
        metric("linear_restart", options.restartLength);
        metric("accepted_commits", 0);
        metric("algebra_self_check_passed", 1);
        metric("algebra_self_check_count", algebra.checks);
        metric("algebra_self_check_maximum_error", algebra.maximumError);

        std::unique_ptr<ReducedSurfaceCoordinates> coordinates;
        if (!identity) {
            auto basis = MakeReducedSurfaceBasis(state, fem, maximumDegree, output);
            coordinates = std::make_unique<ReducedSurfaceCoordinates>(basis.physicalBasis,
                stateFactors, residualFactors, shapeValue.offset, shapeRow.offset, state.Problem().GetCommunicator());
            pd::SaveMatrix(file("state_basis.txt"), coordinates->StateBasis());
            pd::SaveMatrix(file("residual_basis.txt"), coordinates->ResidualBasis());
            pd::SaveMatrix(file("state_transform.txt"), coordinates->StateTransform());
            pd::SaveMatrix(file("residual_transform.txt"), coordinates->ResidualTransform());
        }
        const int reducedSize = identity ? size : coordinates->ReducedSize();
        metric("mode_count", identity ? shapeValue.size : coordinates->ModeCount());
        metric("reduced_size", reducedSize);
        const auto restrictResidual = [&](const mfem::Vector &full, mfem::Vector &reduced) {
            if (identity) reduced = full; else coordinates->RestrictResidual(full, reduced);
        };
        const auto excluded = [&](const mfem::Vector &full, mfem::Vector &result) {
            mfem::Vector reduced, retained;
            restrictResidual(full, reduced);
            if (identity) retained = reduced; else coordinates->InjectResidual(reduced, retained);
            result = full; result -= retained;
        };
        mfem::Vector reducedResidual, reducedDirection(reducedSize), reducedAction, reducedDefect;
        restrictResidual(residual, reducedResidual);
        reducedDirection = 0.;
        mfem::Vector rhs(reducedResidual); rhs *= -1.;
        const auto solve = [&](auto &operation, auto &inverse) {
            auto backend = solver::linear::prepareLinearBackend(
                solver::linear::FGMRES({.restartLength = options.restartLength, .printLevel = -1}),
                operation, inverse, state.Problem().GetCommunicator());
            const auto result = backend.Solve(rhs, reducedDirection, options.linear);
            metric("linear_status", int(result.status));
            metric("linear_iterations", result.iterations);
            metric("linear_seconds", result.solveSeconds);
            metric("linear_reported_true_norm", result.trueResidualNorm);
        };
        if (identity) solve(*state.normalizedOperator, *state.normalizedInverse);
        else {
            ReducedSurfaceOperator operation(*state.normalizedOperator, *coordinates);
            ReducedSurfacePreconditioner inverse(*state.normalizedInverse, *coordinates);
            solve(operation, inverse);
        }
        restore();
        if (identity) direction = reducedDirection; else coordinates->InjectState(reducedDirection, direction);
        if (!d::Finite(direction) || d::Norm(direction) == 0.) throw std::runtime_error("Invalid reduced direction.");
        state.normalizedOperator->Mult(direction, action);
        if (!d::Finite(action)) throw std::runtime_error("Nonfinite full Jacobian action.");
        restrictResidual(action, reducedAction);
        mfem::Vector defect(action); defect += residual;
        reducedDefect = reducedAction; reducedDefect += reducedResidual;
        mfem::Vector excludedResidual, excludedDefect;
        excluded(residual, excludedResidual); excluded(defect, excludedDefect);
        const double fullBaseNorm = d::Norm(residual), reducedBaseNorm = d::Norm(reducedResidual);
        for (double norm : {fullBaseNorm, reducedBaseNorm, d::Norm(defect), d::Norm(reducedDefect),
                            d::Norm(direction), d::Norm(excludedResidual), d::Norm(excludedDefect)})
            if (!std::isfinite(norm)) throw std::runtime_error("Nonfinite reduced diagnostic norm.");
        if (!(fullBaseNorm > 0.) || !(reducedBaseNorm > 0.))
            throw std::runtime_error("This comparison requires nonzero full and retained residuals.");
        metric("full_base_norm", fullBaseNorm);
        metric("reduced_base_norm", reducedBaseNorm);
        metric("excluded_base_norm", d::Norm(excludedResidual));
        metric("full_relative_defect", d::Ratio(d::Norm(defect), fullBaseNorm));
        metric("reduced_relative_defect", d::Ratio(d::Norm(reducedDefect), reducedBaseNorm));
        metric("full_verified_gate", d::Norm(defect) <= options.linear.ConvergenceThreshold(fullBaseNorm));
        metric("reduced_verified_gate", d::Norm(reducedDefect) <= options.linear.ConvergenceThreshold(reducedBaseNorm));
        metric("full_merit_slope", d::Dot(residual, action, 0, size));
        metric("reduced_merit_slope", d::Dot(reducedResidual, reducedAction, 0, reducedSize));
        metric("direction_norm", d::Norm(direction));
        state.normalizedOperator->DenormalizeState(direction, physicalDirection);
        for (const auto &[name, vector] : std::initializer_list<std::pair<const char *, const mfem::Vector *>>{
                 {"base", &base}, {"physical_base", &physicalBase}, {"residual", &residual},
                 {"direction", &direction}, {"physical_direction", &physicalDirection}, {"action", &action},
                 {"defect", &defect}, {"state_factors", &stateFactors}, {"residual_factors", &residualFactors},
                 {"reduced_residual", &reducedResidual}, {"reduced_direction", &reducedDirection},
                 {"reduced_action", &reducedAction}, {"reduced_defect", &reducedDefect},
                 {"excluded_residual", &excludedResidual}, {"excluded_defect", &excludedDefect}})
            d::Save(file(std::string(name) + ".txt"), *vector);
        summary.flush();
        std::cout << "Reduced surface degree " << maximumDegree << ": linear iterations complete; full relative defect "
                  << d::Ratio(d::Norm(defect), fullBaseNorm) << ", reduced "
                  << d::Ratio(d::Norm(reducedDefect), reducedBaseNorm) << std::endl;
        MeasureShapeFieldPhysical(state, fem, physicalBase, physicalDirection, output);
        auto physical = d::Open(file("physical.csv"));
        physical << "sample,quadrature,metric,value\n";
        const auto measurePhysical = [&](int sample) {
            coupled_physical_detail::PreparedSnapshotView<State> view(state);
            experiment::polytrope_validation::PhysicalState snapshot(view, fem);
            for (int q : {14, 18})
                for (const auto &[key, value] : MeasurePhysical(snapshot, n, K, q))
                    physical << sample << ',' << q << ',' << key << ',' << value << '\n';
            physical.flush();
        };
        measurePhysical(-1);
        auto trials = d::Open(file("trials.csv"));
        trials << "sample,alpha,prepared,reason,stage,full_norm,reduced_norm,excluded_shape_norm,"
                  "full_merit_pass,reduced_merit_pass,full_remainder_norm,reduced_remainder_norm\n";
        bool measuredTrial = false;
        int sample = 0;
        for (double alpha : {1., .25, .0625, 1./1024.}) {
            // The affine base is retained verbatim: no projection of q0 and no
            // enthalpy repair, SOC, or accepted-state update belongs in this test.
            mfem::Vector candidate(base); candidate.Add(alpha, direction);
            const auto prepared = prepare(candidate);
            d::Save(file("candidate_a" + std::to_string(sample) + ".txt"), candidate);
            if (prepared) {
                const mfem::Vector trial(state.trialNormalizedResidual);
                if (!d::Finite(trial)) throw std::runtime_error("Nonfinite reduced trial residual.");
                mfem::Vector reducedTrial, omitted, remainder(trial), reducedRemainder;
                restrictResidual(trial, reducedTrial); excluded(trial, omitted);
                remainder -= residual; remainder.Add(-alpha, action);
                restrictResidual(remainder, reducedRemainder);
                const double fullNorm = d::Norm(trial), reducedNorm = d::Norm(reducedTrial);
                trials << sample << ',' << alpha << ",1,-1,-1," << fullNorm << ',' << reducedNorm << ','
                       << d::Norm(omitted) << ',' << (fullNorm <= (1.-1e-4*alpha)*fullBaseNorm) << ','
                       << (reducedNorm <= (1.-1e-4*alpha)*reducedBaseNorm) << ','
                       << d::Norm(remainder) << ',' << d::Norm(reducedRemainder) << '\n';
                d::Save(file("residual_a" + std::to_string(sample) + ".txt"), trial);
                d::Save(file("reduced_residual_a" + std::to_string(sample) + ".txt"), reducedTrial);
                mfem::Vector physicalCandidate;
                state.normalizedOperator->DenormalizeState(candidate, physicalCandidate);
                d::Save(file("physical_candidate_a" + std::to_string(sample) + ".txt"), physicalCandidate);
                if (!measuredTrial) { measurePhysical(sample); measuredTrial = true; }
            } else {
                trials << sample << ',' << alpha << ",0," << int(prepared.error().reason) << ','
                       << int(prepared.error().stage) << ",nan,nan,nan,0,0,nan,nan\n";
            }
            trials.flush();
            ++sample;
        }
        restore();
        mfem::Vector restoredAction(size);
        state.normalizedOperator->Mult(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("Restored Jacobian changed.");
        finish();
        metric("probe_exception", 0);
        summary.flush();
    } catch (...) {
        const auto failure = std::current_exception();
        metric("probe_exception", 1);
        if (!finalized) finish();
        std::rethrow_exception(failure);
    }
}
} // namespace experiment::plateau_diagnostics