#pragma once

// Include after import mean_field. Physical replay is deliberately separate
// from timed trial selection: these quadratures do not influence acceptance.
#include "plateau_physical_diagnostics.hpp"

#include <cstring>
#include <exception>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <type_traits>
#include <vector>

namespace experiment::plateau_diagnostics {
namespace coupled_physical_detail {
inline bool Finite(const mfem::Vector &values) {
    for (int i = 0; i < values.Size(); ++i)
        if (!std::isfinite(values(i))) return false;
    return true;
}

inline bool Identical(const mfem::Vector &a, const mfem::Vector &b) {
    return a.Size() == b.Size() &&
        (a.Size() == 0 || std::memcmp(a.GetData(), b.GetData(), sizeof(mfem::real_t) * a.Size()) == 0);
}

inline mfem::Vector Load(const std::filesystem::path &path, int size) {
    std::ifstream input(path);
    long long storedSize = -1;
    if (!(input >> storedSize) || storedSize != size)
        throw std::invalid_argument("Coupled physical snapshot has an incompatible size: " + path.string());
    mfem::Vector values(size);
    for (int i = 0; i < size; ++i)
        if (!(input >> values(i)) || !std::isfinite(values(i)))
            throw std::invalid_argument("Coupled physical snapshot is incomplete or nonfinite: " + path.string());
    std::string trailing;
    if (input >> trailing || input.bad())
        throw std::invalid_argument("Coupled physical snapshot has trailing data or a read error: " + path.string());
    return values;
}

// PhysicalState's structural constructor uses these historical accepted-state
// member names. This adapter supplies an explicitly prepared, uncommitted
// snapshot instead; it never aliases a writable accepted vector. Reusing that
// constructor preserves the same reduced-field scatter and Piola transforms.
template <class State>
struct PreparedSnapshotView final {
    State &state;
    decltype(std::declval<State &>().normalizedOperator) &normalizedOperator;
    const mfem::Vector &acceptedNormalizedResidual;

    explicit PreparedSnapshotView(State &source)
        : state(source), normalizedOperator(source.normalizedOperator),
          acceptedNormalizedResidual(source.trialNormalizedResidual) {
        if (!normalizedOperator->IsPrepared() || !Finite(acceptedNormalizedResidual))
            throw std::logic_error("Physical snapshot reconstruction requires a fresh prepared trial.");
    }
    decltype(auto) Problem() const { return state.Problem(); }
    const mfem::Vector &AcceptedPhysicalState() const { return normalizedOperator->GetPhysicalState(); }
};
} // namespace coupled_physical_detail

template <class State>
void MeasureCoupledTrialSnapshots(State &state, const mean_field::fem::FEM &fem,
                                 const std::filesystem::path &output, double n, double K) {
    namespace d = coupled_physical_detail;
    int ranks = 0;
    MPI_Comm_size(state.Problem().GetCommunicator(), &ranks);
    if (ranks != 1 || !std::isfinite(n) || n < 1. || !std::isfinite(K) || K <= 0.)
        throw std::invalid_argument("Coupled physical replay requires one rank, finite n>=1 and finite K>0.");

    constexpr std::string_view prefix = "coupled_trial_selected_physical_a";
    constexpr std::string_view suffix = ".txt";
    std::vector<std::pair<int, std::filesystem::path>> snapshots;
    for (const auto &entry : std::filesystem::directory_iterator(output)) {
        const std::string name = entry.path().filename().string();
        if (!entry.is_regular_file() || !name.starts_with(prefix) || !name.ends_with(suffix)) continue;
        const std::string digits = name.substr(prefix.size(), name.size() - prefix.size() - suffix.size());
        if (digits.empty() || !std::all_of(digits.begin(), digits.end(), [](char ch) { return ch >= '0' && ch <= '9'; }))
            throw std::invalid_argument("Coupled physical snapshot has an invalid sample index: " + name);
        snapshots.emplace_back(std::stoi(digits), entry.path());
    }
    std::sort(snapshots.begin(), snapshots.end(), [](const auto &a, const auto &b) { return a.first < b.first; });
    for (std::size_t i = 1; i < snapshots.size(); ++i)
        if (snapshots[i - 1].first == snapshots[i].first)
            throw std::invalid_argument("Coupled physical replay found duplicate sample indices.");

    const auto path = output / "coupled_trial_physical.csv";
    if (std::filesystem::exists(path))
        throw std::invalid_argument("Coupled physical replay output already exists: " + path.string());
    std::ofstream stream(path);
    stream.exceptions(std::ios::failbit | std::ios::badbit);
    stream << std::setprecision(17) << "sample,source,quadrature,metric,value\n";
    const auto write = [&](int sample, const std::string &source, int quadrature, const std::string &metric, double value) {
        stream << sample << ',' << source << ',' << quadrature << ',' << metric << ',' << value << '\n';
    };
    const mfem::Vector accepted(state.acceptedNormalizedState), acceptedResidual(state.acceptedNormalizedResidual),
        acceptedPhysical(state.AcceptedPhysicalState()), correction(state.normalizedCorrection),
        physicalCorrection(state.physicalCorrection), volumeDirection(state.volumeDisplacementDirection);
    const auto &normalization = state.normalizedOperator->GetNormalization();
    const mfem::Vector stateFactors(normalization.StateFactors()), residualFactors(normalization.ResidualFactors());
    const double minimumJacobian = state.acceptedMinimumJacobianDeterminant;
    const auto acceptedUnchanged = [&] {
        return d::Identical(accepted, state.acceptedNormalizedState) &&
            d::Identical(acceptedResidual, state.acceptedNormalizedResidual) &&
            d::Identical(acceptedPhysical, state.AcceptedPhysicalState()) &&
            minimumJacobian == state.acceptedMinimumJacobianDeterminant;
    };
    bool restorationAttempted = false;
    const auto restore = [&] {
        restorationAttempted = true;
        const bool before = acceptedUnchanged();
        state.RestoreAccepted();
        const bool after = acceptedUnchanged();
        const bool factorsUnchanged = d::Identical(stateFactors, normalization.StateFactors()) &&
            d::Identical(residualFactors, normalization.ResidualFactors());
        const bool directionsUnchanged = d::Identical(correction, state.normalizedCorrection) &&
            d::Identical(physicalCorrection, state.physicalCorrection) &&
            d::Identical(volumeDirection, state.volumeDisplacementDirection);
        const bool ready = state.normalizedOperator->IsPrepared() && state.physicalInverse->IsCurrent() &&
            state.normalizedInverse->IsCurrent();
        write(-1, "restoration", 0, "accepted_buffers_unchanged_before_restoration", before);
        write(-1, "restoration", 0, "accepted_buffers_unchanged_after_restoration", after);
        write(-1, "restoration", 0, "normalization_factors_unchanged", factorsUnchanged);
        write(-1, "restoration", 0, "direction_buffers_unchanged", directionsUnchanged);
        write(-1, "restoration", 0, "accepted_operator_and_inverse_current", ready);
        stream.flush();
        if (!before || !after || !factorsUnchanged || !directionsUnchanged || !ready)
            throw std::runtime_error("Coupled physical replay failed its accepted-state restoration invariant.");
    };

    try {
        const double targetMass = state.Problem().GetStellarModel()
            .template specification<mean_field::models::FixedTotalMass>().targetMass().value();
        for (const auto &[sample, snapshotPath] : snapshots) {
            const std::string source = snapshotPath.filename().string();
            const auto saved = d::Load(snapshotPath, state.Problem().StateSize());
            state.normalizedOperator->NormalizeState(saved, state.trialNormalizedState);
            state.normalizedOperator->DenormalizeState(state.trialNormalizedState, state.candidatePhysicalState);
            const auto preparation = state.PrepareTrial();
            write(sample, source, 0, "fresh_preparation_passed", bool(preparation));
            if (!preparation) {
                write(sample, source, 0, "rejection_reason", int(preparation.error().reason));
                write(sample, source, 0, "rejection_stage", int(preparation.error().stage));
                throw std::runtime_error("Selected coupled trial is inadmissible during physical replay: " + source);
            }
            const auto &prepared = state.normalizedOperator->GetPhysicalState();
            if (prepared.Size() != saved.Size() || !d::Finite(prepared) || !d::Finite(state.trialNormalizedResidual))
                throw std::runtime_error("Selected coupled trial produced invalid prepared physical data: " + source);
            long double squaredDifference = 0.;
            double maximumDifference = 0., maximumScaledDifference = 0.;
            for (int i = 0; i < saved.Size(); ++i) {
                const double difference = std::abs(prepared(i) - saved(i));
                squaredDifference += static_cast<long double>(difference) * difference;
                maximumDifference = std::max(maximumDifference, difference);
                maximumScaledDifference = std::max(maximumScaledDifference, difference / std::max(1., std::abs(saved(i))));
            }
            d::PreparedSnapshotView<State> view(state);
            experiment::polytrope_validation::PhysicalState physical(view, fem);
            for (const int quadrature : {14, 18}) {
                auto measurements = MeasurePhysical(physical, n, K, quadrature);
                measurements["saved_prepared_l2_difference"] = static_cast<double>(std::sqrt(squaredDifference));
                measurements["saved_prepared_maximum_difference"] = maximumDifference;
                measurements["saved_prepared_maximum_scaled_difference"] = maximumScaledDifference;
                measurements["minimum_preparation_jacobian"] = preparation->minimumJacobianDeterminant;
                measurements["normalized_bordered_residual"] = physical.normalizedBorderedResidualNorm;
                measurements["normalized_unbordered_residual"] = physical.normalizedUnborderedResidualNorm;
                measurements["normalized_border_action"] = physical.normalizedCentralBorderActionNorm;
                measurements["central_border"] = physical.centralBorder;
                measurements["target_mass"] = targetMass;
                measurements["mass_absolute_error"] = std::abs(measurements.at("mass") - targetMass);
                for (const auto &[metric, value] : measurements) write(sample, source, quadrature, metric, value);
                stream.flush();
            }
        }
        write(-1, "replay", 0, "selected_snapshots_measured", static_cast<double>(snapshots.size()));
        restore();
    } catch (...) {
        const auto failure = std::current_exception();
        if (!restorationAttempted) restore();
        std::rethrow_exception(failure);
    }
}
} // namespace experiment::plateau_diagnostics