#pragma once

#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <limits>
#include <map>
#include <numbers>
#include <stdexcept>
#include <string>

#include "polytrope_analytic_reference.hpp"
#include "polytrope_physical_state.hpp"

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

    // Exact fields evaluated on the actual mapped mesh. This checks the same
    // integration and physical-point locator used for a numerical solution,
    // without constructing a Newton context or using its Lane-Emden seed.
    class AnalyticMeshState final {
    public:
        const mean_field::fem::FEM &finiteElements;
        mean_field::mapping::GridFunctionMappingEvaluator mapping;
        N1Reference reference;

        AnalyticMeshState(const mean_field::fem::FEM &fem, const N1Reference &analytic)
            : finiteElements(fem),
              mapping(*fem.domainMapperStateless, *fem.displacement, *fem.compactificationCoordinate),
              reference(analytic) {
        }

        [[nodiscard]] bool isStellar(const int element) const {
            using Schema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
            return Schema::template attribute_belongs_to<mean_field::utils::domain::Stellar>(
                finiteElements.mesh->GetAttribute(element)
            );
        }

        mean_field::mapping::MappingStatus Evaluate(
            const int element, const mfem::IntegrationPoint &point, PhysicalPoint &result
        ) {
            auto *transformation = finiteElements.mesh->GetElementTransformation(element);
            transformation->SetIntPoint(&point);
            const auto status = mapping.EvaluatePoint(*transformation, point, result.mapping);
            if (status != mean_field::mapping::MappingStatus::valid) return status;
            const double radius = result.mapping.physical_position.Norml2();
            const auto analytic = reference.AtRadius(radius);
            result.element = element;
            result.attribute = transformation->Attribute;
            result.stellarMaterial = isStellar(element);
            result.rho = analytic.density;
            result.h = analytic.enthalpy;
            result.phi = analytic.potential;
            result.rotationPotential = 0.0;
            result.gravityGradientPhysical.SetSize(3);
            result.enthalpyGradientPhysical.SetSize(3);
            result.potentialGradientPhysical.SetSize(3);
            for (int component = 0; component < 3; ++component) {
                const double gradient = radius > 0.0
                    ? analytic.radialPotentialGradient * result.mapping.physical_position(component) / radius : 0.0;
                result.gravityGradientPhysical(component) = gradient;
                result.potentialGradientPhysical(component) = gradient;
                result.enthalpyGradientPhysical(component) = radius <= reference.radius ? -gradient : 0.0;
            }
            return status;
        }
    };

    struct ErrorIntegral final {
        long double errorSquared{0.0L};
        long double referenceSquared{0.0L};
        double maximumScaledError{0.0};

        void Add(const double value, const double exact, const double weight, const double scale) {
            const long double difference = static_cast<long double>(value) - exact;
            errorSquared += weight * difference * difference;
            referenceSquared += static_cast<long double>(weight) * exact * exact;
            maximumScaledError = std::max(maximumScaledError, std::abs(value - exact) / scale);
        }

        [[nodiscard]] double RelativeL2() const {
            return referenceSquared > 0.0L ? static_cast<double>(std::sqrt(errorSquared / referenceSquared))
                                          : std::numeric_limits<double>::quiet_NaN();
        }
    };

    template <typename SampleState>
    Measurements MeasureVolumes(SampleState &state, const N1Reference &reference, const int quadratureOrder) {
        long double volume = 0.0L, mass = 0.0L, binding = 0.0L, forceBinding = 0.0L;
        long double pressureIntegral = 0.0L, enthalpyPressureIntegral = 0.0L, kinetic = 0.0L;
        long double momentOfInertia = 0.0L, bernoulliOffsetMean = 0.0L, bernoulliCenteredSquared = 0.0L;
        long double closureSquared = 0.0L, closureDensityInnerProduct = 0.0L;
        long double gradientMismatch = 0.0L, hydrostaticGradient = 0.0L;
        std::array<long double, 3> firstMoment{};
        ErrorIntegral densityError, enthalpyError, potentialError, pressureError, gravityError;
        double minimumDensity = std::numeric_limits<double>::infinity();
        double minimumEnthalpy = std::numeric_limits<double>::infinity();
        double minimumDeterminant = std::numeric_limits<double>::infinity();
        double bernoulliMinimum = std::numeric_limits<double>::infinity();
        double bernoulliMaximum = -std::numeric_limits<double>::infinity();
        std::uint64_t samples = 0, negativeDensitySamples = 0, negativeEnthalpySamples = 0;
        PhysicalPoint point;
        const double densityScale = reference.CentralDensity();
        const double enthalpyScale = reference.CentralEnthalpy();
        const double pressureScale = reference.PolytropicConstant() * densityScale * densityScale;
        const double gravityScale = reference.gravitationalConstant * reference.mass /
                                    (reference.radius * reference.radius);

        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) {
                const auto &ip = rule.IntPoint(q);
                if (state.Evaluate(element, ip, point) != mean_field::mapping::MappingStatus::valid) {
                    throw std::runtime_error("Invalid physical mapping in polytrope volume verification.");
                }
                transformation->SetIntPoint(&ip);
                // Explicitly check the orientation of both the reference map
                // and the deformation before integrating physical volume.
                const double referenceDeterminant = transformation->Jacobian().Det();
                const double weight = ip.weight * referenceDeterminant * point.mapping.mapping_determinant;
                if (!(referenceDeterminant > 0.0) || !std::isfinite(referenceDeterminant)) {
                    throw std::runtime_error("Non-positive or non-finite reference Jacobian in polytrope volume verification.");
                }
                if (!(weight > 0.0) || !std::isfinite(weight) || !std::isfinite(point.rho) ||
                    !std::isfinite(point.h) || !std::isfinite(point.phi) || !std::isfinite(point.rotationPotential)) {
                    throw std::runtime_error("Non-finite field or non-positive physical integration weight.");
                }
                ++samples;
                const auto &position = point.mapping.physical_position;
                const double radius = position.Norml2();
                const auto exact = reference.AtRadius(radius);
                // Do not clamp numerical density/enthalpy. Negative values are
                // reported, and pressure consistency is checked independently.
                const double pressure = reference.PolytropicConstant() * point.rho * point.rho;
                const double pressureFromEnthalpy = point.h * point.h / (4.0 * reference.PolytropicConstant());
                const long double specificBernoulli = static_cast<long double>(point.h) + point.phi - point.rotationPotential;
                volume += weight;
                mass += weight * point.rho;
                binding += 0.5L * weight * point.rho * point.phi;
                forceBinding -= weight * point.rho * (position * point.gravityGradientPhysical);
                pressureIntegral += weight * pressure;
                enthalpyPressureIntegral += weight * pressureFromEnthalpy;
                // RigidRotation::potential is positive +|Omega x r|^2/2.
                kinetic += weight * point.rho * point.rotationPotential;
                momentOfInertia += weight * point.rho * (position(0) * position(0) + position(1) * position(1));
                // Weighted Welford accumulation about the fixed analytic
                // Bernoulli constant C=-h_c resolves small spatial variations
                // without subtracting two O(h_c^2) second moments. No fitted
                // potential offset is applied to any physical field/error.
                const long double bernoulliOffset = specificBernoulli + enthalpyScale;
                const long double bernoulliDelta = bernoulliOffset - bernoulliOffsetMean;
                bernoulliOffsetMean += (static_cast<long double>(weight) / volume) * bernoulliDelta;
                bernoulliCenteredSquared += weight * bernoulliDelta * (bernoulliOffset - bernoulliOffsetMean);
                const double closure = point.h - 2.0 * reference.PolytropicConstant() * point.rho;
                closureSquared += weight * closure * closure;
                closureDensityInnerProduct += weight * point.rho * closure;
                bernoulliMinimum = std::min(bernoulliMinimum, static_cast<double>(specificBernoulli));
                bernoulliMaximum = std::max(bernoulliMaximum, static_cast<double>(specificBernoulli));
                minimumDensity = std::min(minimumDensity, point.rho);
                minimumEnthalpy = std::min(minimumEnthalpy, point.h);
                minimumDeterminant = std::min(minimumDeterminant, point.mapping.mapping_determinant);
                negativeDensitySamples += point.rho < 0.0;
                negativeEnthalpySamples += point.h < 0.0;
                densityError.Add(point.rho, exact.density, weight, densityScale);
                enthalpyError.Add(point.h, exact.enthalpy, weight, enthalpyScale);
                potentialError.Add(point.phi, exact.potential, weight, enthalpyScale);
                pressureError.Add(pressure, exact.pressure, weight, pressureScale);
                for (int component = 0; component < 3; ++component) {
                    if (!std::isfinite(point.gravityGradientPhysical(component)) ||
                        !std::isfinite(point.potentialGradientPhysical(component)) ||
                        !std::isfinite(point.enthalpyGradientPhysical(component))) {
                        throw std::runtime_error("Non-finite physical field gradient in polytrope volume verification.");
                    }
                    firstMoment[component] += weight * point.rho * position(component);
                    const double exactGradient = radius > 0.0
                        ? exact.radialPotentialGradient * position(component) / radius : 0.0;
                    gravityError.Add(point.gravityGradientPhysical(component), exactGradient, weight, gravityScale);
                    const double mismatch = point.gravityGradientPhysical(component) - point.potentialGradientPhysical(component);
                    gradientMismatch += weight * mismatch * mismatch;
                    const double hydrostatic = point.enthalpyGradientPhysical(component) + point.gravityGradientPhysical(component);
                    hydrostaticGradient += weight * hydrostatic * hydrostatic;
                }
            }
        }
        if (!(volume > 0.0L) || !(mass > 0.0L) || !(binding < 0.0L)) {
            throw std::runtime_error("Polytrope verification requires positive stellar volume/mass and negative binding energy.");
        }
        Measurements result{
            {"quadrature_order", static_cast<double>(quadratureOrder)}, {"stellar_samples", static_cast<double>(samples)},
            {"volume", static_cast<double>(volume)}, {"mass", static_cast<double>(mass)},
            {"mass_relative_error", static_cast<double>(std::abs(mass / reference.mass - 1.0L))},
            {"volume_radius", static_cast<double>(std::cbrt(3.0L * volume / (4.0L * std::numbers::pi_v<long double>)))},
            {"binding_energy", static_cast<double>(binding)}, {"force_binding_energy", static_cast<double>(forceBinding)},
            {"pressure_integral", static_cast<double>(pressureIntegral)},
            {"enthalpy_pressure_integral", static_cast<double>(enthalpyPressureIntegral)},
            {"kinetic_energy", static_cast<double>(kinetic)},
            {"virial_signed", static_cast<double>((2.0L * kinetic + binding + 3.0L * pressureIntegral) / std::abs(binding))},
            {"virial_error", static_cast<double>(std::abs(2.0L * kinetic + binding + 3.0L * pressureIntegral) / std::abs(binding))},
            {"virial_ratio", static_cast<double>((2.0L * kinetic + 3.0L * pressureIntegral) / std::abs(binding))},
            {"force_virial_error", static_cast<double>(std::abs(2.0L * kinetic + forceBinding + 3.0L * pressureIntegral) / std::abs(binding))},
            {"enthalpy_virial_error", static_cast<double>(std::abs(2.0L * kinetic + binding + 3.0L * enthalpyPressureIntegral) / std::abs(binding))},
            {"enthalpy_force_virial_error", static_cast<double>(std::abs(2.0L * kinetic + forceBinding + 3.0L * enthalpyPressureIntegral) / std::abs(binding))},
            {"gravity_energy_consistency", static_cast<double>(std::abs(binding - forceBinding) / std::abs(binding))},
            {"binding_relative_error", static_cast<double>(std::abs(binding / reference.BindingEnergy() - 1.0L))},
            {"pressure_integral_relative_error", static_cast<double>(std::abs(pressureIntegral / reference.PressureIntegral() - 1.0L))},
            {"pressure_integral_eos_disagreement", static_cast<double>(std::abs(pressureIntegral - enthalpyPressureIntegral) / reference.PressureIntegral())},
            {"closure_projection_pressure_gap", static_cast<double>(closureSquared / (4.0L * reference.PolytropicConstant()))},
            {"closure_density_inner_product", static_cast<double>(closureDensityInnerProduct)},
            {"closure_projection_pressure_gap_relative_defect", static_cast<double>((enthalpyPressureIntegral - pressureIntegral -
                closureSquared / (4.0L * reference.PolytropicConstant())) / reference.PressureIntegral())},
            {"moment_of_inertia", static_cast<double>(momentOfInertia)},
            {"moment_of_inertia_relative_error", static_cast<double>(std::abs(momentOfInertia / reference.MomentOfInertia() - 1.0L))},
            {"bernoulli_mean", static_cast<double>(bernoulliOffsetMean - enthalpyScale)},
            {"bernoulli_mean_scaled_error", static_cast<double>(std::abs(bernoulliOffsetMean) / enthalpyScale)},
            {"bernoulli_scaled_range", (bernoulliMaximum - bernoulliMinimum) / enthalpyScale},
            {"bernoulli_scaled_rms_variation", static_cast<double>(std::sqrt(std::max(0.0L, bernoulliCenteredSquared / volume)) / enthalpyScale)},
            {"eos_enthalpy_scaled_rms", static_cast<double>(std::sqrt(closureSquared / volume) / enthalpyScale)},
            {"gravity_gradient_vs_broken_potential_gradient_scaled_rms", static_cast<double>(std::sqrt(gradientMismatch / volume) / gravityScale)},
            {"nonrotating_hydrostatic_gradient_scaled_rms", static_cast<double>(std::sqrt(hydrostaticGradient / volume) / gravityScale)},
            {"minimum_density", minimumDensity}, {"minimum_enthalpy", minimumEnthalpy},
            {"negative_density_samples", static_cast<double>(negativeDensitySamples)},
            {"negative_enthalpy_samples", static_cast<double>(negativeEnthalpySamples)},
            {"minimum_mapping_determinant", minimumDeterminant}
        };
        result["volume_radius_relative_error"] = std::abs(result.at("volume_radius") / reference.radius - 1.0);
        for (const auto &[name, error] : std::array<std::pair<const char *, const ErrorIntegral *>, 5>{{
                 {"density", &densityError}, {"enthalpy", &enthalpyError}, {"potential", &potentialError},
                 {"pressure", &pressureError}, {"gravity_gradient", &gravityError}}}) {
            result[std::string(name) + "_relative_l2_error"] = error->RelativeL2();
            result[std::string(name) + "_maximum_scaled_error"] = error->maximumScaledError;
        }
        for (int component = 0; component < 3; ++component) {
            result["center_of_mass_" + std::to_string(component)] = static_cast<double>(firstMoment[component] / mass);
        }
        return result;
    }

    template <typename SampleState>
    Measurements MeasureSurfaceAndCorners(SampleState &state, const N1Reference &reference, const int order) {
        auto &mesh = *state.finiteElements.mesh;
        PhysicalPoint mapped;
        double minimumRadius = std::numeric_limits<double>::infinity(), maximumRadius = 0.0;
        double maximumSurfaceEnthalpy = 0.0, maximumSurfacePotentialError = 0.0;
        long double area = 0.0L, radiusIntegral = 0.0L, radiusError = 0.0L;
        std::uint64_t boundarySamples = 0, invalidCornerSamples = 0, cornerSamples = 0;
        double minimumCornerDeterminant = std::numeric_limits<double>::infinity();
        double maximumCornerCondition = 0.0;
        for (int boundary = 0; boundary < mesh.GetNBE(); ++boundary) {
            if (mesh.GetBdrAttribute(boundary) != 1) continue; // Canonical sandbox stellar surface.
            // The tagged stellar surface is an interior material interface
            // when a vacuum region is present. Resolve the actual mesh face
            // to retain both adjacent traces in that case.
            const int faceIndex = mesh.GetBdrElementFaceIndex(boundary);
            if (faceIndex < 0) throw std::runtime_error("Missing stellar surface mesh face.");
            auto *face = mesh.GetFaceElementTransformations(faceIndex);
            if (face == nullptr || face->Elem1 == nullptr) throw std::runtime_error("Missing stellar surface transformation.");
            const bool first = state.isStellar(face->Elem1No);
            if (!first && (face->Elem2 == nullptr || !state.isStellar(face->Elem2No))) {
                throw std::runtime_error("Stellar surface has no stellar-material trace.");
            }
            const int element = first ? face->Elem1No : face->Elem2No;
            const auto &rule = mfem::IntRules.Get(face->GetGeometryType(), order);
            for (int q = 0; q < rule.GetNPoints(); ++q) {
                const auto &ip = rule.IntPoint(q);
                // Field evaluation may reuse MFEM's cached element transforms;
                // restore both adjacent face traces before each quadrature point.
                face = mesh.GetFaceElementTransformations(faceIndex);
                face->SetAllIntPoints(&ip);
                const auto volumePoint = first ? face->Elem1->GetIntPoint() : face->Elem2->GetIntPoint();
                mfem::DenseMatrix referenceFaceJacobian(face->Jacobian());
                if (state.Evaluate(element, volumePoint, mapped) != mean_field::mapping::MappingStatus::valid) {
                    throw std::runtime_error("Invalid mapping on stellar surface.");
                }
                mfem::DenseMatrix physicalFaceJacobian(3, 2);
                mfem::Mult(mapped.mapping.mapping_jacobian, referenceFaceJacobian, physicalFaceJacobian);
                const double weight = ip.weight * physicalFaceJacobian.Weight();
                const double radius = mapped.mapping.physical_position.Norml2();
                if (!(weight > 0.0) || !std::isfinite(weight) || !std::isfinite(radius) ||
                    !std::isfinite(mapped.h) || !std::isfinite(mapped.phi)) {
                    throw std::runtime_error("Non-finite field or non-positive physical surface integration weight.");
                }
                area += weight;
                radiusIntegral += weight * radius;
                radiusError += weight * (radius - reference.radius) * (radius - reference.radius);
                minimumRadius = std::min(minimumRadius, radius);
                maximumRadius = std::max(maximumRadius, radius);
                maximumSurfaceEnthalpy = std::max(maximumSurfaceEnthalpy, std::abs(mapped.h) / reference.CentralEnthalpy());
                maximumSurfacePotentialError = std::max(maximumSurfacePotentialError,
                    std::abs(mapped.phi + reference.CentralEnthalpy()) / reference.CentralEnthalpy());
                ++boundarySamples;
            }
        }
        for (int element = 0; element < mesh.GetNE(); ++element) {
            if (!state.isStellar(element)) continue;
            const auto geometry = mesh.GetElementBaseGeometry(element);
            const auto &vertices = *mfem::Geometries.GetVertices(geometry);
            const auto &center = mfem::Geometries.GetCenter(geometry);
            for (int vertex = 0; vertex < vertices.GetNPoints(); ++vertex) {
                for (const double inset : {0.0, 0.005, 0.02}) {
                    const auto &v = vertices.IntPoint(vertex);
                    mfem::IntegrationPoint ip;
                    ip.Set3((1.0-inset)*v.x+inset*center.x, (1.0-inset)*v.y+inset*center.y, (1.0-inset)*v.z+inset*center.z);
                    ++cornerSamples;
                    if (state.Evaluate(element, ip, mapped) != mean_field::mapping::MappingStatus::valid) {
                        ++invalidCornerSamples;
                        continue;
                    }
                    auto *transformation = mesh.GetElementTransformation(element);
                    transformation->SetIntPoint(&ip);
                    mfem::DenseMatrix totalJacobian(3);
                    mfem::Mult(mapped.mapping.mapping_jacobian, transformation->Jacobian(), totalJacobian);
                    minimumCornerDeterminant = std::min(minimumCornerDeterminant, mapped.mapping.mapping_determinant);
                    const double determinant = totalJacobian.Det();
                    const double smallest = totalJacobian.CalcSingularvalue(2);
                    const double largest = totalJacobian.CalcSingularvalue(0);
                    if (!(determinant > 0.0) || !std::isfinite(determinant) || !(smallest > 0.0) ||
                        !std::isfinite(smallest) || !std::isfinite(largest)) ++invalidCornerSamples;
                    else maximumCornerCondition = std::max(maximumCornerCondition, largest / smallest);
                }
            }
        }
        if (!(area > 0.0L) || !std::isfinite(area)) throw std::runtime_error("No finite positive stellar surface area.");
        return {{"surface_samples", static_cast<double>(boundarySamples)}, {"surface_area", static_cast<double>(area)},
                {"surface_radius_minimum", minimumRadius}, {"surface_radius_maximum", maximumRadius},
                {"surface_radius_area_mean", static_cast<double>(radiusIntegral / area)},
                {"surface_radius_relative_rms_error", static_cast<double>(std::sqrt(radiusError / area) / reference.radius)},
                {"surface_radius_relative_range", (maximumRadius - minimumRadius) / reference.radius},
                {"surface_enthalpy_maximum_scaled", maximumSurfaceEnthalpy},
                {"surface_potential_maximum_scaled_error", maximumSurfacePotentialError},
                {"stellar_corner_samples", static_cast<double>(cornerSamples)},
                {"invalid_stellar_corner_samples", static_cast<double>(invalidCornerSamples)},
                {"minimum_stellar_corner_mapping_determinant", minimumCornerDeterminant},
                {"maximum_stellar_corner_element_condition", maximumCornerCondition}};
    }
} // namespace experiment::polytrope_validation