#pragma once

#include <chrono>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <stdexcept>

// Include after import mean_field. This callback only reads the supplied state;
// its WithState owner performs the usual accepted-state restoration afterward.
namespace experiment::plateau_diagnostics {
namespace polynomial_certificate_detail {
inline const char *Status(mean_field::solver::PolynomialSignStatus status) {
    using Status = mean_field::solver::PolynomialSignStatus;
    switch (status) {
    case Status::certified_nonnegative: return "certified_nonnegative";
    case Status::unresolved: return "unresolved";
    case Status::negative: return "negative";
    case Status::non_finite: return "non_finite";
    }
    return "unknown";
}
inline std::ofstream Open(const std::filesystem::path &path) {
    if (std::filesystem::exists(path))
        throw std::runtime_error("Refusing to overwrite a polynomial endpoint certificate.");
    std::ofstream stream(path);
    stream.exceptions(std::ios::badbit | std::ios::failbit);
    stream << std::setprecision(17);
    return stream;
}
inline bool Identical(const mfem::Vector &left,const mfem::Vector &right) {
    return left.Size() == right.Size() && (left.Size() == 0 ||
        std::memcmp(left.GetData(),right.GetData(),static_cast<std::size_t>(left.Size())*sizeof(double)) == 0);
}
}

template <class State>
void CertifyPolynomialEnthalpy(State &state,const mfem::Vector &fullNormalized,
                              const std::filesystem::path &output) {
    namespace detail = polynomial_certificate_detail;
    using Clock = std::chrono::steady_clock;
    int ranks = 0;
    if (MPI_Comm_size(state.Problem().GetCommunicator(),&ranks) != MPI_SUCCESS || ranks != 1)
        throw std::invalid_argument("Polynomial endpoint output requires exactly one MPI rank.");
    const mfem::Vector acceptedPhysical(state.AcceptedPhysicalState()),
        acceptedNormalized(state.acceptedNormalizedState), acceptedResidual(state.acceptedNormalizedResidual),
        correction(state.normalizedCorrection), trialState(state.trialNormalizedState),
        trialResidual(state.trialNormalizedResidual), candidatePhysical(state.candidatePhysicalState);
    const auto generation = state.Problem().GetPreparationGeneration();
    const auto start = Clock::now();
    mean_field::solver::PreparedStellarEnthalpyStrategy<State> prepared(state);
    const auto certificateStart = Clock::now();
    const auto report = prepared.CheckState(fullNormalized);
    const double checkSeconds = std::chrono::duration<double>(Clock::now()-certificateStart).count();
    const double workspaceSeconds = std::chrono::duration<double>(certificateStart-start).count();
    auto certificate = detail::Open(output / "polynomial_endpoint_certificate.csv");
    certificate << "status,lower_bound,upper_bound,visited_subcells,maximum_depth,witness_element,x,y,z,certificate_seconds,workspace_seconds\n"
        << detail::Status(report.status) << ',' << report.lowerBound << ',' << report.upperBound << ','
        << report.visitedSubcells << ',' << report.maximumDepth << ',' << report.witnessElement << ','
        << report.witnessReferencePoint[0] << ',' << report.witnessReferencePoint[1] << ','
        << report.witnessReferencePoint[2] << ',' << checkSeconds << ',' << workspaceSeconds << '\n';
    const bool sameAcceptedPhysical = detail::Identical(acceptedPhysical,state.AcceptedPhysicalState());
    const bool sameAcceptedNormalized = detail::Identical(acceptedNormalized,state.acceptedNormalizedState);
    const bool sameAcceptedResidual = detail::Identical(acceptedResidual,state.acceptedNormalizedResidual);
    const bool sameCorrection = detail::Identical(correction,state.normalizedCorrection);
    const bool sameTrialState = detail::Identical(trialState,state.trialNormalizedState);
    const bool sameTrialResidual = detail::Identical(trialResidual,state.trialNormalizedResidual);
    const bool sameCandidate = detail::Identical(candidatePhysical,state.candidatePhysicalState);
    const bool samePreparation = generation == state.Problem().GetPreparationGeneration();
    auto invariants = detail::Open(output / "polynomial_endpoint_invariants.csv");
    invariants << "metric,value\naccepted_physical_bitwise_identical," << sameAcceptedPhysical
        << "\naccepted_normalized_bitwise_identical," << sameAcceptedNormalized
        << "\naccepted_residual_bitwise_identical," << sameAcceptedResidual
        << "\ncorrection_bitwise_identical," << sameCorrection
        << "\ntrial_state_bitwise_identical," << sameTrialState
        << "\ntrial_residual_bitwise_identical," << sameTrialResidual
        << "\ncandidate_physical_bitwise_identical," << sameCandidate
        << "\npreparation_generation_unchanged," << samePreparation
        << "\nhelper_trial_preparations,0\nhelper_jacobian_applications,0\nhelper_linear_solves,0\n";
    auto scope = detail::Open(output / "polynomial_endpoint_scope.txt");
    scope << "The helper constructs a positivity workspace and certifies only the supplied normalized state.\n"
          << "It does not prepare that state, apply a Jacobian, repair a direction, or solve any equations.\n"
          << "The existing WithState wrapper restores the accepted context after this callback returns;\n"
          << "that standard restoration is outside the helper counters and timings reported here.\n"
          << "Bounds are polynomial enclosures, not measured minimum/maximum values. Witness coordinates\n"
          << "identify a proved negative point when status is negative; unresolved is never a certificate.\n";
    if (!(sameAcceptedPhysical && sameAcceptedNormalized && sameAcceptedResidual && sameCorrection &&
          sameTrialState && sameTrialResidual && sameCandidate && samePreparation))
        throw std::runtime_error("Polynomial endpoint certification changed a context buffer or preparation.");
}
}