#pragma once

// Diagnostic-only. Include after `import mean_field;`. No candidate is committed.
#include "plateau_enthalpy_step_diagnostics.hpp"
#include <array>
#include <exception>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <string_view>
#include <utility>

namespace experiment::plateau_diagnostics {
namespace shape_curvature_detail {
inline double Norm(const mfem::Vector &v, int offset, int size) {
    long double sum = 0.;
    for (int i = offset; i < offset + size; ++i)
        sum += static_cast<long double>(v(i)) * v(i);
    return std::sqrt(sum);
}
inline long double Dot(const mfem::Vector &a, const mfem::Vector &b, int offset, int size) {
    long double sum = 0.;
    for (int i = offset; i < offset + size; ++i)
        sum += static_cast<long double>(a(i)) * b(i);
    return sum;
}
inline double SafeRatio(double numerator, double denominator) {
    if (!std::isfinite(numerator) || !std::isfinite(denominator) || denominator <= 0.)
        return std::numeric_limits<double>::quiet_NaN();
    const double ratio = numerator / denominator;
    return std::isfinite(ratio) ? ratio : std::numeric_limits<double>::quiet_NaN();
}
inline double Cosine(const mfem::Vector &a, const mfem::Vector &b, int offset, int size) {
    const double na = Norm(a, offset, size), nb = Norm(b, offset, size);
    if (!(na > 0.) || !(nb > 0.))
        return std::numeric_limits<double>::quiet_NaN();
    long double dot = 0.;
    for (int i = offset; i < offset + size; ++i)
        dot += static_cast<long double>(a(i)) * b(i);
    return std::clamp(static_cast<double>(dot / na / nb), -1., 1.);
}
inline bool Finite(const mfem::Vector &v) {
    for (int i = 0; i < v.Size(); ++i)
        if (!std::isfinite(v(i)))
            return false;
    return true;
}
inline std::ofstream Open(const std::filesystem::path &path) {
    std::ofstream file(path);
    file.exceptions(std::ios::failbit | std::ios::badbit);
    file << std::setprecision(17);
    return file;
}
inline const char *ReasonName(mean_field::operators::StellarEquilibriumPreparationRejectionReason reason) {
    using R = mean_field::operators::StellarEquilibriumPreparationRejectionReason;
    switch (reason) {
    case R::inverted_geometry: return "inverted_geometry";
    case R::non_finite_geometry: return "non_finite_geometry";
    case R::thermodynamic_domain: return "thermodynamic_domain";
    case R::non_finite_thermodynamics: return "non_finite_thermodynamics";
    case R::inadmissible_physics: return "inadmissible_physics";
    case R::non_finite_physics: return "non_finite_physics";
    }
    return "unknown";
}
inline const char *StageName(mean_field::operators::StellarEquilibriumPreparationStage stage) {
    using S = mean_field::operators::StellarEquilibriumPreparationStage;
    switch (stage) {
    case S::unspecified: return "unspecified";
    case S::generated_geometry: return "generated_geometry";
    case S::gravity: return "gravity";
    case S::barotropic_closure: return "barotropic_closure";
    case S::hydrostatic_equilibrium: return "hydrostatic_equilibrium";
    case S::displacement_residual: return "displacement_residual";
    case S::pressure_force: return "pressure_force";
    case S::gravity_displacement_force: return "gravity_displacement_force";
    case S::rotational_displacement_force: return "rotational_displacement_force";
    case S::displacement_composition: return "displacement_composition";
    case S::mass_normalization: return "mass_normalization";
    case S::model_specification: return "model_specification";
    }
    return "unknown";
}
} // namespace shape_curvature_detail

// The caller supplies the actual frozen base and selected correction in the
// CURRENT normalization. Fields means every value coordinate except the one
// manifest surface block, including scalar multipliers and angular velocity.
// The usual path costs 22 preparations, three Jv, and no inverse applications.
// On return the frozen base is prepared, caller correction workspaces are
// restored, and the physical inverse may be stale until the caller refreshes it.
template <class State>
void ProbeShapeCurvature(State &s, const mean_field::fem::FEM &f, double n,
                         const mfem::Vector &normalizedBase, const mfem::Vector &normalizedDirection,
                         const std::filesystem::path &output) {
    namespace d = shape_curvature_detail;
    using RejectionReason = mean_field::operators::StellarEquilibriumPreparationRejectionReason;
    const double nan = std::numeric_limits<double>::quiet_NaN();
    const double eps = std::numeric_limits<mfem::real_t>::epsilon();
    int ranks = 0;
    MPI_Comm_size(s.Problem().GetCommunicator(), &ranks);
    if (ranks != 1)
        throw std::invalid_argument("Shape-curvature diagnostics require one MPI rank.");
    const int size = s.Problem().StateSize();
    if (size != s.Problem().EquationSize() || normalizedBase.Size() != size || normalizedDirection.Size() != size ||
        !d::Finite(normalizedBase) || !d::Finite(normalizedDirection))
        throw std::invalid_argument("Shape-curvature diagnostics require finite compatible root vectors.");
    // Copy inputs before touching state: a caller may pass aliases of its workspaces.
    const mfem::Vector base(normalizedBase);
    std::array<mfem::Vector, 3> directions{mfem::Vector(normalizedDirection), mfem::Vector(size),
                                            mfem::Vector(normalizedDirection)};
    constexpr std::array<const char *, 3> names{"full", "shape", "fields"};
    directions[1] = 0.;
    const auto values = s.Problem().GetManifest().valueBlocks();
    const auto rows = s.Problem().GetManifest().residualBlocks();
    int shapeBlocks = 0, shapeOffset = -1, shapeSize = -1;
    for (const auto &block : values) {
        if (!std::string_view(block.stableId).starts_with("surface"))
            continue;
        ++shapeBlocks;
        shapeOffset = block.offset;
        shapeSize = block.size;
        for (int i = block.offset; i < block.offset + block.size; ++i) {
            directions[1](i) = directions[0](i);
            directions[2](i) = 0.;
        }
    }
    if (shapeBlocks != 1 || shapeOffset < 0 || shapeSize <= 0)
        throw std::invalid_argument("Shape-curvature diagnostics require exactly one nonempty surface value block.");
    std::filesystem::create_directories(output);
    if (std::filesystem::exists(output / "shape_curvature_summary.csv"))
        throw std::invalid_argument("Shape-curvature output already exists; use a distinct directory per state.");
    auto summary = d::Open(output / "shape_curvature_summary.csv");
    auto events = d::Open(output / "shape_curvature_preparations.csv");
    auto restores = d::Open(output / "shape_curvature_restorations.csv");
    auto directionFile = d::Open(output / "shape_curvature_directions.csv");
    auto actionsFile = d::Open(output / "shape_curvature_actions.csv");
    auto boundsFile = d::Open(output / "shape_curvature_bounds.csv");
    auto remainderFile = d::Open(output / "shape_curvature_remainders.csv");
    auto mixedFile = d::Open(output / "shape_curvature_mixed.csv");
    summary << "metric,value\n";
    events << "preparation,kind,alpha,admissible,rejection_reason,rejection_name,stage,stage_name,eos_code_applicable,"
              "eos_code,reported_minimum_jacobian,residual_norm\n";
    restores << "preparation,kind,block,base_residual_norm,restoration_difference_norm\n";
    directionFile << "direction,value_block,offset,size,normalized_norm\n";
    actionsFile << "block,base_residual_norm,full_action_norm,shape_action_norm,fields_action_norm,"
                   "split_error_norm,relative_split_error,shape_fields_cosine\n";
    boundsFile << "direction,metric,value\n";
    remainderFile << "alpha,direction,block,admissible,state_increment_resolved,alpha_squared_valid,base_residual_norm,trial_residual_norm,"
                     "action_norm,remainder_norm,remainder_over_alpha_squared,difference_quotient_error,"
                     "relative_difference_quotient_error,heuristic_roundoff_scale,remainder_over_roundoff_scale,"
                     "above_roundoff_scale,expected_step_norm,actual_step_norm,state_addition_error_norm\n";
    mixedFile << "alpha,block,all_trials_admissible,all_state_increments_resolved,alpha_squared_valid,full_remainder_norm,"
                 "shape_remainder_norm,fields_remainder_norm,mixed_remainder_norm,mixed_over_alpha_squared,"
                 "heuristic_roundoff_scale,mixed_over_roundoff_scale,above_roundoff_scale,"
                 "full_shape_cosine,full_fields_cosine,full_mixed_cosine,shape_fields_cosine,"
                 "shape_mixed_cosine,fields_mixed_cosine,component_norm_sum_over_full,linear_prediction_norm,"
                 "full_trial_norm,linear_dot_full,linear_dot_shape,linear_dot_fields,linear_dot_mixed,"
                 "merit_identity_error,projection_decomposition_error\n";
    summary << "machine_epsilon," << eps << "\nshape_offset," << shapeOffset << "\nshape_size," << shapeSize
            << "\nbase_state_norm," << base.Norml2() << "\nfull_direction_norm," << directions[0].Norml2() << '\n';
    mfem::Vector directionSplitError(directions[0]);
    directionSplitError -= directions[1];
    directionSplitError -= directions[2];
    summary << "direction_split_error_norm," << directionSplitError.Norml2() << '\n';
    for (int k = 0; k < 3; ++k)
        for (const auto &block : values)
            directionFile << names[k] << ',' << block.stableId << ',' << block.offset << ',' << block.size << ','
                          << d::Norm(directions[k], block.offset, block.size) << '\n';
    directionFile.flush();

    const mfem::Vector savedCorrection(s.normalizedCorrection), savedPhysicalCorrection(s.physicalCorrection),
        savedVolumeDirection(s.volumeDisplacementDirection), savedAcceptedState(s.acceptedNormalizedState),
        savedAcceptedResidual(s.acceptedNormalizedResidual), savedAcceptedPhysical(s.AcceptedPhysicalState());
    auto restoreWorkspaces = [&]() {
        s.normalizedCorrection = savedCorrection;
        s.physicalCorrection = savedPhysicalCorrection;
        s.volumeDisplacementDirection = savedVolumeDirection;
    };
    int preparationCount = 0;
    mfem::Vector baseResidual(size), restorationEnvelope(size);
    restorationEnvelope = 0.;
    bool haveBaseResidual = false;
    auto prepare = [&](const mfem::Vector &state, const char *kind, double alpha) {
        s.trialNormalizedState = state;
        ++preparationCount;
        auto result = s.PrepareTrial();
        events << preparationCount << ',' << kind << ',' << alpha << ',' << bool(result) << ',';
        if (result) {
            events << "-1,none,-1,none,0,-1," << result->minimumJacobianDeterminant << ','
                   << s.trialNormalizedResidual.Norml2() << '\n';
        } else {
            const auto &error = result.error();
            const bool eos = error.reason == RejectionReason::thermodynamic_domain ||
                             error.reason == RejectionReason::non_finite_thermodynamics;
            events << int(error.reason) << ',' << d::ReasonName(error.reason) << ',' << int(error.stage) << ','
                   << d::StageName(error.stage) << ',' << eos << ',' << (eos ? int(error.thermodynamicErrorCode) : -1)
                   << ',' << error.minimumJacobianDeterminant << ',' << nan << '\n';
        }
        events.flush();
        return result;
    };
    auto restoreBase = [&](const char *kind) {
        const auto result = prepare(base, kind, 0.);
        if (!result || !d::Finite(s.trialNormalizedResidual))
            throw std::runtime_error("Shape-curvature diagnostic could not restore its finite frozen base.");
        if (!haveBaseResidual) {
            baseResidual = s.trialNormalizedResidual;
            haveBaseResidual = true;
        }
        mfem::Vector difference(s.trialNormalizedResidual);
        difference -= baseResidual;
        for (int i = 0; i < size; ++i)
            restorationEnvelope(i) = std::max(restorationEnvelope(i), std::abs(difference(i)));
        for (const auto &block : rows)
            restores << preparationCount << ',' << kind << ',' << block.stableId << ','
                     << d::Norm(baseResidual, block.offset, block.size) << ','
                     << d::Norm(difference, block.offset, block.size) << '\n';
        restores.flush();
    };

    try {
        std::array<mfem::Vector, 3> actions{mfem::Vector(size), mfem::Vector(size), mfem::Vector(size)};
        for (int k = 0; k < 3; ++k) {
            restoreBase(names[k]); // Every J action is preceded by an exact base preparation.
            s.normalizedOperator->Mult(directions[k], actions[k]);
            if (!d::Finite(actions[k]))
                throw std::runtime_error("Nonfinite shape-curvature Jacobian action.");
        }
        mfem::Vector splitError(actions[0]);
        splitError -= actions[1];
        splitError -= actions[2];
        for (const auto &block : rows) {
            const int o = block.offset, count = block.size;
            const double full = d::Norm(actions[0], o, count), shape = d::Norm(actions[1], o, count),
                         fields = d::Norm(actions[2], o, count), error = d::Norm(splitError, o, count);
            const double scale = full + shape + fields;
            actionsFile << block.stableId << ',' << d::Norm(baseResidual, o, count) << ',' << full << ',' << shape
                        << ',' << fields << ',' << error << ',' << (scale > 0. ? error / scale : nan) << ','
                        << d::Cosine(actions[1], actions[2], o, count) << '\n';
        }
        actionsFile.flush();
        const double actionLinearityScale =
            128. * eps * (1. + actions[0].Norml2() + actions[1].Norml2() + actions[2].Norml2());
        const bool linearityResolved = splitError.Norml2() <= actionLinearityScale;
        summary << "base_residual_norm," << baseResidual.Norml2() << "\naction_split_error_norm,"
                << splitError.Norml2() << "\naction_linearity_heuristic_scale," << actionLinearityScale
                << "\naction_linearity_within_heuristic," << linearityResolved << '\n';

        mfem::Vector physicalBase;
        s.normalizedOperator->DenormalizeState(base, physicalBase);
        double commonCap = 1.;
        for (int k = 0; k < 3; ++k) {
            s.normalizedCorrection = directions[k];
            const auto geometry = s.EstimateLargestSafeStepSize(1., .9);
            boundsFile << names[k] << ",geometry_safe_step," << geometry.stepSize << '\n'
                       << names[k] << ",geometry_boundary_step," << geometry.boundaryStepSize << '\n'
                       << names[k] << ",geometry_limited," << geometry.limitedByGeometry << '\n'
                       << names[k] << ",geometry_limiting_element," << geometry.limitingElement << '\n'
                       << names[k] << ",geometry_limiting_rule," << geometry.limitingRule << '\n'
                       << names[k] << ",geometry_limiting_point," << geometry.limitingQuadraturePoint << '\n';
            if (!std::isfinite(geometry.stepSize) || geometry.stepSize <= 0.)
                commonCap = nan;
            else if (std::isfinite(commonCap))
                commonCap = std::min(commonCap, geometry.stepSize);
            // Shape-only has exactly zero enthalpy direction. Full and fields
            // are scanned separately and should have identical sampled bounds.
            if (k != 1) {
                const auto enthalpy = MeasureEnthalpyStep(s, f, n, &physicalBase);
                for (const auto &[key, value] : enthalpy)
                    boundsFile << names[k] << ',' << key << ',' << value << '\n';
                const double bound = enthalpy.at("alpha_thermodynamic");
                if (std::isnan(bound) || bound <= 0.)
                    commonCap = nan;
                else if (std::isfinite(bound) && std::isfinite(commonCap))
                    commonCap = std::min(commonCap, .8 * bound);
            }
        }
        boundsFile.flush();
        summary << "common_alpha_cap," << commonCap << '\n';
        // This is an explicitly heuristic representation threshold, not a
        // certificate of residual assembly accuracy or a physical tolerance.
        const double updateFloor = 64. * eps * (1. + base.Norml2());
        summary << "heuristic_state_update_floor," << updateFloor << '\n';
        const bool usableCap = linearityResolved && std::isfinite(commonCap) && commonCap > 0. &&
                               commonCap * directions[0].Norml2() > updateFloor;
        summary << "common_step_regime_usable," << usableCap << '\n';
        summary.flush();
        if (usableCap) {
            constexpr std::array<double, 6> fractions{1., .1, .01, .001, .0001, .00001};
            for (double fraction : fractions) {
                const double alpha = commonCap * fraction;
                if (!std::isfinite(alpha) || alpha <= 0.) {
                    summary << "skipped_nonpositive_or_nonfinite_alpha," << alpha << '\n';
                    continue;
                }
                const double alphaSquared = alpha * alpha;
                const bool alphaSquaredValid = std::isfinite(alphaSquared) && alphaSquared > 0.;
                mfem::Vector linearPrediction(baseResidual), fullTrialResidual(size);
                linearPrediction.Add(alpha, actions[0]);
                fullTrialResidual = nan;
                std::array<mfem::Vector, 3> remainders{mfem::Vector(size), mfem::Vector(size), mfem::Vector(size)};
                std::array<bool, 3> admissible{}, resolved{};
                std::array<mfem::Vector, 3> floors{mfem::Vector(static_cast<int>(rows.size())),
                    mfem::Vector(static_cast<int>(rows.size())), mfem::Vector(static_cast<int>(rows.size()))};
                for (int k = 0; k < 3; ++k) {
                    mfem::Vector trial(base);
                    trial.Add(alpha, directions[k]);
                    mfem::Vector step(trial);
                    step -= base;
                    mfem::Vector additionError(step);
                    additionError.Add(-alpha, directions[k]);
                    const double expectedStep = alpha * directions[k].Norml2();
                    const double actualStep = step.Norml2(), additionErrorNorm = additionError.Norml2();
                    resolved[k] = directions[k].Norml2() == 0. ||
                                  (expectedStep > updateFloor && actualStep > 0. && additionErrorNorm < .01 * expectedStep);
                    const auto result = prepare(trial, names[k], alpha);
                    admissible[k] = bool(result) && d::Finite(s.trialNormalizedResidual);
                    if (admissible[k]) {
                        if (k == 0)
                            fullTrialResidual = s.trialNormalizedResidual;
                        remainders[k] = s.trialNormalizedResidual;
                        remainders[k] -= baseResidual;
                        remainders[k].Add(-alpha, actions[k]);
                    } else {
                        remainders[k] = nan;
                    }
                    int rowIndex = 0;
                    for (const auto &block : rows) {
                        const int o = block.offset, count = block.size;
                        const double fn = d::Norm(baseResidual, o, count), jn = d::Norm(actions[k], o, count);
                        const double floor = std::max(8. * d::Norm(restorationEnvelope, o, count),
                                                      128. * eps * (1. + fn + alpha * jn));
                        floors[k](rowIndex++) = floor;
                        const double rn = admissible[k] ? d::Norm(remainders[k], o, count) : nan;
                        const double trialNorm = admissible[k] ? d::Norm(s.trialNormalizedResidual, o, count) : nan;
                        remainderFile << alpha << ',' << names[k] << ',' << block.stableId << ',' << admissible[k]
                                      << ',' << resolved[k] << ',' << alphaSquaredValid << ',' << fn << ',' << trialNorm << ',' << jn << ',' << rn
                                      << ',' << d::SafeRatio(rn, alphaSquared) << ',' << d::SafeRatio(rn, alpha) << ','
                                      << d::SafeRatio(d::SafeRatio(rn, alpha), jn) << ',' << floor << ',' << d::SafeRatio(rn, floor)
                                      << ',' << (admissible[k] && resolved[k] && alphaSquaredValid && rn > floor) << ',' << expectedStep
                                      << ',' << actualStep << ',' << additionErrorNorm << '\n';
                    }
                    remainderFile.flush();
                }
                const bool allAdmissible = admissible[0] && admissible[1] && admissible[2];
                const bool allResolved = resolved[0] && resolved[1] && resolved[2];
                mfem::Vector mixed(size);
                mixed = nan;
                if (allAdmissible) {
                    mixed = remainders[0];
                    mixed -= remainders[1];
                    mixed -= remainders[2]; // Subtract vectors before norms, preserving cancellation.
                }
                int rowIndex = 0;
                for (const auto &block : rows) {
                    const int o = block.offset, count = block.size;
                    const double rf = allAdmissible ? d::Norm(remainders[0], o, count) : nan,
                                 rs = allAdmissible ? d::Norm(remainders[1], o, count) : nan,
                                 rp = allAdmissible ? d::Norm(remainders[2], o, count) : nan,
                                 rm = allAdmissible ? d::Norm(mixed, o, count) : nan;
                    const double floor = floors[0](rowIndex) + floors[1](rowIndex) + floors[2](rowIndex);
                    ++rowIndex;
                    mixedFile << alpha << ',' << block.stableId << ',' << allAdmissible << ',' << allResolved << ','
                              << alphaSquaredValid << ',' << rf << ',' << rs << ',' << rp << ',' << rm << ','
                              << d::SafeRatio(rm, alphaSquared) << ',' << floor << ',' << d::SafeRatio(rm, floor) << ','
                              << (allAdmissible && allResolved && alphaSquaredValid && rm > floor);
                    mixedFile << ',' << (allAdmissible ? d::Cosine(remainders[0], remainders[1], o, count) : nan)
                              << ',' << (allAdmissible ? d::Cosine(remainders[0], remainders[2], o, count) : nan)
                              << ',' << (allAdmissible ? d::Cosine(remainders[0], mixed, o, count) : nan)
                              << ',' << (allAdmissible ? d::Cosine(remainders[1], remainders[2], o, count) : nan)
                              << ',' << (allAdmissible ? d::Cosine(remainders[1], mixed, o, count) : nan)
                              << ',' << (allAdmissible ? d::Cosine(remainders[2], mixed, o, count) : nan)
                              << ',' << d::SafeRatio(rs + rp + rm, rf);
                    const long double lf = admissible[0] ? d::Dot(linearPrediction, remainders[0], o, count) : nan,
                                      ls = admissible[1] ? d::Dot(linearPrediction, remainders[1], o, count) : nan,
                                      lp = admissible[2] ? d::Dot(linearPrediction, remainders[2], o, count) : nan,
                                      lm = allAdmissible ? d::Dot(linearPrediction, mixed, o, count) : nan;
                    const long double identityError = admissible[0] ?
                        d::Dot(fullTrialResidual, fullTrialResidual, o, count) -
                        d::Dot(linearPrediction, linearPrediction, o, count) - 2. * lf -
                        d::Dot(remainders[0], remainders[0], o, count) : nan;
                    mixedFile << ',' << d::Norm(linearPrediction, o, count) << ','
                              << (admissible[0] ? d::Norm(fullTrialResidual, o, count) : nan) << ','
                              << lf << ',' << ls << ',' << lp << ',' << lm << ',' << identityError << ','
                              << (allAdmissible ? lf - ls - lp - lm : nan) << '\n';
                }
                mixedFile.flush();
            }
        }
        restoreBase("final_restore");
        summary << "preparations," << preparationCount << "\nrestoration_envelope_norm,"
                << restorationEnvelope.Norml2() << '\n';
        for (const auto &entry : std::array<std::pair<const char *, const mfem::Vector *>, 3>{{
                 {"accepted_normalized_state_change", &s.acceptedNormalizedState},
                 {"accepted_residual_change", &s.acceptedNormalizedResidual},
                 {"accepted_physical_state_change", &s.AcceptedPhysicalState()}}}) {
            const mfem::Vector *reference = entry.second == &s.acceptedNormalizedState ? &savedAcceptedState :
                (entry.second == &s.acceptedNormalizedResidual ? &savedAcceptedResidual : &savedAcceptedPhysical);
            mfem::Vector difference(*entry.second);
            difference -= *reference;
            summary << entry.first << ',' << difference.Norml2() << '\n';
            if (difference.Norml2() != 0.)
                throw std::runtime_error("Shape-curvature diagnostics detected a changed accepted vector.");
        }
        summary.flush();
        restoreWorkspaces();
    } catch (...) {
        const auto original = std::current_exception();
        try {
            restoreBase("exception_restore");
        } catch (...) {
            restoreWorkspaces();
            throw;
        }
        restoreWorkspaces();
        std::rethrow_exception(original);
    }
}
} // namespace experiment::plateau_diagnostics