#pragma once

// Diagnostic-only extension of the Q2 slice observer. A monotonic partition
// from derivative roots handles degrees three and four without sampling a sign
// indicator or relying on a particular closed-form cubic/quartic formula.
#include "negative_density_regions.hpp"

#include <string_view>

namespace experiment::plateau_diagnostics::density_order_detail {
namespace region = negative_density_regions_detail;
using Real = region::Real;

struct Polynomial final {
    std::array<Real, 5> coefficients{};
    int degree{0};
    [[nodiscard]] Real Value(Real x) const {
        Real value = coefficients[degree], correction = 0;
        for (int j = degree - 1; j >= 0; --j) {
            const Real product = value * x;
            const Real productError = std::fma(value, x, -product);
            const Real sum = product + coefficients[j];
            const Real sumError = std::abs(product) >= std::abs(coefficients[j])
                ? (product - sum) + coefficients[j] : (coefficients[j] - sum) + product;
            correction = std::fma(correction, x, productError + sumError);
            value = sum;
        }
        // Compensated Horner retains the sign of shallow pockets without
        // assuming long double has a wider mantissa on Apple arm64.
        return value + correction;
    }
    [[nodiscard]] Real Scale() const {
        Real result = 0;
        for (int j = 0; j <= degree; ++j) result += std::abs(coefficients[j]);
        return result;
    }
    [[nodiscard]] Polynomial Derivative() const {
        Polynomial result;
        result.degree = std::max(0, degree - 1);
        for (int j = 1; j <= degree; ++j) result.coefficients[j - 1] = j * coefficients[j];
        return result;
    }
    [[nodiscard]] Real Integral(Real left, Real right) const {
        constexpr int binomial[5][5]{{1}, {1, 1}, {1, 2, 1}, {1, 3, 3, 1}, {1, 4, 6, 4, 1}};
        const Real width = right - left, middle = left + width / 2;
        Real mean = Value(middle);
        for (int k = 2; k <= degree; k += 2) {
            Real translated = 0;
            for (int j = degree; j >= k; --j)
                translated = std::fma(translated, middle, binomial[j][k] * coefficients[j]);
            mean += translated * std::pow(width / 2, k) / (k + 1);
        }
        return width * mean;
    }
};

inline Polynomial Interpolate(const std::array<double, 5>& values, int degree) {
    if (degree < 2 || degree > 4) throw std::invalid_argument("Density slice degree must be two, three, or four.");
    std::array<Real, 5> nodes{}, divided{}, basis{};
    for (int i = 0; i <= degree; ++i) {
        nodes[i] = static_cast<double>(i) / degree;
        divided[i] = values[i];
        if (!std::isfinite(values[i])) throw std::domain_error("Nonfinite density interpolation value.");
    }
    for (int width = 1; width <= degree; ++width)
        for (int i = degree; i >= width; --i)
            divided[i] = (divided[i] - divided[i - 1]) / (nodes[i] - nodes[i - width]);
    Polynomial result;
    result.degree = degree;
    basis[0] = 1;
    for (int k = 0; k <= degree; ++k) {
        for (int j = 0; j <= k; ++j) result.coefficients[j] += divided[k] * basis[j];
        if (k == degree) break;
        for (int j = k + 1; j > 0; --j) basis[j] = basis[j - 1] - nodes[k] * basis[j];
        basis[0] *= -nodes[k];
    }
    return result;
}

inline std::vector<Real> Roots(Polynomial polynomial) {
    while (polynomial.degree > 0 && polynomial.coefficients[polynomial.degree] == 0) --polynomial.degree;
    const Real scale = polynomial.Scale();
    if (!std::isfinite(scale)) throw std::domain_error("Nonfinite density polynomial coefficients.");
    if (polynomial.degree == 0 || scale == 0) return {};
    int exponent = 0;
    (void)std::frexp(scale, &exponent);
    // A power-of-two rescaling preserves the represented polynomial. Dividing
    // by an arbitrary norm can create a pocket at an exact repeated root.
    for (int j = 0; j <= polynomial.degree; ++j)
        polynomial.coefficients[j] = std::scalbn(polynomial.coefficients[j], -exponent);
    if (polynomial.degree == 1) {
        const Real root = -polynomial.coefficients[0] / polynomial.coefficients[1];
        return std::isfinite(root) && root >= 0 && root <= 1 ? std::vector<Real>{root} : std::vector<Real>{};
    }
    std::vector<Real> partition{0, 1};
    for (Real point : Roots(polynomial.Derivative())) if (point > 0 && point < 1) partition.push_back(point);
    std::sort(partition.begin(), partition.end());
    partition.erase(std::unique(partition.begin(), partition.end()), partition.end());
    std::vector<Real> result;
    for (Real point : partition) if (polynomial.Value(point) == 0) result.push_back(point);
    for (std::size_t i = 1; i < partition.size(); ++i) {
        Real left = partition[i - 1], right = partition[i];
        Real fleft = polynomial.Value(left), fright = polynomial.Value(right);
        if (fleft == 0 || fright == 0 || (fleft < 0) == (fright < 0)) continue;
        // No value tolerance declares a critical point to be zero. Doing so
        // could erase both roots of a shallow pocket next to that critical point.
        for (int iteration = 0; iteration < 128; ++iteration) {
            const Real middle = left + (right - left) / 2;
            if (middle == left || middle == right) break;
            const Real value = polynomial.Value(middle);
            if (value == 0) { left = right = middle; break; }
            if ((value < 0) == (fleft < 0)) { left = middle; fleft = value; }
            else { right = middle; fright = value; }
        }
        result.push_back(std::abs(polynomial.Value(left)) <= std::abs(polynomial.Value(right)) ? left : right);
    }
    std::sort(result.begin(), result.end());
    result.erase(std::unique(result.begin(), result.end()), result.end());
    return result;
}

inline std::vector<region::Interval> Partition(const Polynomial& polynomial) {
    std::vector<Real> points{0, 1};
    for (Real root : Roots(polynomial)) if (root > 0 && root < 1) points.push_back(root);
    std::sort(points.begin(), points.end());
    points.erase(std::unique(points.begin(), points.end()), points.end());
    std::vector<region::Interval> result;
    for (std::size_t j = 1; j < points.size(); ++j) {
        const Real left = points[j - 1], right = points[j];
        const bool negative = polynomial.Value(left + (right - left) / 2) < 0;
        for (Real fraction : {Real{0.25}, Real{0.75}}) {
            const Real value = polynomial.Value(left + fraction * (right - left));
            if ((negative && value > 0) || (!negative && value < 0))
                throw std::runtime_error("Density polynomial root partition has inconsistent signs.");
        }
        result.push_back({left, right, negative});
    }
    return result;
}

inline void RootSelfChecks(const std::filesystem::path& path) {
    auto out = region::Open(path);
    out << "case,degree,negative_length,expected_negative_length,negative_mass,expected_negative_mass,passed\n";
    const auto check = [&](const char* name, Polynomial polynomial, Real expectedLength, Real expectedMass) {
        Real length = 0, mass = 0;
        for (const auto& interval : Partition(polynomial)) if (interval.negative) {
            length += interval.right - interval.left;
            mass -= polynomial.Integral(interval.left, interval.right);
        }
        const bool passed = std::abs(length - expectedLength) < 2e-11L * std::max(Real{1e-8}, expectedLength) &&
            std::abs(mass - expectedMass) < 2e-10L * std::max(Real{1e-8}, expectedMass);
        out << name << ',' << polynomial.degree << ',' << length << ',' << expectedLength << ',' << mass << ',' << expectedMass << ',' << passed << '\n';
        if (!passed) throw std::runtime_error(std::string("General density root self-check failed: ") + name);
    };
    check("positive_constant", {{{1}}, 0}, 0, 0);
    check("negative_constant", {{{-2}}, 0}, 1, 2);
    check("linear_crossing", {{{-0.25L, 1}}, 1}, 0.25L, 0.03125L);
    check("quadratic_interior_pocket", {{{0.05L, -0.6L, 1.5L}}, 2}, std::sqrt(Real{6}) / 15, std::sqrt(Real{6}) / 2250);
    check("cubic_interior_pocket", {{{0, 3.0L / 16, -1, 1}}, 3}, 0.5L, 1.0L / 96);
    check("quartic_two_pockets", {{{105.0L / 4096, -11.0L / 32, 43.0L / 32, -2, 1}}, 4}, 0.5L, 19.0L / 15360);
    check("quartic_double_roots_nonnegative", {{{9.0L / 256, -3.0L / 8, 11.0L / 8, -2, 1}}, 4}, 0, 0);
    // Independently check the power-basis reconstruction, including its highest
    // degree, rather than testing only roots of hand-entered coefficient arrays.
    for (int degree : {2, 3, 4}) {
        std::array<double, 5> values{};
        for (int i = 0; i <= degree; ++i) values[i] = std::pow(static_cast<double>(i) / degree, degree);
        const auto p = Interpolate(values, degree);
        for (Real t : {Real{0.13}, Real{0.37}, Real{0.81}})
            if (std::abs(p.Value(t) - std::pow(t, degree)) > 2e-14L)
                throw std::runtime_error("Density polynomial interpolation self-check failed.");
        out << "reconstruct_t_power_" << degree << ',' << degree << ",0,0,0,0,1\n";
    }
}

struct Outputs final {
    std::ofstream elements, summary, intervals, surface;
    explicit Outputs(const std::filesystem::path& output)
        : elements(region::Open(output / "density_order_elements.csv")),
          summary(region::Open(output / "density_order_summary.csv")),
          intervals(region::Open(output / "density_order_intervals.csv")),
          surface(region::Open(output / "density_order_surface.csv")) {
        constexpr auto extra = "target_mass,eos_squared_integral,signed_eos_integral,negative_mass_uncertainty_scale,"
            "first_moment_x,first_moment_y,first_moment_z,inertia_z,minimum_enthalpy,negative_enthalpy_samples,"
            "eos_defined_everywhere_sampled,resolved_negative_intervals,uncertain_negative_intervals,"
            "resolved_witness_negative_volume,resolved_witness_negative_mass,uncertain_negative_volume,uncertain_negative_mass,";
        elements << "control,degree,tangential_order,interval_order,element,attribute,axis,orientation,surface_face," << extra;
        region::Header(elements);
        summary << "control,degree,tangential_order,interval_order," << extra;
        region::Header(summary);
        intervals << "control,degree,element,axis,orientation,u,v,t_left,t_right,c0,c1,c2,c3,c4,negative_volume,negative_mass,"
                     "minimum_density,witness_h,reconstruction_tolerance,resolved_negative_witness\n";
        surface << "control,degree,face,element,u,v,rho,h,rho_eos,error,eos_defined,reconstruction_tolerance,resolved_negative\n";
    }
};

struct Extra final {
    Real targetMass{0}, eosSquared{0}, signedEos{0}, uncertainty{0};
    std::array<Real, 3> firstMoment{};
    Real inertiaZ{0};
    double minimumH{std::numeric_limits<double>::infinity()};
    std::size_t negativeH{0};
    std::size_t resolvedNegativeIntervals{0}, uncertainNegativeIntervals{0};
    Real resolvedVolume{0}, resolvedMass{0}, uncertainVolume{0}, uncertainMass{0};
    void Add(const Extra& other) {
        targetMass += other.targetMass; eosSquared += other.eosSquared; signedEos += other.signedEos; uncertainty += other.uncertainty;
        inertiaZ += other.inertiaZ; minimumH = std::min(minimumH, other.minimumH); negativeH += other.negativeH;
        resolvedNegativeIntervals += other.resolvedNegativeIntervals; uncertainNegativeIntervals += other.uncertainNegativeIntervals;
        resolvedVolume += other.resolvedVolume; resolvedMass += other.resolvedMass;
        uncertainVolume += other.uncertainVolume; uncertainMass += other.uncertainMass;
        for (int d = 0; d < 3; ++d) firstMoment[d] += other.firstMoment[d];
    }
    void Write(std::ostream& out) const {
        const Real nan = std::numeric_limits<Real>::quiet_NaN();
        out << (negativeH == 0 ? targetMass : nan) << ',' << (negativeH == 0 ? eosSquared : nan) << ','
            << (negativeH == 0 ? signedEos : nan) << ',' << uncertainty << ',';
        for (Real value : firstMoment) out << value << ',';
        out << inertiaZ << ',' << minimumH << ',' << negativeH << ',' << (negativeH == 0) << ','
            << resolvedNegativeIntervals << ',' << uncertainNegativeIntervals << ',' << resolvedVolume << ',' << resolvedMass << ','
            << uncertainVolume << ',' << uncertainMass << ',';
    }
};

inline double EosTarget(double h, double n, double K) {
    if (!std::isfinite(h)) throw std::domain_error("Nonfinite enthalpy in density order observer.");
    const double rho = h > 0 ? std::pow(h / ((n + 1) * K), n) : 0;
    if (!std::isfinite(rho)) throw std::domain_error("Nonfinite EOS target in density order observer.");
    return rho;
}

inline void Observe(polytrope_validation::PhysicalState& state, const mfem::ParGridFunction& density,
                    int degree, std::string_view label, double n, double K, Outputs& output) {
    const auto& fem = state.finiteElements;
    const auto axes = region::ChooseAxes(state);
    mean_field::mapping::GridFunctionMappingEvaluator mapper(state.domainMapper, state.displacement, *fem.compactificationCoordinate);
    mean_field::mapping::MappingPointContext mapped;
    const auto& intervalRule = mfem::IntRules.Get(mfem::Geometry::SEGMENT, 14);
    for (int order : {14, 22, 30}) {
        const auto& tangentRule = mfem::IntRules.Get(mfem::Geometry::SEGMENT, order);
        region::Measure global;
        Extra globalExtra;
        for (int element = 0; element < fem.mesh->GetNE(); ++element) {
            if (!state.isStellar(element)) continue;
            const auto& axis = axes[static_cast<std::size_t>(element)];
            auto* transformation = fem.mesh->GetElementTransformation(element);
            region::Measure local;
            Extra localExtra;
            for (int iu = 0; iu < tangentRule.GetNPoints(); ++iu) for (int iv = 0; iv < tangentRule.GetNPoints(); ++iv) {
                const auto& up = tangentRule.IntPoint(iu);
                const auto& vp = tangentRule.IntPoint(iv);
                const double u = up.x, v = vp.x;
                std::array<double, 5> values{};
                for (int i = 0; i <= degree; ++i) values[i] = density.GetValue(element, region::Point(axis, u, v, static_cast<double>(i) / degree));
                const auto polynomial = Interpolate(values, degree);
                const Real tolerance = 2048 * std::numeric_limits<double>::epsilon() * polynomial.Scale();
                const auto verify = [&](Real t, double rho) {
                    const Real error = std::abs(polynomial.Value(t) - rho);
                    local.maximumReconstructionError = std::max(local.maximumReconstructionError, static_cast<double>(error));
                    local.maximumReconstructionTolerance = std::max(local.maximumReconstructionTolerance, static_cast<double>(tolerance));
                    if (!std::isfinite(rho) || error > tolerance) throw std::runtime_error("Density order slice reconstruction failed.");
                };
                for (Real t : {Real{0.07}, Real{0.19}, Real{0.41}, Real{0.63}, Real{0.87}, Real{0.97}})
                    verify(t, density.GetValue(element, region::Point(axis, u, v, t)));
                ++local.lines;
                const auto critical = Roots(polynomial.Derivative());
                const auto evaluate = [&](Real t, region::Measure& piece, bool negative) {
                    const auto point = region::Point(axis, u, v, t);
                    if (mapper.EvaluatePoint(*transformation, point, mapped) != mean_field::mapping::MappingStatus::valid || mapped.compactified)
                        throw std::domain_error("Invalid density order mapping.");
                    const double rho = density.GetValue(element, point), h = state.enthalpy.GetValue(element, point);
                    verify(t, rho);
                    if (!std::isfinite(h)) throw std::domain_error("Nonfinite density order enthalpy.");
                    if (rho < piece.minimumRho) {
                        piece.minimumRho = rho; piece.witnessH = h; piece.witnessElement = element;
                        piece.witnessT = static_cast<double>(t); piece.witnessU = u; piece.witnessV = v;
                        for (int d = 0; d < 3; ++d) piece.witnessPosition[d] = mapped.physical_position(d);
                    }
                    if (negative) {
                        const double radius = mapped.physical_position.Norml2();
                        piece.negativeRadiusMin = std::min(piece.negativeRadiusMin, radius);
                        piece.negativeRadiusMax = std::max(piece.negativeRadiusMax, radius);
                    }
                    transformation->SetIntPoint(&point);
                    const double reference = transformation->Jacobian().Det();
                    const Real determinant = static_cast<Real>(reference) * mapped.mapping_determinant;
                    if (!(reference > 0) || !std::isfinite(reference) || !(mapped.mapping_determinant > 0) ||
                        !std::isfinite(mapped.mapping_determinant) || !std::isfinite(determinant))
                        throw std::domain_error("Invalid density order physical measure.");
                    return std::pair{determinant, h};
                };
                for (const auto& part : Partition(polynomial)) {
                    region::Measure piece;
                    (void)evaluate(part.left, piece, part.negative);
                    (void)evaluate(part.right, piece, part.negative);
                    for (Real point : critical) if (point > part.left && point < part.right) (void)evaluate(point, piece, part.negative);
                    for (int iq = 0; iq < intervalRule.GetNPoints(); ++iq) {
                        const auto& qp = intervalRule.IntPoint(iq);
                        const Real t = part.left + (part.right - part.left) * qp.x;
                        const auto [determinant, h] = evaluate(t, piece, part.negative);
                        const Real weight = static_cast<Real>(up.weight) * vp.weight * qp.weight * (part.right - part.left) * determinant;
                        const Real rho = polynomial.Value(t), target = EosTarget(h, n, K);
                        if (!(weight > 0) || !std::isfinite(weight) || !std::isfinite(rho) ||
                            (part.negative && rho > 0) || (!part.negative && rho < 0))
                            throw std::runtime_error("Density order interval quadrature has an inconsistent sign.");
                        piece.volume += weight; piece.mass += weight * rho; ++piece.points;
                        localExtra.targetMass += weight * target;
                        localExtra.eosSquared += weight * (rho - target) * (rho - target);
                        localExtra.signedEos += weight * (rho - target);
                        localExtra.uncertainty += weight * tolerance;
                        localExtra.minimumH = std::min(localExtra.minimumH, h);
                        localExtra.negativeH += h < 0;
                        for (int d = 0; d < 3; ++d) localExtra.firstMoment[d] += weight * rho * mapped.physical_position(d);
                        localExtra.inertiaZ += weight * rho * (static_cast<Real>(mapped.physical_position(0)) * mapped.physical_position(0) +
                                                              static_cast<Real>(mapped.physical_position(1)) * mapped.physical_position(1));
                    }
                    if (part.negative) {
                        piece.negativeIntervals = 1; piece.negativeVolume = piece.volume; piece.negativeMass = -piece.mass;
                        piece.negativeTMin = part.left; piece.negativeTMax = part.right;
                        const bool resolved = piece.minimumRho < -tolerance;
                        if (resolved) {
                            ++localExtra.resolvedNegativeIntervals;
                            localExtra.resolvedVolume += piece.negativeVolume; localExtra.resolvedMass += piece.negativeMass;
                        } else {
                            ++localExtra.uncertainNegativeIntervals;
                            localExtra.uncertainVolume += piece.negativeVolume; localExtra.uncertainMass += piece.negativeMass;
                        }
                        if (order == 30) {
                            output.intervals << label << ',' << degree << ',' << element << ',' << axis.coordinate << ',' << axis.orientation
                                << ',' << u << ',' << v << ',' << part.left << ',' << part.right;
                            for (Real coefficient : polynomial.coefficients) output.intervals << ',' << coefficient;
                            output.intervals << ',' << piece.negativeVolume << ',' << piece.negativeMass << ',' << piece.minimumRho << ','
                                << piece.witnessH << ',' << tolerance << ',' << resolved << '\n';
                        }
                    }
                    local.Add(piece);
                }
            }
            if (!(local.volume > 0) || !std::isfinite(local.mass) || !std::isfinite(local.negativeMass))
                throw std::runtime_error("Nonfinite density order region integral.");
            output.elements << label << ',' << degree << ',' << order << ",14," << element << ',' << fem.mesh->GetAttribute(element) << ','
                << axis.coordinate << ',' << axis.orientation << ',' << axis.face << ',';
            localExtra.Write(output.elements); region::Write(output.elements, local);
            global.Add(local); globalExtra.Add(localExtra);
        }
        if (!(global.volume > 0)) throw std::runtime_error("Density order observer found no stellar volume.");
        output.summary << label << ',' << degree << ',' << order << ",14,";
        globalExtra.Write(output.summary); region::Write(output.summary, global);
        output.summary.flush(); output.elements.flush(); output.intervals.flush();
    }
    // These are one-sided field traces on every stellar/nonstellar face, not
    // surface-area samples. They expose extrapolation defects missed by nodes.
    for (int face = 0; face < fem.mesh->GetNumFaces(); ++face) {
        int left = -1, right = -1;
        fem.mesh->GetFaceElements(face, &left, &right);
        const bool leftStar = left >= 0 && state.isStellar(left), rightStar = right >= 0 && state.isStellar(right);
        if (leftStar == rightStar) continue;
        const int element = leftStar ? left : right;
        auto* transformation = fem.mesh->GetFaceElementTransformations(face);
        mfem::Array<int> dofs;
        auto* dofTransform = density.ParFESpace()->GetElementDofs(element, dofs);
        mfem::Vector coefficients, shape(dofs.Size());
        density.GetSubVector(dofs, coefficients);
        if (dofTransform) dofTransform->InvTransformPrimal(coefficients);
        const auto* fe = density.ParFESpace()->GetFE(element);
        for (double u : {0., .25, .5, .75, 1.}) for (double v : {0., .25, .5, .75, 1.}) {
            mfem::IntegrationPoint facePoint, point;
            facePoint.Set2(u, v);
            if (leftStar) transformation->Loc1.Transform(facePoint, point); else transformation->Loc2.Transform(facePoint, point);
            const double rho = density.GetValue(element, point), h = state.enthalpy.GetValue(element, point);
            if (!std::isfinite(h)) throw std::domain_error("Nonfinite density order surface enthalpy.");
            const double target = h >= 0 ? EosTarget(h, n, K) : std::numeric_limits<double>::quiet_NaN();
            if (!std::isfinite(rho)) throw std::domain_error("Nonfinite density order surface trace.");
            fe->CalcShape(point, shape);
            Real magnitude = 0;
            for (int i = 0; i < shape.Size(); ++i) magnitude += std::abs(static_cast<Real>(shape(i)) * coefficients(i));
            const Real tolerance = 2048 * std::numeric_limits<double>::epsilon() * magnitude;
            output.surface << label << ',' << degree << ',' << face << ',' << element << ',' << u << ',' << v << ','
                << rho << ',' << h << ',' << target << ',' << rho - target << ',' << (h >= 0) << ',' << tolerance << ',' << (rho < -tolerance) << '\n';
        }
    }
    output.surface.flush();
}

} // namespace experiment::plateau_diagnostics::density_order_detail