#pragma once

// Include after import mean_field. This observer does not change accepted
// coefficients, the prepared operator, or the Newton correction.
#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <limits>
#include <map>
#include <stdexcept>
#include <string>

#include <mfem.hpp>
#include <mpi.h>

namespace experiment::plateau_diagnostics {
// Computes the forward h>=0 boundary along the CURRENT normalizedCorrection
// on exactly the production closure and pressure-force quadrature points.
// n must match the active polytropic EOS. Scalar
// enthalpy coefficients pull back unchanged under the moving-domain map,
// so h_q(alpha)=h_q(0)+alpha*dh_q: no geometry derivative is needed here.
//
// Source convention: prepared_barotropic_closure.cpp get_eos_rule (lines
// 102-157) and TryPrepare's basis contraction/domain check (lines 593-632).
// The production operator rejects h<0 strictly, despite the standalone
// Polytrope EOS extension to h<=0. Pressure-force uses the PressureForce
// query with extraOrder=ceil(n*p_h), as in prepared_pressure_force.cpp
// lines 226-282, and also rejects h<0 (lines 787-803). Its tensor and dense
// contraction paths sample the same polynomial and can differ at roundoff.
// The union bound concerns h>=0 only; geometry and other finite-data tests
// can still impose stricter limits.
//
// alpha_h=+infinity means no sampled enthalpy upper bound in this direction.
// alpha_h=NaN means the accepted base was already invalid or a sampled value
// was nonfinite. alpha_h=0 means a sampled h=0 has dh<0. A step exactly at
// the reported boundary can differ at roundoff due to floating-point state
// addition and contraction. alpha_h remains the closure alias;
// closure_alpha_h, pressure_alpha_h, and alpha_thermodynamic are explicit.
// No candidate is prepared or committed here. physicalBase optionally
// selects a caller-owned frozen trial state instead of the accepted state;
// its coefficients must use this context's physical root layout.
template <class State>
std::map<std::string, double> MeasureEnthalpyStep(State &state, const mean_field::fem::FEM &fem, const double n,
                                                  const mfem::Vector *physicalBase = nullptr) {
    if (!std::isfinite(n) || n < 1.0)
        throw std::invalid_argument("Enthalpy step diagnostics require finite n>=1.");
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    using ClosureDomain = mean_field::field::FieldDomainT<mean_field::field::Density>;
    using PressureDomain = mean_field::field::FieldDomainT<mean_field::field::Enthalpy>;
    using EnthalpyField = mean_field::field::Field<mean_field::field::Enthalpy>;
    int ranks = 0;
    MPI_Comm_size(fem.mesh->GetComm(), &ranks);
    if (ranks != 1)
        throw std::invalid_argument("Enthalpy step diagnostics require exactly one MPI rank.");
    if (!fem.enthalpyFes || !fem.densityFes || !fem.displacementFes || !fem.quadratureFactory) {
        throw std::invalid_argument("Enthalpy step diagnostics received an incomplete FEM.");
    }
    const double maximumExtra = n * mean_field::field::Enthalpy::Scalar::familyOrder;
    if (!std::isfinite(maximumExtra) || maximumExtra > std::numeric_limits<int>::max()) {
        throw std::invalid_argument("Enthalpy step diagnostics received an invalid effective EOS order.");
    }
    int offset = -1, size = -1;
    for (const auto &block : state.Problem().GetManifest().valueBlocks()) {
        if (block.stableId == "specific_enthalpy") {
            offset = block.offset;
            size = block.size;
            break;
        }
    }
    mfem::Vector physicalCorrection;
    state.normalizedOperator->DenormalizeState(state.normalizedCorrection, physicalCorrection);
    const auto &accepted = physicalBase ? *physicalBase : state.AcceptedPhysicalState();
    mfem::ParGridFunction h(fem.enthalpyFes.get()), dh(fem.enthalpyFes.get());
    const auto adapter =
        mean_field::field::make_field_dof_grid_function_adapter<mean_field::field::Enthalpy, DomainSchema>(
            *fem.enthalpyFes);
    if (offset < 0 || size != adapter.dof_map().reduced_size() || offset + size > accepted.Size() ||
        physicalCorrection.Size() != accepted.Size()) {
        throw std::invalid_argument("Enthalpy step diagnostics found an incompatible enthalpy block.");
    }
    mfem::Vector reducedH(size), reducedDh(size);
    for (int i = 0; i < size; ++i) {
        reducedH(i) = accepted(offset + i);
        reducedDh(i) = physicalCorrection(offset + i);
    }
    adapter.scatter(reducedH, h);
    adapter.scatter(reducedDh, dh);

    const double nan = std::numeric_limits<double>::quiet_NaN();
    const double inf = std::numeric_limits<double>::infinity();
    const auto scan = [&](const bool pressure) -> std::map<std::string, double> {
        const int extraOrder =
            static_cast<int>(std::ceil((pressure ? n : n - 1.0) * mean_field::field::Enthalpy::Scalar::familyOrder));
        double alpha = inf, minH = inf, minDh = inf, minFullStepH = inf;
        double limitingH = nan, limitingDh = nan;
        int limitingElement = -1, limitingPoint = -1, limitingAttribute = -1;
        std::array<double, 3> limitingIp{nan, nan, nan};
        int minimumHElement = -1, minimumHPoint = -1;
        int rulePointCountMin = std::numeric_limits<int>::max(), rulePointCountMax = 0;
        std::uint64_t samples = 0, nonfiniteSamples = 0, negativeBaseSamples = 0, decreasingSamples = 0,
                      negativeFullStepSamples = 0;
        mfem::Array<int> dofs;
        mfem::Vector localH, localDh, shape;
        for (int element = 0; element < fem.mesh->GetNE(); ++element) {
            auto *transformation = fem.mesh->GetElementTransformation(element);
            const bool inSupport =
                pressure ? DomainSchema::template attribute_belongs_to<PressureDomain>(transformation->Attribute)
                         : DomainSchema::template attribute_belongs_to<ClosureDomain>(transformation->Attribute);
            if (!inSupport)
                continue;
            const auto &densityElement = *fem.densityFes->GetFE(element);
            const auto &enthalpyElement = *fem.enthalpyFes->GetFE(element);
            if (enthalpyElement.GetOrder() != mean_field::field::Enthalpy::Scalar::familyOrder ||
                (!pressure && densityElement.GetOrder() != mean_field::field::Density::Scalar::familyOrder) ||
                (pressure && fem.displacementFes->GetFE(element)->GetOrder() !=
                                 mean_field::field::Displacement::Vector::familyOrder)) {
                throw std::invalid_argument(
                    "Enthalpy step diagnostics require the registered production field orders.");
            }
            const auto query =
                pressure ? EnthalpyField::template make_query<mean_field::field::Enthalpy::Form::PressureForce>(
                               mean_field::quadrature::QuadratureRole::discretization, transformation->OrderW(),
                               std::array<int, 1>{extraOrder}, mean_field::utils::DOMAINS::STELLAR,
                               mean_field::quadrature::MappingKind::general)
                         : EnthalpyField::template make_query<mean_field::field::Enthalpy::Form::EosClosureSource>(
                               mean_field::quadrature::QuadratureRole::discretization, transformation->OrderW(),
                               std::array<int, 1>{extraOrder}, mean_field::utils::DOMAINS::STELLAR,
                               mean_field::quadrature::MappingKind::general);
            const auto resolution = fem.quadratureFactory->get(query, transformation->GetGeometryType());
            if (!resolution.integration_rule)
                throw std::runtime_error("Missing production EOS quadrature rule.");
            const auto &rule = *resolution.integration_rule;
            rulePointCountMin = std::min(rulePointCountMin, rule.GetNPoints());
            rulePointCountMax = std::max(rulePointCountMax, rule.GetNPoints());
            const auto basis = fem.GetReferenceTables().GetScalarTable(enthalpyElement, rule);
            auto *dofTransformation = fem.enthalpyFes->GetElementDofs(element, dofs);
            h.GetSubVector(dofs, localH);
            dh.GetSubVector(dofs, localDh);
            if (dofTransformation) {
                dofTransformation->InvTransformPrimal(localH);
                dofTransformation->InvTransformPrimal(localDh);
            }
            shape.SetSize(enthalpyElement.GetDof());
            for (int q = 0; q < rule.GetNPoints(); ++q) {
                ++samples;
                basis->GetValues().GetRow(q, shape);
                const double hq = localH * shape;
                const double dhq = localDh * shape;
                if (!std::isfinite(hq) || !std::isfinite(dhq)) {
                    ++nonfiniteSamples;
                    continue;
                }
                if (hq < minH) {
                    minH = hq;
                    minimumHElement = element;
                    minimumHPoint = q;
                }
                minDh = std::min(minDh, dhq);
                minFullStepH = std::min(minFullStepH, hq + dhq);
                negativeBaseSamples += hq < 0.0;
                negativeFullStepSamples += hq + dhq < 0.0;
                if (dhq < 0.0) {
                    ++decreasingSamples;
                    const double candidate = -hq / dhq;
                    if (candidate < alpha) {
                        alpha = candidate;
                        limitingH = hq;
                        limitingDh = dhq;
                        limitingElement = element;
                        limitingPoint = q;
                        limitingAttribute = transformation->Attribute;
                        const auto &ip = rule.IntPoint(q);
                        limitingIp = {ip.x, ip.y, ip.z};
                    }
                }
            }
        }
        const bool valid = samples > 0 && nonfiniteSamples == 0 && negativeBaseSamples == 0;
        return {{"alpha_h", valid ? alpha : nan},
                {"alpha_h_raw_crossing", alpha},
                {"alpha_h_clipped_to_one", valid ? std::min(1.0, alpha) : nan},
                {"enthalpy_base_admissible", valid ? 1.0 : 0.0},
                {"eos_extra_order", static_cast<double>(extraOrder)},
                {"samples", static_cast<double>(samples)},
                {"nonfinite_samples", static_cast<double>(nonfiniteSamples)},
                {"negative_base_samples", static_cast<double>(negativeBaseSamples)},
                {"decreasing_samples", static_cast<double>(decreasingSamples)},
                {"predicted_negative_full_step_samples", static_cast<double>(negativeFullStepSamples)},
                {"minimum_h", samples > nonfiniteSamples ? minH : nan},
                {"minimum_dh", samples > nonfiniteSamples ? minDh : nan},
                {"predicted_minimum_h_full_step", samples > nonfiniteSamples ? minFullStepH : nan},
                {"minimum_h_element", static_cast<double>(minimumHElement)},
                {"minimum_h_point", static_cast<double>(minimumHPoint)},
                {"limiting_element", static_cast<double>(limitingElement)},
                {"limiting_point", static_cast<double>(limitingPoint)},
                {"limiting_attribute", static_cast<double>(limitingAttribute)},
                {"limiting_h", limitingH},
                {"limiting_dh", limitingDh},
                {"limiting_ip_x", limitingIp[0]},
                {"limiting_ip_y", limitingIp[1]},
                {"limiting_ip_z", limitingIp[2]},
                {"quadrature_points_per_element_minimum", samples ? static_cast<double>(rulePointCountMin) : nan},
                {"quadrature_points_per_element_maximum", static_cast<double>(rulePointCountMax)}};
    };
    auto result = scan(false);
    const auto pressure = scan(true);
    result["closure_alpha_h"] = result.at("alpha_h");
    result["closure_samples"] = result.at("samples");
    for (const auto &[name, value] : pressure)
        result["pressure_" + name] = value;
    const bool unionValid =
        result.at("enthalpy_base_admissible") != 0.0 && pressure.at("enthalpy_base_admissible") != 0.0;
    const double unionAlpha = unionValid ? std::min(result.at("alpha_h"), pressure.at("alpha_h")) : nan;
    result["alpha_thermodynamic"] = unionAlpha;
    result["alpha_thermodynamic_clipped_to_one"] = unionValid ? std::min(1.0, unionAlpha) : nan;
    result["thermodynamic_base_admissible"] = unionValid ? 1.0 : 0.0;
    // 0=closure, 1=pressure-force, -1=invalid or no finite upper bound.
    result["thermodynamic_limiting_family"] =
        unionValid && std::isfinite(unionAlpha) ? (pressure.at("alpha_h") < result.at("alpha_h") ? 1.0 : 0.0) : -1.0;
    result["thermodynamic_minimum_h"] = std::isfinite(result.at("minimum_h")) && std::isfinite(pressure.at("minimum_h"))
                                            ? std::min(result.at("minimum_h"), pressure.at("minimum_h"))
                                            : nan;
    return result;
}
} // namespace experiment::plateau_diagnostics