#pragma once

// Include after import mean_field. These element-local projections diagnose the
// weak closure at a frozen geometry and enthalpy; they are never state updates.
#include <algorithm>
#include <array>
#include <cmath>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "polytrope_physical_state.hpp"

namespace experiment::plateau_diagnostics {

struct DensityProjectionControl final {
    std::string label;
    // -1 denotes the production query, whose resolved order is recorded for
    // each element. Explicit controls use ordinary MFEM rules of this order.
    int quadratureOrder;
    bool productionQuadrature;
    std::unique_ptr<mfem::ParGridFunction> density;
};

struct NegativeDensityProjectionReport final {
    // The spaces remain borrowed from PhysicalState::finiteElements. Values on
    // vacuum elements are zero storage, not a physical extension of the star.
    std::vector<DensityProjectionControl> controls;
};

namespace negative_density_projection_detail {

inline double Norm(const mfem::Vector& vector) {
    long double squared = 0.0L;
    for (int i = 0; i < vector.Size(); ++i) {
        if (!std::isfinite(vector(i))) throw std::domain_error("Nonfinite projection vector.");
        squared += static_cast<long double>(vector(i)) * vector(i);
    }
    const double norm = static_cast<double>(std::sqrt(squared));
    if (!std::isfinite(norm)) throw std::domain_error("Projection norm overflow.");
    return norm;
}

inline double Relative(const double numerator, const double denominator) {
    return denominator > 0.0 ? numerator / denominator
                            : (numerator == 0.0 ? 0.0 : std::numeric_limits<double>::infinity());
}

inline void WriteVector(const std::filesystem::path& path, const mfem::Vector& vector) {
    std::ofstream stream(path);
    stream.exceptions(std::ios::failbit | std::ios::badbit);
    stream << std::setprecision(17) << vector.Size() << '\n';
    for (int i = 0; i < vector.Size(); ++i) stream << vector(i) << '\n';
}

// The independent, small Cholesky solve keeps this diagnostic separate from
// the production density inverse. A bad mass matrix must not be regularized:
// that would change the equation whose negative solution we are investigating.
inline mfem::DenseMatrix Factor(const mfem::DenseMatrix& mass) {
    const int size = mass.Height();
    mfem::DenseMatrix lower(size);
    lower = 0.0;
    double scale = 0.0;
    for (int i = 0; i < size; ++i) scale = std::max(scale, mass(i, i));
    if (!(scale > 0.0) || !std::isfinite(scale))
        throw std::domain_error("Invalid density projection mass diagonal.");
    for (int i = 0; i < size; ++i) {
        for (int j = 0; j <= i; ++j) {
            long double value = mass(i, j);
            for (int k = 0; k < j; ++k) value -= static_cast<long double>(lower(i, k)) * lower(j, k);
            if (i == j) {
                if (!std::isfinite(value) || value <= 64.0L * std::numeric_limits<double>::epsilon() * scale)
                    throw std::domain_error("Density projection mass is singular or numerically unresolved.");
                lower(i, j) = static_cast<double>(std::sqrt(value));
            } else {
                lower(i, j) = static_cast<double>(value / lower(j, j));
                if (!std::isfinite(lower(i, j))) throw std::domain_error("Nonfinite mass factor.");
            }
        }
    }
    return lower;
}

inline mfem::Vector Solve(const mfem::DenseMatrix& lower, const mfem::Vector& rhs) {
    mfem::Vector result(rhs);
    for (int i = 0; i < result.Size(); ++i) {
        long double value = result(i);
        for (int j = 0; j < i; ++j) value -= static_cast<long double>(lower(i, j)) * result(j);
        result(i) = static_cast<double>(value / lower(i, i));
    }
    for (int i = result.Size() - 1; i >= 0; --i) {
        long double value = result(i);
        for (int j = i + 1; j < result.Size(); ++j) value -= static_cast<long double>(lower(j, i)) * result(j);
        result(i) = static_cast<double>(value / lower(i, i));
    }
    (void)Norm(result);
    return result;
}

inline mfem::Vector Defect(const mfem::DenseMatrix& mass, const mfem::Vector& coefficients,
                          const mfem::Vector& rhs) {
    mfem::Vector result(rhs.Size());
    for (int i = 0; i < rhs.Size(); ++i) {
        long double value = -static_cast<long double>(rhs(i));
        for (int j = 0; j < coefficients.Size(); ++j)
            value += static_cast<long double>(mass(i, j)) * coefficients(j);
        result(i) = static_cast<double>(value);
    }
    return result;
}

inline void VerifySolve(const mfem::DenseMatrix& mass, const mfem::Vector& coefficients,
                        const mfem::Vector& rhs, const mfem::Vector& defect) {
    long double massSquared = 0.0L;
    for (int i = 0; i < mass.Height(); ++i)
        for (int j = 0; j < mass.Width(); ++j)
            massSquared += static_cast<long double>(mass(i, j)) * mass(i, j);
    const long double scale = Norm(rhs) + std::sqrt(massSquared) * Norm(coefficients);
    if (!std::isfinite(scale) || Norm(defect) > 1.0e-11L * scale)
        throw std::runtime_error("Local density projection failed its independent normal-equation check.");
}

inline bool Identical(const mfem::Vector& first, const mfem::Vector& second) {
    return first.Size() == second.Size() &&
           (first.Size() == 0 || std::memcmp(first.GetData(), second.GetData(),
                                           sizeof(double) * static_cast<std::size_t>(first.Size())) == 0);
}

struct Totals final {
    int elements{0};
    long double volume{0.0L}, projectedSavedSquared{0.0L}, projectedTargetSquared{0.0L}, savedTargetSquared{0.0L};
    long double projectedDefectSquared{0.0L}, savedDefectSquared{0.0L};
    double minimumSavedCoefficient{std::numeric_limits<double>::infinity()};
    double minimumProjectedCoefficient{std::numeric_limits<double>::infinity()};
    double maximumCoefficientDifference{0.0}, maximumConstantCoefficientError{0.0}, maximumConstantSampleError{0.0};
};

} // namespace negative_density_projection_detail

inline NegativeDensityProjectionReport MeasureNegativeDensityProjection(
    experiment::polytrope_validation::PhysicalState& state, const double n, const double K,
    const std::filesystem::path& output) {
    namespace detail = negative_density_projection_detail;
    using namespace mean_field;
    const eos::Polytrope equationOfState(n, K);
    const auto& fem = state.finiteElements;
    if (!fem.mesh || !fem.densityFes || !fem.enthalpyFes || !fem.compactificationCoordinate ||
        !fem.quadratureFactory || fem.mesh->Dimension() != 3 || !fem.mesh->Conforming() ||
        fem.densityFes->GetVDim() != 1 ||
        dynamic_cast<const mfem::L2_FECollection*>(fem.densityFes->FEColl()) == nullptr ||
        state.density.ParFESpace() != fem.densityFes.get() || state.enthalpy.ParFESpace() != fem.enthalpyFes.get())
        throw std::invalid_argument("Density projection diagnostics require the current scalar conforming L2 density space.");
    int ranks = 0;
    MPI_Comm_size(fem.mesh->GetComm(), &ranks);
    if (ranks != 1) throw std::invalid_argument("Density projection diagnostics require one MPI rank.");

    const double extra = (n - 1.0) * static_cast<double>(field::Enthalpy::Scalar::familyOrder);
    if (!std::isfinite(extra) || extra < 0.0 || extra > std::numeric_limits<int>::max())
        throw std::invalid_argument("Invalid production EOS quadrature order.");
    const int extraOrder = static_cast<int>(std::ceil(extra));
    std::filesystem::create_directories(output);
    const auto elementsPath = output / "negative_density_projection_elements.csv";
    const auto summaryPath = output / "negative_density_projection_summary.csv";
    const auto productionResidualPath = output / "projection_production_closure_residual_full.txt";
    if (std::filesystem::exists(elementsPath) || std::filesystem::exists(summaryPath) ||
        std::filesystem::exists(productionResidualPath))
        throw std::invalid_argument("Density projection diagnostics require fresh output files.");

    NegativeDensityProjectionReport report;
    for (const auto& [label, order] : std::array<std::pair<const char*, int>, 3>{
             std::pair{"production", -1}, std::pair{"q18", 18}, std::pair{"q26", 26}}) {
        const auto vectorPath = output / (std::string("negative_density_projection_") + label + "_coefficients.txt");
        if (std::filesystem::exists(vectorPath))
            throw std::invalid_argument("Density projection diagnostics would overwrite saved coefficients.");
        auto projected = std::make_unique<mfem::ParGridFunction>(fem.densityFes.get());
        *projected = 0.0;
        report.controls.push_back({label, order, order < 0, std::move(projected)});
    }
    const mfem::Vector savedDensity(state.density), savedEnthalpy(state.enthalpy), savedPotential(state.potential),
        savedGravity(state.gravityGradientReference), savedDisplacement(state.displacement);
    // Even the observer's mapping cache is private here. No trial preparation,
    // normalization, accepted buffer, or prepared operator is involved.
    mapping::GridFunctionMappingEvaluator mapper(state.domainMapper, state.displacement, *fem.compactificationCoordinate);
    mfem::ParGridFunction constant(fem.densityFes.get());
    mfem::ConstantCoefficient one(1.0);
    constant.ProjectCoefficient(one);
    mfem::Vector fullClosureResidual(fem.densityFes->GetVSize());
    fullClosureResidual = 0.0;
    std::array<detail::Totals, 3> totals;
    std::vector<bool> visited(static_cast<std::size_t>(fem.densityFes->GetVSize()), false);

    std::ofstream elements(elementsPath), summary(summaryPath);
    elements.exceptions(std::ios::failbit | std::ios::badbit);
    summary.exceptions(std::ios::failbit | std::ios::badbit);
    elements << std::setprecision(17);
    summary << std::setprecision(17) << "control,metric,value\n";
    elements << "control,requested_order,quadrature_order,quadrature_points,element,attribute,density_order,enthalpy_order,"
                "basis_family,dofs,volume,minimum_h,minimum_target_density,saved_coefficient_minimum,saved_coefficient_maximum,"
                "projected_coefficient_minimum,projected_coefficient_maximum,coefficient_difference_norm,coefficient_relative_difference,"
                "coefficient_maximum_difference,rhs_norm,projected_normal_defect_norm,projected_normal_relative_defect,"
                "saved_normal_defect_norm,saved_normal_relative_defect,projection_vs_saved_l2,projection_vs_target_l2,saved_vs_target_l2,"
                "minimum_projected_sampled_density,minimum_saved_sampled_density,constant_normal_defect_norm,"
                "constant_coefficient_error_maximum,constant_sample_error_maximum\n";

    for (int element = 0; element < fem.mesh->GetNE(); ++element) {
        if (!state.isStellar(element)) continue;
        const auto* densityElement = fem.densityFes->GetFE(element);
        const auto* enthalpyElement = fem.enthalpyFes->GetFE(element);
        if (!densityElement || !enthalpyElement || densityElement->GetMapType() != mfem::FiniteElement::VALUE ||
            densityElement->GetOrder() != field::Density::Scalar::familyOrder ||
            enthalpyElement->GetOrder() != field::Enthalpy::Scalar::familyOrder)
            throw std::invalid_argument("Density projection FE elements do not match the production scalar fields.");
        mfem::Array<int> dofs;
        if (fem.densityFes->GetElementDofs(element, dofs) != nullptr)
            throw std::invalid_argument("Density projection does not support transformed element DOFs.");
        if (dofs.Size() != densityElement->GetDof()) throw std::logic_error("Density element layout mismatch.");
        for (int i = 0; i < dofs.Size(); ++i) {
            if (dofs[i] < 0 || dofs[i] >= fullClosureResidual.Size() || visited[static_cast<std::size_t>(dofs[i])])
                throw std::invalid_argument("Density projection requires independent, unsigned element-local DOFs.");
            visited[static_cast<std::size_t>(dofs[i])] = true;
        }
        mfem::Vector saved, constantExpected;
        state.density.GetSubVector(dofs, saved);
        constant.GetSubVector(dofs, constantExpected);
        (void)detail::Norm(saved);
        auto* transformation = fem.mesh->GetElementTransformation(element);

        for (std::size_t controlIndex = 0; controlIndex < report.controls.size(); ++controlIndex) {
            auto& control = report.controls[controlIndex];
            auto& total = totals[controlIndex];
            const mfem::IntegrationRule* rule = nullptr;
            int resolvedOrder = control.quadratureOrder;
            if (control.productionQuadrature) {
                // Reproduce the shared production rule, including the density
                // mass degree when it exceeds the nonlinear source degree.
                const auto query = field::Field<field::Enthalpy>::make_query<field::Enthalpy::Form::EosClosureSource>(
                    quadrature::QuadratureRole::discretization, transformation->OrderW(), std::array<int, 1>{extraOrder},
                    utils::DOMAINS::STELLAR, quadrature::MappingKind::general);
                const auto sourceResolution = fem.quadratureFactory->get(query, transformation->GetGeometryType());
                const auto massQuery = field::Field<field::Density>::make_query<field::Density::Form::EosClosureMass>(
                    quadrature::QuadratureRole::discretization, transformation->OrderW(), {},
                    utils::DOMAINS::STELLAR, quadrature::MappingKind::general);
                const auto massResolution = fem.quadratureFactory->get(massQuery, transformation->GetGeometryType());
                const auto &resolution = sourceResolution.resolution.order >= massResolution.resolution.order
                                             ? sourceResolution
                                             : massResolution;
                rule = resolution.integration_rule;
                resolvedOrder = resolution.resolution.order;
            } else {
                rule = &mfem::IntRules.Get(transformation->GetGeometryType(), resolvedOrder);
            }
            if (!rule || rule->GetNPoints() == 0) throw std::runtime_error("Missing density projection quadrature.");
            const int size = dofs.Size(), points = rule->GetNPoints();
            mfem::DenseMatrix mass(size), samples(points, size);
            mfem::Vector rhs(size), constantRhs(size), shape(size), weights(points), targets(points);
            mass = 0.0;
            rhs = 0.0;
            constantRhs = 0.0;
            double minimumH = std::numeric_limits<double>::infinity();
            double minimumTarget = minimumH;
            long double volume = 0.0L;
            for (int q = 0; q < points; ++q) {
                const auto& ip = rule->IntPoint(q);
                mapping::MappingPointContext mapped;
                if (mapper.EvaluatePoint(*transformation, ip, mapped) != mapping::MappingStatus::valid || mapped.compactified)
                    throw std::domain_error("Invalid stellar mapping during density projection.");
                const double h = state.enthalpy.GetValue(element, ip);
                if (!std::isfinite(h) || h < 0.0)
                    throw std::domain_error("Density projection encountered negative or nonfinite enthalpy; EOS was not clipped.");
                const double target = eos::evaluate<eos::quantity::Density>(
                    equationOfState, dimensions::SpecificEnthalpyValue{h}).value();
                transformation->SetIntPoint(&ip);
                const double referenceDeterminant = transformation->Jacobian().Det();
                const double weight = ip.weight * referenceDeterminant * mapped.mapping_determinant;
                if (!(referenceDeterminant > 0.0) || !std::isfinite(referenceDeterminant) ||
                    !(mapped.mapping_determinant > 0.0) || !std::isfinite(mapped.mapping_determinant) ||
                    !(weight > 0.0) || !std::isfinite(weight) || target < 0.0 || !std::isfinite(target))
                    throw std::domain_error("Invalid density projection weight or EOS value.");
                densityElement->CalcShape(ip, shape);
                (void)detail::Norm(shape);
                weights(q) = weight;
                targets(q) = target;
                volume += weight;
                minimumH = std::min(minimumH, h);
                minimumTarget = std::min(minimumTarget, target);
                for (int i = 0; i < size; ++i) {
                    samples(q, i) = shape(i);
                    rhs(i) += weight * shape(i) * target;
                    constantRhs(i) += weight * shape(i);
                    for (int j = 0; j < size; ++j) mass(i, j) += weight * shape(i) * shape(j);
                }
            }
            const auto lower = detail::Factor(mass);
            const auto coefficients = detail::Solve(lower, rhs);
            const auto constantCoefficients = detail::Solve(lower, constantRhs);
            const auto projectedDefect = detail::Defect(mass, coefficients, rhs);
            const auto savedDefect = detail::Defect(mass, saved, rhs);
            const auto constantDefect = detail::Defect(mass, constantCoefficients, constantRhs);
            detail::VerifySolve(mass, coefficients, rhs, projectedDefect);
            detail::VerifySolve(mass, constantCoefficients, constantRhs, constantDefect);
            control.density->SetSubVector(dofs, coefficients);
            if (control.productionQuadrature) fullClosureResidual.SetSubVector(dofs, savedDefect);

            mfem::Vector difference(coefficients);
            difference -= saved;
            double minimumSaved = std::numeric_limits<double>::infinity(), maximumSaved = -minimumSaved;
            double minimumProjected = minimumSaved, maximumProjected = -minimumSaved;
            double maximumDifference = 0.0, constantCoefficientError = 0.0, constantSampleError = 0.0;
            for (int i = 0; i < size; ++i) {
                minimumSaved = std::min(minimumSaved, saved(i));
                maximumSaved = std::max(maximumSaved, saved(i));
                minimumProjected = std::min(minimumProjected, coefficients(i));
                maximumProjected = std::max(maximumProjected, coefficients(i));
                maximumDifference = std::max(maximumDifference, std::abs(difference(i)));
                constantCoefficientError = std::max(constantCoefficientError, std::abs(constantCoefficients(i) - constantExpected(i)));
            }
            long double projectedSavedSquared = 0.0L, projectedTargetSquared = 0.0L, savedTargetSquared = 0.0L;
            double minimumProjectedSample = std::numeric_limits<double>::infinity(), minimumSavedSample = minimumProjectedSample;
            for (int q = 0; q < points; ++q) {
                long double projectedValue = 0.0L, savedValue = 0.0L, constantValue = 0.0L;
                for (int i = 0; i < size; ++i) {
                    projectedValue += static_cast<long double>(samples(q, i)) * coefficients(i);
                    savedValue += static_cast<long double>(samples(q, i)) * saved(i);
                    constantValue += static_cast<long double>(samples(q, i)) * constantCoefficients(i);
                }
                const long double differenceValue = projectedValue - savedValue;
                const long double projectedError = projectedValue - targets(q), savedError = savedValue - targets(q);
                projectedSavedSquared += weights(q) * differenceValue * differenceValue;
                projectedTargetSquared += weights(q) * projectedError * projectedError;
                savedTargetSquared += weights(q) * savedError * savedError;
                minimumProjectedSample = std::min(minimumProjectedSample, static_cast<double>(projectedValue));
                minimumSavedSample = std::min(minimumSavedSample, static_cast<double>(savedValue));
                constantSampleError = std::max(constantSampleError, static_cast<double>(std::abs(constantValue - 1.0L)));
            }
            if (!std::isfinite(constantCoefficientError) || !std::isfinite(constantSampleError) ||
                constantCoefficientError > 1.0e-10 || constantSampleError > 1.0e-10)
                throw std::runtime_error("Density projection failed independent constant reconstruction.");
            const double rhsNorm = detail::Norm(rhs), projectedDefectNorm = detail::Norm(projectedDefect);
            const double savedDefectNorm = detail::Norm(savedDefect);
            elements << control.label << ',' << control.quadratureOrder << ',' << resolvedOrder << ',' << points << ','
                     << element << ',' << fem.mesh->GetAttribute(element) << ',' << densityElement->GetOrder() << ','
                     << enthalpyElement->GetOrder() << ',' << std::quoted(fem.densityFes->FEColl()->Name()) << ',' << size << ','
                     << volume << ',' << minimumH << ',' << minimumTarget << ',' << minimumSaved << ',' << maximumSaved << ','
                     << minimumProjected << ',' << maximumProjected << ',' << detail::Norm(difference) << ','
                     << detail::Relative(detail::Norm(difference), detail::Norm(saved)) << ',' << maximumDifference << ','
                     << rhsNorm << ',' << projectedDefectNorm << ',' << detail::Relative(projectedDefectNorm, rhsNorm) << ','
                     << savedDefectNorm << ',' << detail::Relative(savedDefectNorm, rhsNorm) << ','
                     << std::sqrt(projectedSavedSquared) << ',' << std::sqrt(projectedTargetSquared) << ','
                     << std::sqrt(savedTargetSquared) << ',' << minimumProjectedSample << ',' << minimumSavedSample << ','
                     << detail::Norm(constantDefect) << ',' << constantCoefficientError << ',' << constantSampleError << '\n';
            ++total.elements;
            total.volume += volume;
            total.projectedSavedSquared += projectedSavedSquared;
            total.projectedTargetSquared += projectedTargetSquared;
            total.savedTargetSquared += savedTargetSquared;
            total.projectedDefectSquared += static_cast<long double>(projectedDefectNorm) * projectedDefectNorm;
            total.savedDefectSquared += static_cast<long double>(savedDefectNorm) * savedDefectNorm;
            total.minimumSavedCoefficient = std::min(total.minimumSavedCoefficient, minimumSaved);
            total.minimumProjectedCoefficient = std::min(total.minimumProjectedCoefficient, minimumProjected);
            total.maximumCoefficientDifference = std::max(total.maximumCoefficientDifference, maximumDifference);
            total.maximumConstantCoefficientError = std::max(total.maximumConstantCoefficientError, constantCoefficientError);
            total.maximumConstantSampleError = std::max(total.maximumConstantSampleError, constantSampleError);
        }
    }
    const bool fieldsIdentical = detail::Identical(savedDensity, state.density) && detail::Identical(savedEnthalpy, state.enthalpy) &&
        detail::Identical(savedPotential, state.potential) && detail::Identical(savedGravity, state.gravityGradientReference) &&
        detail::Identical(savedDisplacement, state.displacement);
    if (!fieldsIdentical) throw std::logic_error("Density projection diagnostic changed a source field.");
    summary << "all,source_fields_bitwise_identical,1\nall,trial_preparations,0\nall,mpi_ranks,1\n"
               "all,unsupported_elements_are_zero_storage,1\nall,negative_enthalpy_is_rejected,1\n"
               "all,sampled_extrema_are_not_polynomial_certificates,1\n";
    for (std::size_t index = 0; index < report.controls.size(); ++index) {
        const auto& control = report.controls[index];
        const auto& total = totals[index];
        if (total.elements == 0) throw std::invalid_argument("Density projection found no stellar elements.");
        const auto metric = [&](const char* name, const auto value) { summary << control.label << ',' << name << ',' << value << '\n'; };
        metric("elements", total.elements);
        metric("local_mass_solves", 2 * total.elements);
        metric("volume", total.volume);
        metric("projection_vs_saved_l2", std::sqrt(total.projectedSavedSquared));
        metric("projection_vs_target_l2", std::sqrt(total.projectedTargetSquared));
        metric("saved_vs_target_l2", std::sqrt(total.savedTargetSquared));
        metric("projected_normal_defect_norm", std::sqrt(total.projectedDefectSquared));
        metric("saved_normal_defect_norm", std::sqrt(total.savedDefectSquared));
        metric("saved_coefficient_minimum", total.minimumSavedCoefficient);
        metric("projected_coefficient_minimum", total.minimumProjectedCoefficient);
        metric("coefficient_maximum_difference", total.maximumCoefficientDifference);
        metric("constant_coefficient_error_maximum", total.maximumConstantCoefficientError);
        metric("constant_sample_error_maximum", total.maximumConstantSampleError);
        detail::WriteVector(output / ("negative_density_projection_" + control.label + "_coefficients.txt"), *control.density);
    }
    detail::WriteVector(productionResidualPath, fullClosureResidual);
    return report;
}

} // namespace experiment::plateau_diagnostics