#pragma once

// Include after the benchmark and physical-snapshot helpers. Each observation
// starts from the same prepared physical state; nothing is accepted or clipped.
namespace experiment::plateau_diagnostics {
template <class State>
void RecordHarmonicLinearTrials(State &state, const mean_field::fem::FEM &fem,
                               const std::filesystem::path &output, double index) {
    namespace d = preconditioner_benchmark_detail;
    const mfem::Vector base(state.trialNormalizedState), residual(state.trialNormalizedResidual),
        savedCorrection(state.normalizedCorrection), accepted(state.AcceptedPhysicalState());
    mfem::Vector physicalBase;
    state.normalizedOperator->DenormalizeState(base, physicalBase);
    const auto restore = [&] {
        state.trialNormalizedState = base;
        if (!state.PrepareTrial()) throw std::runtime_error("Linear trial base cannot be restored");
    };
    auto trials = d::Open(output / "linear_direction_trials.csv");
    trials << "role,alpha,prepared,reason,stage,base_norm,trial_norm,relative_norm,linear_defect_relative,"
              "direction_norm,action_norm,merit_slope,remainder_norm,remainder_over_alpha,"
              "predicted_norm,signed_remainder_projection,merit_identity_error,decrease_screen\n";
    auto blocks = d::Open(output / "linear_direction_trial_blocks.csv");
    blocks << "role,alpha,block,base_norm,action_norm,trial_norm,remainder_norm\n";
    auto bounds = d::Open(output / "linear_direction_enthalpy_bounds.csv");
    bounds << "role,metric,value\n";
    for (const std::string role : {"baseline", "candidate"}) {
        restore();
        const auto correction = coupled_physical_detail::Load(output / (role + "_normalized_correction.txt"), base.Size());
        mfem::Vector action(base.Size()), defect(base.Size());
        state.normalizedOperator->Mult(correction, action);
        defect = action; defect += residual;
        const double relative = defect.Norml2() / residual.Norml2();
        const double slope = residual * action;
        state.normalizedCorrection = correction;
        for (const auto &[key, value] : MeasureEnthalpyStep(state, fem, index, &physicalBase))
            bounds << role << ',' << key << ',' << value << '\n';
        d::Save(output / (role + "_trial_jacobian_action.txt"), action);
        for (const double alpha : {1., .25, .0625, .015625, .00390625, .0009765625, .000244140625}) {
            restore();
            state.trialNormalizedState = base;
            state.trialNormalizedState.Add(alpha, correction);
            const auto prepared = state.PrepareTrial();
            trials << role << ',' << alpha << ',' << prepared.has_value() << ','
                << (prepared ? -1 : int(prepared.error().reason)) << ','
                << (prepared ? -1 : int(prepared.error().stage)) << ',' << residual.Norml2() << ',';
            if (!prepared) {
                trials << "nan,nan," << relative << ',' << correction.Norml2() << ',' << action.Norml2()
                       << ',' << slope << ",nan,nan,nan,nan,nan,0\n";
                continue;
            }
            const mfem::Vector trial(state.trialNormalizedResidual);
            mfem::Vector prediction(residual), remainder(trial);
            prediction.Add(alpha, action);
            remainder -= prediction;
            const double projection = prediction * remainder;
            const double identity = trial * trial - prediction * prediction - 2. * projection - remainder * remainder;
            trials << trial.Norml2() << ',' << trial.Norml2() / residual.Norml2() << ',' << relative << ','
                << correction.Norml2() << ',' << action.Norml2() << ',' << slope << ',' << remainder.Norml2()
                << ',' << remainder.Norml2() / alpha << ',' << prediction.Norml2() << ',' << projection << ','
                << identity << ',' << (trial * trial <= (1. - 1e-4 * alpha) * (residual * residual)) << '\n';
            for (const auto &block : state.Problem().GetManifest().residualBlocks())
                blocks << role << ',' << alpha << ',' << block.stableId << ','
                    << d::BlockNorm(residual, block.offset, block.size) << ','
                    << d::BlockNorm(action, block.offset, block.size) << ','
                    << d::BlockNorm(trial, block.offset, block.size) << ','
                    << d::BlockNorm(remainder, block.offset, block.size) << '\n';
        }
        trials.flush(); blocks.flush(); bounds.flush();
    }
    restore();
    state.normalizedCorrection = savedCorrection;
    if (!d::Identical(accepted, state.AcceptedPhysicalState()) ||
        !d::Identical(residual, state.trialNormalizedResidual))
        throw std::runtime_error("Linear trial observations altered their base or accepted state");
    auto invariants = d::Open(output / "linear_direction_trial_invariants.csv");
    invariants << "metric,value\naccepted_unchanged,1\nbase_residual_restored,1\n";
}
} // namespace experiment::plateau_diagnostics