#pragma once

#include <array>
#include <chrono>
#include <exception>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <limits>
#include <span>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <vector>

// Include after import mean_field. Frozen probes never substitute their base
// into accepted buffers: the context's original seed supplies the safe anchor.
namespace experiment::plateau_diagnostics {
namespace polynomial_probe_detail {
inline const char *Status(mean_field::solver::PolynomialSignStatus status) {
    using S = mean_field::solver::PolynomialSignStatus;
    switch (status) {
    case S::certified_nonnegative: return "certified_nonnegative";
    case S::unresolved: return "unresolved";
    case S::negative: return "negative";
    case S::non_finite: return "non_finite";
    }
    return "unknown";
}
inline const char *RepairStatus(mean_field::solver::EnthalpyDirectionRepairStatus status) {
    using S = mean_field::solver::EnthalpyDirectionRepairStatus;
    switch (status) {
    case S::feasible: return "feasible";
    case S::invalid_base: return "invalid_base";
    case S::invalid_anchor: return "invalid_anchor";
    case S::invalid_direction: return "invalid_direction";
    case S::unresolved: return "unresolved";
    case S::no_descent: return "no_descent";
    }
    return "unknown";
}
inline std::ofstream Open(const std::filesystem::path &path) {
    if (std::filesystem::exists(path)) throw std::runtime_error("Refusing to overwrite positivity probe output.");
    std::ofstream stream(path);
    stream.exceptions(std::ios::badbit | std::ios::failbit);
    stream << std::setprecision(17);
    return stream;
}
inline void Save(const std::filesystem::path &path,const mfem::Vector &values) {
    auto stream = Open(path); stream << values.Size() << '\n';
    for (int i = 0; i < values.Size(); ++i) stream << values(i) << '\n';
}
inline double Difference(const mfem::Vector &a,const mfem::Vector &b) {
    if (a.Size() != b.Size()) return std::numeric_limits<double>::infinity();
    mfem::Vector delta(a); delta -= b; return delta.Norml2();
}
}

template <class State>
void ProbePolynomialEnthalpy(State &s,const mean_field::fem::FEM &f,
                            const mfem::Vector &normalizedBase,const mfem::Vector &normalizedDirection,
                            const std::filesystem::path &output) {
    using namespace mean_field;
    namespace detail = polynomial_probe_detail;
    using Schema = utils::domain::CoreEnvelopeVacuumDomainSchema;
    using Domain = field::FieldDomainT<field::Enthalpy>;
    int ranks = 0;
    if (MPI_Comm_size(s.Problem().GetCommunicator(),&ranks) != MPI_SUCCESS || ranks != 1)
        throw std::invalid_argument("Frozen polynomial positivity diagnostics require one rank.");
    const auto start = std::chrono::steady_clock::now();
    const mfem::Vector acceptedPhysical(s.AcceptedPhysicalState()), acceptedNormalized(s.acceptedNormalizedState),
        acceptedResidual(s.acceptedNormalizedResidual), savedCorrection(s.normalizedCorrection),
        base(normalizedBase), originalDirection(normalizedDirection);
    int preparations = 0;
    auto restore = [&]() {
        s.trialNormalizedState = base;
        ++preparations;
        if (!s.PrepareTrial()) throw std::runtime_error("Frozen positivity base failed production preparation.");
    };
    auto finish = [&]() {
        restore(); s.physicalInverse->Refresh(); s.normalizedCorrection = savedCorrection;
    };
    try {
        restore();
        const mfem::Vector baseResidual(s.trialNormalizedResidual);
        using Problem = std::remove_cvref_t<decltype(s.Problem())>;
        using Form = typename Problem::FormType;
        constexpr int hIndex = utils::blocks::type_index_v<utils::blocks::enthalpy::specific::value,typename Form::value_blocks>;
        constexpr utils::blocks::value_block<hIndex> block{};
        const int offset = s.Problem().GetManifest().layout().offset(block);
        const int size = s.Problem().GetManifest().layout().size(block);
        std::vector<int> elements;
        for (int e = 0; e < f.mesh->GetNE(); ++e)
            if (Schema::template attribute_belongs_to<Domain>(f.mesh->GetAttribute(e))) elements.push_back(e);
        const auto &surface = s.Problem().GetPressureSurfaceRows().reduced_dofs();
        solver::PreparedEnthalpyAdmissibility prepared(*f.enthalpyFes,
            field::make_field_dof_map<field::Enthalpy,Schema>(*f.enthalpyFes),elements,
            std::span<const int>(surface.GetData(),surface.Size()));
        mfem::Vector physicalBase,physicalDirection,h(size),p(size),anchor(size);
        s.normalizedOperator->DenormalizeState(base,physicalBase);
        s.normalizedOperator->DenormalizeState(originalDirection,physicalDirection);
        for (int i = 0; i < size; ++i) {
            h(i) = physicalBase(offset+i); p(i) = physicalDirection(offset+i);
            anchor(i) = acceptedPhysical(offset+i);
        }
        auto certificates = detail::Open(output / "polynomial_certificates.csv");
        certificates << "state,status,lower_bound,upper_bound,visited_subcells,maximum_depth,witness_element,x,y,z,seconds\n";
        auto check = [&](const char *name,const mfem::Vector &values) {
            const auto before = std::chrono::steady_clock::now();
            const auto report = prepared.Check(values);
            certificates << name << ',' << detail::Status(report.status) << ',' << report.lowerBound << ','
                << report.upperBound << ',' << report.visitedSubcells << ',' << report.maximumDepth << ','
                << report.witnessElement << ',' << report.witnessReferencePoint[0] << ','
                << report.witnessReferencePoint[1] << ',' << report.witnessReferencePoint[2] << ','
                << std::chrono::duration<double>(std::chrono::steady_clock::now()-before).count() << '\n';
            certificates.flush(); return report;
        };
        check("frozen_base",h); check("context_seed_anchor",anchor);
        mfem::Vector fullCandidate(h); fullCandidate += p;
        check("original_full_candidate",fullCandidate);
        const auto repair = prepared.MinimumRepairFraction(h,anchor,p);
        auto summary = detail::Open(output / "polynomial_repair.csv");
        summary << "status,minimum_safe_fraction,base_status,anchor_status,candidate_status\n"
            << detail::RepairStatus(repair.status) << ',' << repair.minimumSafeFraction << ','
            << detail::Status(repair.base.status) << ',' << detail::Status(repair.anchor.status) << ','
            << detail::Status(repair.candidate.status) << '\n';
        summary.flush();
        auto trials = detail::Open(output / "polynomial_repair_trials.csv");
        trials << "fraction,alpha,certificate,prepared,reason,stage,eos_code,residual_norm,linear_model_norm,merit_slope,direction_change\n";
        auto evaluate = [&](double fraction,double alpha,const mfem::Vector &direction) {
            restore();
            mfem::Vector action(s.Problem().EquationSize()); s.normalizedOperator->Mult(direction,action);
            const double slope = baseResidual*action;
            mfem::Vector model(action); model *= alpha; model += baseResidual;
            mfem::Vector candidate(base); candidate.Add(alpha,direction);
            mfem::Vector physicalCandidate,reduced(size);
            s.normalizedOperator->DenormalizeState(candidate,physicalCandidate);
            for (int i = 0; i < size; ++i) reduced(i) = physicalCandidate(offset+i);
            const auto certificate = prepared.Check(reduced);
            s.trialNormalizedState = candidate; ++preparations;
            const auto preparation = s.PrepareTrial();
            trials << fraction << ',' << alpha << ',' << detail::Status(certificate.status) << ','
                << static_cast<bool>(preparation) << ',';
            if (preparation) trials << "-1,-1,-1," << s.trialNormalizedResidual.Norml2();
            else trials << int(preparation.error().reason) << ',' << int(preparation.error().stage) << ','
                << int(preparation.error().thermodynamicErrorCode) << ",nan";
            trials << ',' << model.Norml2() << ',' << slope << ','
                << detail::Difference(direction,originalDirection) << '\n'; trials.flush();
        };
        evaluate(0.0,1.0,originalDirection);
        if (repair.status == solver::EnthalpyDirectionRepairStatus::feasible) {
            mfem::Vector repairedH;
            const std::array fractions{repair.minimumSafeFraction,
                0.5*(repair.minimumSafeFraction+1.0),1.0};
            for (std::size_t i = 0; i < fractions.size(); ++i) {
                const double fraction = fractions[i];
                const auto report = prepared.ApplyRepairFraction(h,anchor,p,fraction,repairedH);
                if (report.status != solver::PolynomialSignStatus::certified_nonnegative) continue;
                mfem::Vector repairedPhysical(physicalDirection),repairedNormalized;
                for (int j = 0; j < size; ++j) repairedPhysical(offset+j) = repairedH(j);
                s.normalizedOperator->NormalizeState(repairedPhysical,repairedNormalized);
                for (int j = 0; j < repairedNormalized.Size(); ++j)
                    if (j < offset || j >= offset+size) repairedNormalized(j) = originalDirection(j);
                detail::Save(output / ("polynomial_repaired_direction_"+std::to_string(i)+".txt"),repairedNormalized);
                for (double alpha : {1.0,0.5,0.125}) evaluate(fraction,alpha,repairedNormalized);
            }
        }
        finish();
        auto checks = detail::Open(output / "polynomial_restoration.csv");
        checks << "metric,value\naccepted_physical_difference," << detail::Difference(acceptedPhysical,s.AcceptedPhysicalState())
            << "\naccepted_normalized_difference," << detail::Difference(acceptedNormalized,s.acceptedNormalizedState)
            << "\naccepted_residual_difference," << detail::Difference(acceptedResidual,s.acceptedNormalizedResidual)
            << "\nrestored_residual_difference," << detail::Difference(baseResidual,s.trialNormalizedResidual)
            << "\nrestored_correction_difference," << detail::Difference(savedCorrection,s.normalizedCorrection)
            << "\npreparations," << preparations << "\nseconds,"
            << std::chrono::duration<double>(std::chrono::steady_clock::now()-start).count() << '\n';
        if (detail::Difference(acceptedPhysical,s.AcceptedPhysicalState()) != 0.0 ||
            detail::Difference(acceptedNormalized,s.acceptedNormalizedState) != 0.0 ||
            detail::Difference(acceptedResidual,s.acceptedNormalizedResidual) != 0.0 ||
            detail::Difference(baseResidual,s.trialNormalizedResidual) != 0.0)
            throw std::runtime_error("Polynomial positivity diagnostics changed an accepted state or failed restoration.");
    } catch (...) {
        const auto error = std::current_exception();
        finish(); std::rethrow_exception(error);
    }
}
}