#pragma once

// Experiment-only observations of a frozen physical state. No sampled field,
// projection, or positivity condition is installed in the equilibrium problem.
#include "negative_density_projection.hpp"

#include <array>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <stdexcept>
#include <string>
#include <vector>

namespace experiment::plateau_diagnostics {
    namespace negative_density_detail {
        inline std::ofstream Open(const std::filesystem::path &path) {
            if (std::filesystem::exists(path))
                throw std::runtime_error("Diagnostic output already exists: " + path.string());
            std::ofstream out(path);
            out.exceptions(std::ios::badbit | std::ios::failbit);
            out << std::setprecision(17);
            return out;
        }

        struct SurfaceFace {
            int face, element, axis;
            double coordinate;
        };

        struct Topology {
            std::vector<SurfaceFace> surface;
            std::vector<int> layer;
            std::vector<std::vector<int>> faces;
        };

        inline Topology BuildTopology(polytrope_validation::PhysicalState &state) {
            auto &mesh = *state.finiteElements.mesh;
            Topology result;
            result.layer.assign(mesh.GetNE(), -1);
            result.faces.resize(mesh.GetNE());
            std::vector<std::vector<int>> neighbors(mesh.GetNE());
            std::queue<int> queue;
            for (int face = 0; face < mesh.GetNumFaces(); ++face) {
                int left = -1, right = -1;
                mesh.GetFaceElements(face, &left, &right);
                const bool leftStar  = left >= 0 && state.isStellar(left);
                const bool rightStar = right >= 0 && state.isStellar(right);
                if (leftStar && rightStar) {
                    neighbors[left].push_back(right);
                    neighbors[right].push_back(left);
                }
                if (leftStar == rightStar)
                    continue;
                const int element    = leftStar ? left : right;
                auto *transformation = mesh.GetFaceElementTransformations(face);
                if (transformation == nullptr)
                    throw std::runtime_error("Missing stellar interface transformation.");
                mfem::IntegrationPoint facePoint, point;
                facePoint.Set2(0.5, 0.5);
                if (leftStar)
                    transformation->Loc1.Transform(facePoint, point);
                else
                    transformation->Loc2.Transform(facePoint, point);
                const std::array<double, 3> coordinate{point.x, point.y, point.z};
                int axis     = -1;
                double value = 0.;
                for (int d = 0; d < 3; ++d) {
                    if (std::abs(coordinate[d]) < 1e-12 || std::abs(coordinate[d] - 1.) < 1e-12) {
                        if (axis >= 0)
                            throw std::runtime_error("Ambiguous reference surface face.");
                        axis  = d;
                        value = std::round(coordinate[d]);
                    }
                }
                if (axis < 0)
                    throw std::runtime_error("Density localization currently requires tensor-product cells.");
                result.faces[element].push_back(static_cast<int>(result.surface.size()));
                result.surface.push_back({face, element, axis, value});
                if (result.layer[element] < 0) {
                    result.layer[element] = 0;
                    queue.push(element);
                }
            }
            while (!queue.empty()) {
                const int element = queue.front();
                queue.pop();
                for (const int neighbor : neighbors[element])
                    if (result.layer[neighbor] < 0) {
                        result.layer[neighbor] = result.layer[element] + 1;
                        queue.push(neighbor);
                    }
            }
            for (int element = 0; element < mesh.GetNE(); ++element)
                if (state.isStellar(element) && result.layer[element] < 0)
                    throw std::runtime_error("Stellar component is disconnected from its surface.");
            return result;
        }

        struct Aggregate {
            std::size_t samples = 0, negative = 0, negativeH = 0, jointNegative = 0;
            long double volume = 0., mass = 0., negativeMass = 0., negativeVolume = 0., negativeHVolume = 0.,
                        closure2 = 0.;
            double minimum = INFINITY, maximum = -INFINITY, minimumH = INFINITY;
            double negativeHMinimum = INFINITY, negativeHMaximum = -INFINITY;
            double negativeRadiusMinimum = INFINITY, negativeRadiusMaximum = -INFINITY;
            double negativeReferenceDepthMinimum = INFINITY, negativeReferenceDepthMaximum = -INFINITY;
            void
            Add(double rho,
                double h,
                double rhoH,
                double radius,
                double depth,
                double weight) {
                ++samples;
                volume += weight;
                mass += weight * rho;
                closure2 += weight * (rho - rhoH) * (rho - rhoH);
                minimum  = std::min(minimum, rho);
                maximum  = std::max(maximum, rho);
                minimumH = std::min(minimumH, h);
                if (h < 0.) {
                    ++negativeH;
                    negativeHVolume += weight;
                }
                if (rho >= 0.)
                    return;
                ++negative;
                if (h < 0.)
                    ++jointNegative;
                negativeMass -= weight * rho;
                negativeVolume += weight;
                negativeHMinimum      = std::min(negativeHMinimum, h);
                negativeHMaximum      = std::max(negativeHMaximum, h);
                negativeRadiusMinimum = std::min(negativeRadiusMinimum, radius);
                negativeRadiusMaximum = std::max(negativeRadiusMaximum, radius);
                if (std::isfinite(depth)) {
                    negativeReferenceDepthMinimum = std::min(negativeReferenceDepthMinimum, depth);
                    negativeReferenceDepthMaximum = std::max(negativeReferenceDepthMaximum, depth);
                }
            }
        };

        inline double Observed(double v) {
            return std::isfinite(v) ? v : std::numeric_limits<double>::quiet_NaN();
        }
        inline void Header(std::ostream &out) {
            out << "samples,negative_samples,negative_h_samples,joint_negative_samples,volume,mass,negative_mass,"
                   "negative_volume,negative_h_volume,closure_squared_integral,minimum_density,maximum_density,minimum_"
                   "h,minimum_h_at_negative_rho,maximum_h_at_negative_rho,negative_radius_minimum,negative_radius_"
                   "maximum,negative_reference_depth_minimum,negative_reference_depth_maximum\n";
        }
        inline void Write(
            std::ostream &out,
            const Aggregate &a
        ) {
            out << a.samples << ',' << a.negative << ',' << a.negativeH << ',' << a.jointNegative << ',' << a.volume
                << ',' << a.mass << ',' << a.negativeMass << ',' << a.negativeVolume << ',' << a.negativeHVolume << ','
                << a.closure2 << ',' << Observed(a.minimum) << ',' << Observed(a.maximum) << ',' << Observed(a.minimumH)
                << ',' << Observed(a.negativeHMinimum) << ',' << Observed(a.negativeHMaximum) << ','
                << Observed(a.negativeRadiusMinimum) << ',' << Observed(a.negativeRadiusMaximum) << ','
                << Observed(a.negativeReferenceDepthMinimum) << ',' << Observed(a.negativeReferenceDepthMaximum)
                << '\n';
        }

        inline double ReferenceDepth(
            const Topology &topology,
            int element,
            const mfem::IntegrationPoint &point
        ) {
            const std::array<double, 3> coordinates{point.x, point.y, point.z};
            double depth = INFINITY;
            for (const int index : topology.faces[element]) {
                const auto &face = topology.surface[index];
                depth            = std::min(depth, std::abs(coordinates[face.axis] - face.coordinate));
            }
            return Observed(depth);
        }
    } // namespace negative_density_detail

    inline void RecordNegativeDensity(
        polytrope_validation::PhysicalState &state,
        double n,
        double K,
        const std::filesystem::path &output
    ) {
        namespace d            = negative_density_detail;
        const auto topology    = d::BuildTopology(state);
        const auto projections = MeasureNegativeDensityProjection(state, n, K, output);
        auto &mesh             = *state.finiteElements.mesh;
        auto topologyFile      = d::Open(output / "density_element_topology.csv");
        topologyFile << "element,attribute,surface_layer,surface_face_count,density_order,enthalpy_order,density_dof_"
                        "minimum,density_dof_maximum\n";
        mfem::Array<int> dofs;
        mfem::Vector coefficients;
        for (int element = 0; element < mesh.GetNE(); ++element)
            if (state.isStellar(element)) {
                auto *transform = state.finiteElements.densityFes->GetElementVDofs(element, dofs);
                state.density.GetSubVector(dofs, coefficients);
                if (transform)
                    transform->InvTransformPrimal(coefficients);
                topologyFile << element << ',' << mesh.GetAttribute(element) << ',' << topology.layer[element] << ','
                             << topology.faces[element].size() << ','
                             << state.finiteElements.densityFes->GetFE(element)->GetOrder() << ','
                             << state.finiteElements.enthalpyFes->GetFE(element)->GetOrder() << ','
                             << coefficients.Min() << ',' << coefficients.Max() << '\n';
            }
        auto globalFile = d::Open(output / "density_summary.csv");
        globalFile << "quadrature,field,";
        d::Header(globalFile);
        auto elementFile = d::Open(output / "density_elements.csv");
        elementFile << "quadrature,field,element,attribute,surface_layer,";
        d::Header(elementFile);
        auto radialFile = d::Open(output / "density_radial_bins.csv");
        radialFile << "quadrature,field,radial_bin_left,radial_bin_right,";
        d::Header(radialFile);
        auto negativeFile = d::Open(output / "negative_density_samples.csv");
        negativeFile << "quadrature,element,attribute,surface_layer,xi,eta,zeta,x,y,z,radius,reference_inward_depth,"
                        "weight,rho,h,rho_from_h";
        for (const auto &control : projections.controls)
            negativeFile << ",rho_projected_" << control.label;
        negativeFile << '\n';
        std::vector<std::string> labels{"state", "eos_pointwise"};
        for (const auto &control : projections.controls)
            labels.push_back(control.label);
        const double factor = (n + 1.) * K;
        mean_field::mapping::MappingPointContext mapped;
        for (const int order : {18, 26}) {
            std::cout << "Density localization quadrature " << order << std::endl;
            std::vector<d::Aggregate> global(labels.size());
            std::vector<std::map<int, d::Aggregate>> radial(labels.size());
            for (int element = 0; element < mesh.GetNE(); ++element) {
                if (!state.isStellar(element))
                    continue;
                auto *transformation = mesh.GetElementTransformation(element);
                if (transformation->GetGeometryType() != mfem::Geometry::CUBE)
                    throw std::runtime_error("Density localization requires hexahedral elements.");
                const auto &rule = mfem::IntRules.Get(transformation->GetGeometryType(), order);
                std::vector<d::Aggregate> local(labels.size());
                for (int q = 0; q < rule.GetNPoints(); ++q) {
                    const auto &point = rule.IntPoint(q);
                    if (state.mapping.EvaluatePoint(*transformation, point, mapped) !=
                        mean_field::mapping::MappingStatus::valid)
                        throw std::runtime_error("Invalid mapping during density localization.");
                    transformation->SetIntPoint(&point);
                    const double weight = point.weight * transformation->Jacobian().Det() * mapped.mapping_determinant;
                    if (!std::isfinite(weight) || weight <= 0.)
                        throw std::runtime_error("Invalid density observation weight.");
                    const double rho = state.density.GetValue(element, point);
                    const double h   = state.enthalpy.GetValue(element, point);
                    if (!std::isfinite(rho) || !std::isfinite(h))
                        throw std::runtime_error("Nonfinite observed density or enthalpy.");
                    const double rhoH   = h > 0. ? std::pow(h / factor, n) : 0.;
                    const double radius = mapped.physical_position.Norml2();
                    const double depth  = d::ReferenceDepth(topology, element, point);
                    std::vector<double> densities{rho, rhoH};
                    for (const auto &control : projections.controls)
                        densities.push_back(control.density->GetValue(element, point));
                    const int bin = static_cast<int>(std::floor(radius * 200.));
                    for (std::size_t field = 0; field < labels.size(); ++field) {
                        if (!std::isfinite(densities[field]))
                            throw std::runtime_error("Nonfinite projected density.");
                        local[field].Add(densities[field], h, rhoH, radius, depth, weight);
                        global[field].Add(densities[field], h, rhoH, radius, depth, weight);
                        radial[field][bin].Add(densities[field], h, rhoH, radius, depth, weight);
                    }
                    // Retain every negative-state point at the finer rule. Volume
                    // statistics above include all points, so this sparse export
                    // cannot silently redefine the denominator of a fraction.
                    if (rho < 0. && order == 26) {
                        negativeFile << order << ',' << element << ',' << mesh.GetAttribute(element) << ','
                                     << topology.layer[element] << ',' << point.x << ',' << point.y << ',' << point.z;
                        for (int axis = 0; axis < 3; ++axis)
                            negativeFile << ',' << mapped.physical_position(axis);
                        negativeFile << ',' << radius << ',' << depth << ',' << weight << ',' << rho << ',' << h << ','
                                     << rhoH;
                        for (std::size_t field = 2; field < densities.size(); ++field)
                            negativeFile << ',' << densities[field];
                        negativeFile << '\n';
                    }
                }
                for (std::size_t field = 0; field < labels.size(); ++field) {
                    elementFile << order << ',' << labels[field] << ',' << element << ',' << mesh.GetAttribute(element)
                                << ',' << topology.layer[element] << ',';
                    d::Write(elementFile, local[field]);
                }
            }
            for (std::size_t field = 0; field < labels.size(); ++field) {
                globalFile << order << ',' << labels[field] << ',';
                d::Write(globalFile, global[field]);
                for (const auto &[bin, value] : radial[field]) {
                    radialFile << order << ',' << labels[field] << ',' << bin / 200. << ',' << (bin + 1) / 200. << ',';
                    d::Write(radialFile, value);
                }
            }
            globalFile.flush();
            elementFile.flush();
            radialFile.flush();
            negativeFile.flush();
        }

        // Surface-inward traces include the face itself, which volume Gauss rules
        // never sample. The direction is a reference-cell coordinate line, not a
        // physical normal or an assumed spherical radial line.
        auto traces = d::Open(output / "density_surface_traces.csv");
        traces << "face,element,tangent_u,tangent_v,inward_coordinate,x,y,z,radius,distance_from_face,rho,h,rho_from_h";
        for (const auto &control : projections.controls)
            traces << ",rho_projected_" << control.label;
        traces << '\n';
        constexpr std::array<double, 5> tangents{0., .25, .5, .75, 1.};
        constexpr std::array<double, 17> depths{0.,  .001, .003, .007, .01, .02, .03, .04, .06,
                                                .08, .1,   .15,  .2,   .3,  .5,  .75, 1.};
        for (const auto &face : topology.surface) {
            auto *transformation = mesh.GetElementTransformation(face.element);
            const int axis1 = (face.axis + 1) % 3, axis2 = (face.axis + 2) % 3;
            for (const double u : tangents)
                for (const double v : tangents) {
                    mfem::Vector boundaryPosition(3);
                    for (const double depth : depths) {
                        std::array<double, 3> coordinate{};
                        coordinate[face.axis] = face.coordinate == 0. ? depth : 1. - depth;
                        coordinate[axis1]     = u;
                        coordinate[axis2]     = v;
                        mfem::IntegrationPoint point;
                        point.Set3(coordinate[0], coordinate[1], coordinate[2]);
                        if (state.mapping.EvaluatePoint(*transformation, point, mapped) !=
                            mean_field::mapping::MappingStatus::valid)
                            throw std::runtime_error("Invalid surface-inward trace mapping.");
                        if (depth == 0.)
                            boundaryPosition = mapped.physical_position;
                        mfem::Vector distance(mapped.physical_position);
                        distance -= boundaryPosition;
                        const double rho = state.density.GetValue(face.element, point),
                                     h   = state.enthalpy.GetValue(face.element, point);
                        traces << face.face << ',' << face.element << ',' << u << ',' << v << ',' << depth;
                        for (int axis = 0; axis < 3; ++axis)
                            traces << ',' << mapped.physical_position(axis);
                        traces << ',' << mapped.physical_position.Norml2() << ',' << distance.Norml2() << ',' << rho
                               << ',' << h << ',' << (h > 0. ? std::pow(h / factor, n) : 0.);
                        for (const auto &control : projections.controls)
                            traces << ',' << control.density->GetValue(face.element, point);
                        traces << '\n';
                    }
                }
        }
        auto scope = d::Open(output / "density_scope.txt");
        scope << "Frozen observations only; no equilibrium state, operator or field space is modified.\n"
              << "density_basis=" << state.finiteElements.densityFes->FEColl()->Name() << '\n'
              << "enthalpy_basis=" << state.finiteElements.enthalpyFes->FEColl()->Name() << '\n'
              << "surface_faces=" << topology.surface.size() << '\n'
              << "Surface layer is graph distance through stellar face neighbors to a stellar/nonstellar interface.\n"
              << "Reference inward depth is available only in cells touching that interface.\n"
              << "Radial bins are in current physical coordinates with the diagnostic's zero center and unit reference "
                 "radius.\n"
              << "Quadrature extrema and trace extrema are observations, not polynomial sign certificates.\n";
    }
} // namespace experiment::plateau_diagnostics