#pragma once

// Experiment-only observer. Include after `import mean_field;`, with the
// repository's experiments directory on the include path. The live-state
// PhysicalState constructor is generic in n and rotation; its saved-file replay
// constructor is intentionally not used here.
#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <limits>
#include <map>
#include <numbers>
#include <stdexcept>
#include <string>
#include <utility>

#include "polytrope_physical_state.hpp"

namespace experiment::plateau_diagnostics {
using Measurements = std::map<std::string, double>;

// Assumes the production zero-surface-pressure problem, hence no external
// pressure surface term in 2T + W + 3 integral(P). Both the gravitational
// energy W=.5 integral(rho phi) and force W=-integral(rho x.grad(phi)) are
// reported. The latter uses the independently represented physical RT flux,
// not the broken gradient of the L2 potential. PhysicalState::Evaluate
// already applies the deformation Piola transform after MFEM's reference
// mesh Piola transform: no additional vector transformation belongs here.
//
// Call independently with orders 14 and 18 to assess quadrature sensitivity.
// These are quadrature-sampled extrema, not global pointwise certificates.
// Negative rho is not clipped: signed mass, W, T, J, and rho(h) mismatch are
// preserved. P(rho) is undefined for negative rho, so its complete integral
// and associated virials are NaN whenever any such samples occur. P(h) uses
// the production EOS extension rho(h)=0, P(h)=0 for h<=0. Negative samples
// and their volume/mass are explicitly reported.
//
// Invalid mapping/field samples invalidate all complete volume integrals;
// they are counted instead of aborting measurement or silently dropping
// their contribution. Parameters that cannot describe a polytrope throw.
inline Measurements MeasurePhysical(experiment::polytrope_validation::PhysicalState &state, const double n,
                                    const double K, const int quadratureOrder = 14) {
    if (!std::isfinite(n) || n < 1.0 || !std::isfinite(K) || K <= 0.0 || quadratureOrder < 0) {
        throw std::invalid_argument(
            "Physical diagnostics require finite n>=1, finite K>0, and nonnegative quadrature order.");
    }
    int ranks = 0;
    MPI_Comm_size(state.finiteElements.mesh->GetComm(), &ranks);
    if (ranks != 1)
        throw std::invalid_argument("Physical diagnostics require exactly one MPI rank.");

    const long double index = n;
    const long double constant = K;
    const long double enthalpyFactor = (index + 1.0L) * constant;
    const auto &omega = state.rotation.angular_velocity();
    const auto &center = state.rotation.center();
    const double nan = std::numeric_limits<double>::quiet_NaN();
    const long double inf = std::numeric_limits<long double>::infinity();
    std::uint64_t samples = 0, validSamples = 0, invalidSamples = 0;
    std::uint64_t negativeRhoSamples = 0, negativeHSamples = 0;
    std::uint64_t invalidEnthalpyEosSamples = 0, invalidDensityEosSamples = 0;
    long double volume = 0.0L, mass = 0.0L, binding = 0.0L, forceBinding = 0.0L;
    long double pressureH = 0.0L, pressureRhoNonnegative = 0.0L, kinetic = 0.0L;
    long double negativeRhoVolume = 0.0L, negativeHVolume = 0.0L, negativeRhoMassMagnitude = 0.0L;
    long double closureSquared = 0.0L, closureIntegral = 0.0L, closureMax = 0.0L;
    long double pressureMismatchSquared = 0.0L, pressureMismatchVolume = 0.0L;
    long double bernoulliMean = 0.0L, bernoulliCenteredSquared = 0.0L, bernoulliRawSquared = 0.0L;
    long double bernoulliMin = inf, bernoulliMax = -inf;
    long double rhoMin = inf, rhoMax = -inf, hMin = inf, hMax = -inf, rhoFromHMax = 0.0L;
    long double minDeterminant = inf, inertiaZ = 0.0L;
    long double gradientMismatchSquared = 0.0L, hydrostaticGradientSquared = 0.0L;
    long double drhodhPositiveMin = inf, drhodhPositiveMax = 0.0L;
    std::array<long double, 3> firstMoment{}, angularMomentum{}, gravitationalForce{};
    experiment::polytrope_validation::PhysicalPoint point;
    mfem::Vector centrifugalGradient(3);

    for (int element = 0; element < state.finiteElements.mesh->GetNE(); ++element) {
        if (!state.isStellar(element))
            continue;
        auto *transformation = state.finiteElements.mesh->GetElementTransformation(element);
        const auto &rule = mfem::IntRules.Get(transformation->GetGeometryType(), quadratureOrder);
        for (int q = 0; q < rule.GetNPoints(); ++q) {
            ++samples;
            const auto &ip = rule.IntPoint(q);
            if (state.Evaluate(element, ip, point) != mean_field::mapping::MappingStatus::valid) {
                ++invalidSamples;
                continue;
            }
            // Evaluate may reuse MFEM's cached transformation; restore ip.
            transformation->SetIntPoint(&ip);
            const double referenceDeterminant = transformation->Jacobian().Det();
            const long double weight =
                static_cast<long double>(ip.weight) * referenceDeterminant * point.mapping.mapping_determinant;
            if (!(referenceDeterminant > 0.0) || !std::isfinite(referenceDeterminant) ||
                !(point.mapping.mapping_determinant > 0.0) || !std::isfinite(point.mapping.mapping_determinant) ||
                !(weight > 0.0L) || !std::isfinite(weight)) {
                ++invalidSamples;
                continue;
            }
            const auto &position = point.mapping.physical_position;
            state.rotation.potential_gradient(position, centrifugalGradient);
            bool finite = std::isfinite(point.rho) && std::isfinite(point.h) && std::isfinite(point.phi) &&
                          std::isfinite(point.rotationPotential);
            for (int d = 0; d < 3; ++d) {
                finite = finite && std::isfinite(position(d)) && std::isfinite(point.gravityGradientPhysical(d)) &&
                         std::isfinite(point.potentialGradientPhysical(d)) &&
                         std::isfinite(point.enthalpyGradientPhysical(d)) && std::isfinite(centrifugalGradient(d));
            }
            if (!finite) {
                ++invalidSamples;
                continue;
            }
            ++validSamples;
            const long double rho = point.rho;
            const long double h = point.h;
            const long double phi = point.phi;
            const long double psi = point.rotationPotential;
            const long double rhoH = h > 0.0L ? std::pow(h / enthalpyFactor, index) : 0.0L;
            const long double pH = h > 0.0L ? rhoH * h / (index + 1.0L) : 0.0L;
            const long double pRho = rho >= 0.0L ? constant * std::pow(rho, 1.0L + 1.0L / index) : 0.0L;
            const bool finiteHEos = std::isfinite(rhoH) && std::isfinite(pH);
            const bool finiteRhoEos = rho >= 0.0L && std::isfinite(pRho);
            if (!finiteHEos)
                ++invalidEnthalpyEosSamples;
            if (rho >= 0.0L && !finiteRhoEos)
                ++invalidDensityEosSamples;

            volume += weight;
            mass += weight * rho;
            binding += 0.5L * weight * rho * phi;
            kinetic += weight * rho * psi;
            const long double specificBernoulli = h + phi - psi;
            if (validSamples == 1) {
                // Exact initialization avoids a spurious first-sample
                // variance from rounding an O(C) Bernoulli value.
                bernoulliMean = specificBernoulli;
            } else {
                const long double delta = specificBernoulli - bernoulliMean;
                bernoulliMean += weight / volume * delta;
                bernoulliCenteredSquared += weight * delta * (specificBernoulli - bernoulliMean);
            }
            const long double rawBernoulli = specificBernoulli - state.bernoulliConstant;
            bernoulliRawSquared += weight * rawBernoulli * rawBernoulli;
            bernoulliMin = std::min(bernoulliMin, specificBernoulli);
            bernoulliMax = std::max(bernoulliMax, specificBernoulli);
            rhoMin = std::min(rhoMin, rho);
            rhoMax = std::max(rhoMax, rho);
            hMin = std::min(hMin, h);
            hMax = std::max(hMax, h);
            minDeterminant = std::min(minDeterminant, static_cast<long double>(point.mapping.mapping_determinant));
            if (rho < 0.0L) {
                ++negativeRhoSamples;
                negativeRhoVolume += weight;
                negativeRhoMassMagnitude -= weight * rho;
            }
            if (h < 0.0L) {
                ++negativeHSamples;
                negativeHVolume += weight;
            }
            if (finiteHEos) {
                pressureH += weight * pH;
                const long double closure = rho - rhoH;
                closureIntegral += weight * closure;
                closureSquared += weight * closure * closure;
                closureMax = std::max(closureMax, std::abs(closure));
                rhoFromHMax = std::max(rhoFromHMax, rhoH);
                if (h > 0.0L) {
                    const long double derivative = index * rhoH / h;
                    drhodhPositiveMin = std::min(drhodhPositiveMin, derivative);
                    drhodhPositiveMax = std::max(drhodhPositiveMax, derivative);
                }
            }
            if (finiteRhoEos)
                pressureRhoNonnegative += weight * pRho;
            if (finiteHEos && finiteRhoEos) {
                const long double pressureMismatch = pRho - pH;
                pressureMismatchSquared += weight * pressureMismatch * pressureMismatch;
                pressureMismatchVolume += weight;
            }

            std::array<long double, 3> radius{}, velocity{};
            for (int d = 0; d < 3; ++d)
                radius[d] = static_cast<long double>(position(d)) - center(d);
            velocity[0] =
                static_cast<long double>(omega(1)) * radius[2] - static_cast<long double>(omega(2)) * radius[1];
            velocity[1] =
                static_cast<long double>(omega(2)) * radius[0] - static_cast<long double>(omega(0)) * radius[2];
            velocity[2] =
                static_cast<long double>(omega(0)) * radius[1] - static_cast<long double>(omega(1)) * radius[0];
            angularMomentum[0] += weight * rho * (radius[1] * velocity[2] - radius[2] * velocity[1]);
            angularMomentum[1] += weight * rho * (radius[2] * velocity[0] - radius[0] * velocity[2]);
            angularMomentum[2] += weight * rho * (radius[0] * velocity[1] - radius[1] * velocity[0]);
            inertiaZ += weight * rho * (radius[0] * radius[0] + radius[1] * radius[1]);
            for (int d = 0; d < 3; ++d) {
                const long double flux = point.gravityGradientPhysical(d);
                forceBinding -= weight * rho * position(d) * flux;
                firstMoment[d] += weight * rho * position(d);
                gravitationalForce[d] -= weight * rho * flux;
                const long double mismatch = flux - point.potentialGradientPhysical(d);
                const long double hydrostatic = point.enthalpyGradientPhysical(d) + flux - centrifugalGradient(d);
                gradientMismatchSquared += weight * mismatch * mismatch;
                hydrostaticGradientSquared += weight * hydrostatic * hydrostatic;
            }
        }
    }

    const bool complete = invalidSamples == 0 && validSamples > 0 && volume > 0.0L && std::isfinite(volume);
    const bool validHEos = complete && invalidEnthalpyEosSamples == 0;
    const bool validRhoEos = complete && negativeRhoSamples == 0 && invalidDensityEosSamples == 0;
    const bool validBinding = complete && std::isfinite(binding) && binding < 0.0L;
    const auto value = [nan](const long double v, const bool valid = true) -> double {
        const double result = static_cast<double>(v);
        return valid && std::isfinite(result) ? result : nan;
    };
    const auto rms = [volume](const long double sum) {
        if (!(volume > 0.0L) || !std::isfinite(volume) || !std::isfinite(sum)) {
            return std::numeric_limits<long double>::quiet_NaN();
        }
        return std::sqrt(std::max(0.0L, sum / volume));
    };
    const auto relative = [nan, &value](const long double a, const long double denominator, const bool valid) {
        return valid && denominator > 0.0L ? value(a / denominator) : nan;
    };
    const bool validPressureMismatch = validHEos && pressureMismatchVolume > 0.0L && invalidDensityEosSamples == 0;
    const double pressureMismatchRms =
        validPressureMismatch ? value(std::sqrt(pressureMismatchSquared / pressureMismatchVolume)) : nan;
    const long double rhoScale = std::max({std::abs(rhoMin), std::abs(rhoMax), rhoFromHMax});
    const long double hScale = std::max(std::abs(hMin), std::abs(hMax));
    const long double absoluteBinding = std::abs(binding);
    const long double volumeRadius = std::cbrt(3.0L * volume / (4.0L * std::numbers::pi_v<long double>));

    Measurements result{
        {"quadrature_order", static_cast<double>(quadratureOrder)},
        {"polytropic_index", n},
        {"polytropic_constant", K},
        {"stellar_samples", static_cast<double>(samples)},
        {"valid_stellar_samples", static_cast<double>(validSamples)},
        {"invalid_stellar_samples", static_cast<double>(invalidSamples)},
        {"complete_volume_measurement", complete ? 1.0 : 0.0},
        {"pressure_from_enthalpy_defined_everywhere", validHEos ? 1.0 : 0.0},
        {"pressure_from_density_defined_everywhere", validRhoEos ? 1.0 : 0.0},
        {"invalid_enthalpy_eos_samples", static_cast<double>(invalidEnthalpyEosSamples)},
        {"invalid_density_eos_samples", static_cast<double>(invalidDensityEosSamples)},
        {"volume", value(volume, complete)},
        {"mass", value(mass, complete)},
        {"volume_radius", value(volumeRadius, complete)},
        {"binding_energy", value(binding, complete)},
        {"force_binding_energy", value(forceBinding, complete)},
        {"kinetic_energy", value(kinetic, complete)},
        {"pressure_integral", value(pressureRhoNonnegative, validRhoEos)},
        {"enthalpy_pressure_integral", value(pressureH, validHEos)},
        {"nonnegative_density_pressure_integral",
         value(pressureRhoNonnegative, complete && invalidDensityEosSamples == 0)},
        {"kinetic_to_binding_ratio", relative(kinetic, absoluteBinding, validBinding)},
        {"gravity_energy_consistency", relative(std::abs(binding - forceBinding), absoluteBinding, validBinding)},
        {"moment_of_inertia_z", value(inertiaZ, complete)},
        {"angular_velocity_norm", state.rotation.angular_velocity().Norml2()},
        {"angular_momentum_z", value(angularMomentum[2], complete)},
        {"minimum_density", value(rhoMin)},
        {"maximum_density", value(rhoMax)},
        {"minimum_enthalpy", value(hMin)},
        {"maximum_enthalpy", value(hMax)},
        {"minimum_mapping_determinant", value(minDeterminant)},
        {"negative_density_samples", static_cast<double>(negativeRhoSamples)},
        {"negative_enthalpy_samples", static_cast<double>(negativeHSamples)},
        {"negative_density_volume_fraction", relative(negativeRhoVolume, volume, complete)},
        {"negative_enthalpy_volume_fraction", relative(negativeHVolume, volume, complete)},
        {"negative_density_mass_magnitude", value(negativeRhoMassMagnitude, complete)},
        {"negative_density_mass_fraction", relative(negativeRhoMassMagnitude, std::abs(mass), complete)},
        {"bernoulli_constant", state.bernoulliConstant},
        {"bernoulli_mean", value(bernoulliMean, complete)},
        {"bernoulli_mean_minus_constant", value(bernoulliMean - state.bernoulliConstant, complete)},
        {"bernoulli_raw_rms", value(rms(bernoulliRawSquared), complete)},
        {"bernoulli_mean_subtracted_rms", value(rms(bernoulliCenteredSquared), complete)},
        {"bernoulli_raw_rms_scaled", relative(rms(bernoulliRawSquared), hScale, complete)},
        {"bernoulli_mean_subtracted_rms_scaled", relative(rms(bernoulliCenteredSquared), hScale, complete)},
        {"bernoulli_sampled_range", value(bernoulliMax - bernoulliMin, complete)},
        {"eos_density_mismatch_integral", value(closureIntegral, validHEos)},
        {"eos_density_mismatch_rms", value(rms(closureSquared), validHEos)},
        {"eos_density_mismatch_maximum", value(closureMax, validHEos)},
        {"eos_density_mismatch_rms_scaled", relative(rms(closureSquared), rhoScale, validHEos)},
        {"eos_density_mismatch_maximum_scaled", relative(closureMax, rhoScale, validHEos)},
        {"eos_density_scale_sampled_maximum", value(rhoScale, validHEos)},
        {"eos_enthalpy_scale_sampled_maximum", value(hScale, complete)},
        {"eos_drhodh_positive_h_minimum", value(drhodhPositiveMin, validHEos)},
        {"eos_drhodh_positive_h_maximum", value(drhodhPositiveMax, validHEos)},
        {"eos_pressure_mismatch_rms_on_nonnegative_density", pressureMismatchRms},
        {"gravity_gradient_vs_broken_potential_gradient_rms", value(rms(gradientMismatchSquared), complete)},
        {"hydrostatic_gradient_rms", value(rms(hydrostaticGradientSquared), complete)}};
    for (const auto &entry :
         std::array<std::pair<std::string, long double>, 2>{{{"", pressureRhoNonnegative}, {"enthalpy_", pressureH}}}) {
        const bool validPressure = entry.first.empty() ? validRhoEos : validHEos;
        const long double energyNumerator = 2.0L * kinetic + binding + 3.0L * entry.second;
        const long double forceNumerator = 2.0L * kinetic + forceBinding + 3.0L * entry.second;
        result[entry.first + "virial_signed"] =
            relative(energyNumerator, absoluteBinding, validBinding && validPressure);
        result[entry.first + "virial_error"] =
            relative(std::abs(energyNumerator), absoluteBinding, validBinding && validPressure);
        // Use |W_energy| for both definitions to compare their numerators.
        result[entry.first + "force_virial_signed"] =
            relative(forceNumerator, absoluteBinding, validBinding && validPressure);
        result[entry.first + "force_virial_error"] =
            relative(std::abs(forceNumerator), absoluteBinding, validBinding && validPressure);
    }
    long double omegaDotAngularMomentum = 0.0L;
    for (int d = 0; d < 3; ++d) {
        result["center_of_mass_" + std::to_string(d)] = complete && mass != 0.0L ? value(firstMoment[d] / mass) : nan;
        result["angular_momentum_about_rotation_center_" + std::to_string(d)] = value(angularMomentum[d], complete);
        result["net_gravitational_force_" + std::to_string(d)] = value(gravitationalForce[d], complete);
        result["rotation_center_" + std::to_string(d)] = center(d);
        result["angular_velocity_" + std::to_string(d)] = omega(d);
        omegaDotAngularMomentum += static_cast<long double>(omega(d)) * angularMomentum[d];
    }
    result["rigid_rotation_energy_identity_absolute_error"] =
        value(std::abs(omegaDotAngularMomentum - 2.0L * kinetic), complete);
    return result;
}
} // namespace experiment::plateau_diagnostics