#pragma once

#include "model_slice_geometry.hpp"
#include "polytrope_physical_state.hpp"

#include <chrono>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>

namespace experiment::model_slices {
struct Options final {
    std::string plane{"both"};
    double offset{0.};
    int subdivisions{6};
    void Validate() const {
        if (subdivisions < 2 || subdivisions > 32)
            throw std::invalid_argument("Slice subdivisions must be between 2 and 32.");
        (void)Plane::Parse(plane == "both" ? "xz" : plane, offset);
    }
};

namespace detail {
inline std::ofstream Open(const std::filesystem::path& path) {
    if (std::filesystem::exists(path)) throw std::invalid_argument("Refusing to overwrite slice output: " + path.string());
    std::ofstream stream(path);
    stream.exceptions(std::ios::failbit | std::ios::badbit);
    stream << std::setprecision(17);
    return stream;
}
inline bool Identical(const mfem::Vector& a, const mfem::Vector& b) {
    return a.Size() == b.Size() && (a.Size() == 0 ||
        std::memcmp(a.GetData(), b.GetData(), std::size_t(a.Size()) * sizeof(double)) == 0);
}
inline mfem::IntegrationPoint IntegrationPoint(const Vector& reference) {
    mfem::IntegrationPoint point;
    point.Set3(reference[0], reference[1], reference[2]);
    return point;
}

class ElementMap final {
public:
    ElementMap(polytrope_validation::PhysicalState& state, int element, bool deformed = true)
        : state(state), deformed(deformed), transformation(*state.finiteElements.mesh->GetElementTransformation(element)) {}
    Point operator()(const Vector& reference) {
        const auto point = IntegrationPoint(reference);
        Vector physical{};
        if (deformed) {
            if (state.mapping.EvaluatePoint(transformation, point, mapped) != mean_field::mapping::MappingStatus::valid || mapped.compactified)
                throw std::domain_error("Model slices require valid, noncompactified stellar geometry.");
            for (int d = 0; d < 3; ++d) physical[d] = mapped.physical_position(d);
        } else {
            transformation.Transform(point, referencePosition);
            for (int d = 0; d < 3; ++d) physical[d] = referencePosition(d);
        }
        for (double value : physical) if (!std::isfinite(value)) throw std::domain_error("Nonfinite slice geometry.");
        return {reference, physical};
    }
private:
    polytrope_validation::PhysicalState& state;
    bool deformed;
    mfem::ElementTransformation& transformation;
    mean_field::mapping::MappingPointContext mapped;
    mfem::Vector referencePosition = mfem::Vector(3);
};

class Intersections final {
public:
    Intersections(ElementMap& map, Plane plane, double tolerance) : map(map), plane(std::move(plane)), tolerance(tolerance) {}
    Point operator()(const Point& a, const Point& b) {
        const Vector& first = a.reference < b.reference ? a.reference : b.reference;
        const Vector& last = a.reference < b.reference ? b.reference : a.reference;
        const std::array<double, 6> key{first[0], first[1], first[2], last[0], last[1], last[2]};
        if (const auto found = cache.find(key); found != cache.end()) return found->second;
        const Point root = Intersect(a, b, plane, tolerance, map);
        cache.emplace(key, root);
        return root;
    }
private:
    ElementMap& map;
    Plane plane;
    double tolerance;
    std::map<std::array<double, 6>, Point> cache;
};

struct Bounds final {
    Vector lower{INFINITY, INFINITY, INFINITY}, upper{-INFINITY, -INFINITY, -INFINITY};
    void Add(const Vector& value) {
        for (int d = 0; d < 3; ++d) { lower[d] = std::min(lower[d], value[d]); upper[d] = std::max(upper[d], value[d]); }
    }
};

inline void ExportPlane(polytrope_validation::PhysicalState& state, const Options& options, const Plane& plane,
                        const std::filesystem::path& directory) {
    const auto start = std::chrono::steady_clock::now();
    if (!std::filesystem::create_directory(directory)) throw std::invalid_argument("Slice directory already exists.");
    auto vertices = Open(directory / "vertices.csv"), triangles = Open(directory / "triangles.csv");
    vertices << "vertex,element,attribute,xi,eta,zeta,x,y,z,rho,h,plane_error\n";
    triangles << "triangle,element,v0,v1,v2\n";
    auto& mesh = *state.finiteElements.mesh;
    const int n = options.subdivisions, width = n + 1;
    const auto index = [width](int i, int j, int k) { return i + width * (j + width * k); };
    Bounds bounds;
    std::size_t vertexCount = 0, triangleCount = 0;
    double maximumPlaneError = 0., maximumChordError = 0., minimumDensity = INFINITY, minimumEnthalpy = INFINITY;
    // Coordinate-scaled tolerance is independent of density units and order.
    double scale = 0.;
    for (int element = 0; element < mesh.GetNE(); ++element) if (state.isStellar(element)) {
        ElementMap map(state, element);
        for (int corner = 0; corner < 8; ++corner) {
            const auto point = map({double(corner & 1), double((corner >> 1) & 1), double((corner >> 2) & 1)});
            scale = std::max(scale, std::hypot(point.physical[0], point.physical[1], point.physical[2]));
        }
    }
    if (!(scale > 0) || !std::isfinite(scale)) throw std::domain_error("Invalid stellar slice scale.");
    const double tolerance = 2e-12 * std::max(scale, std::abs(plane.offset));
    for (int element = 0; element < mesh.GetNE(); ++element) if (state.isStellar(element)) {
        if (mesh.GetElementBaseGeometry(element) != mfem::Geometry::CUBE)
            throw std::invalid_argument("Model slices currently require hexahedral stellar elements.");
        ElementMap map(state, element);
        Intersections intersections(map, plane, tolerance);
        std::vector<Point> lattice(static_cast<std::size_t>(width * width * width));
        for (int k = 0; k <= n; ++k) for (int j = 0; j <= n; ++j) for (int i = 0; i <= n; ++i) {
            auto& point = lattice[index(i, j, k)];
            point = map({double(i) / n, double(j) / n, double(k) / n});
            bounds.Add(point.physical);
        }
        for (int k = 0; k < n; ++k) for (int j = 0; j < n; ++j) for (int i = 0; i < n; ++i) {
            std::array<Point, 8> cube;
            for (int c = 0; c < 8; ++c) cube[c] = lattice[index(i + (c & 1), j + ((c >> 1) & 1), k + ((c >> 2) & 1))];
            for (const auto& indices : cubeTetrahedra) {
                std::array<Point, 4> tetrahedron;
                for (int c = 0; c < 4; ++c) tetrahedron[c] = cube[indices[c]];
                for (const auto& triangle : CutTetrahedron(tetrahedron, plane, tolerance, intersections)) {
                    triangles << triangleCount++ << ',' << element << ',' << vertexCount << ',' << vertexCount + 1 << ',' << vertexCount + 2 << '\n';
                    Vector referenceCenter{}, physicalCenter{};
                    for (const auto& point : triangle) {
                        const auto ip = IntegrationPoint(point.reference);
                        const double rho = state.density.GetValue(element, ip), h = state.enthalpy.GetValue(element, ip);
                        if (!std::isfinite(rho) || !std::isfinite(h)) throw std::domain_error("Nonfinite slice field value.");
                        const double error = std::abs(plane.Distance(point));
                        if (error > tolerance) throw std::logic_error("A slice vertex is off its physical plane.");
                        maximumPlaneError = std::max(maximumPlaneError, error);
                        minimumDensity = std::min(minimumDensity, rho); minimumEnthalpy = std::min(minimumEnthalpy, h);
                        vertices << vertexCount++ << ',' << element << ',' << mesh.GetAttribute(element);
                        for (double value : point.reference) vertices << ',' << value;
                        for (double value : point.physical) vertices << ',' << value;
                        vertices << ',' << rho << ',' << h << ',' << error << '\n';
                        for (int d = 0; d < 3; ++d) {
                            referenceCenter[d] += point.reference[d] / 3.; physicalCenter[d] += point.physical[d] / 3.;
                        }
                        bounds.Add(point.physical);
                    }
                    maximumChordError = std::max(maximumChordError, Distance(map(referenceCenter).physical, physicalCenter));
                }
            }
        }
    }
    vertices.flush(); triangles.flush();

    std::array<std::size_t, 2> segmentCounts{};
    std::array<int, 2> coplanar{};
    std::array<double, 2> surfaceError{}, surfaceChordError{};
    for (int reference = 0; reference < 2; ++reference) {
        auto surface = Open(directory / (reference ? "reference_surface.csv" : "surface.csv"));
        surface << "segment,face,element,x0,y0,z0,x1,y1,z1,maximum_plane_error\n";
        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), rightStar = right >= 0 && state.isStellar(right);
            if (leftStar == rightStar) continue;
            const int element = leftStar ? left : right;
            auto* faceMap = mesh.GetFaceElementTransformations(face);
            if (faceMap->GetGeometryType() != mfem::Geometry::SQUARE)
                throw std::invalid_argument("Model slices currently require quadrilateral material faces.");
            ElementMap map(state, element, !reference);
            Intersections intersections(map, plane, tolerance);
            std::vector<Point> lattice(static_cast<std::size_t>(width * width));
            for (int j = 0; j <= n; ++j) for (int i = 0; i <= n; ++i) {
                mfem::IntegrationPoint facePoint, elementPoint;
                facePoint.Set2(double(i) / n, double(j) / n);
                if (leftStar) faceMap->Loc1.Transform(facePoint, elementPoint); else faceMap->Loc2.Transform(facePoint, elementPoint);
                lattice[i + width * j] = map({elementPoint.x, elementPoint.y, elementPoint.z});
            }
            for (int j = 0; j < n; ++j) for (int i = 0; i < n; ++i) {
                const std::array<Point, 4> square{lattice[i + width * j], lattice[i + 1 + width * j],
                    lattice[i + width * (j + 1)], lattice[i + 1 + width * (j + 1)]};
                for (const auto& indices : {std::array<int, 3>{0, 1, 3}, std::array<int, 3>{0, 3, 2}}) {
                    const Triangle triangle{square[indices[0]], square[indices[1]], square[indices[2]]};
                    for (const auto& segment : CutSurfaceTriangle(triangle, plane, tolerance, intersections, coplanar[reference])) {
                        const double error = std::max(std::abs(plane.Distance(segment[0])), std::abs(plane.Distance(segment[1])));
                        if (error > tolerance) throw std::logic_error("A surface intersection is off its physical plane.");
                        surfaceError[reference] = std::max(surfaceError[reference], error);
                        surfaceChordError[reference] = std::max(surfaceChordError[reference],
                            Distance(map(Blend(segment[0].reference, segment[1].reference, .5)).physical,
                                     Blend(segment[0].physical, segment[1].physical, .5)));
                        surface << segmentCounts[reference]++ << ',' << face << ',' << element;
                        for (const auto& point : segment) { for (double value : point.physical) surface << ',' << value; bounds.Add(point.physical); }
                        surface << ',' << error << '\n';
                    }
                }
            }
        }
    }
    auto metadata = Open(directory / "metadata.csv");
    metadata << "metric,value\nplane," << plane.name << "\nnormal_axis," << plane.normal
             << "\nhorizontal_axis," << plane.horizontal << "\nvertical_axis," << plane.vertical
             << "\noffset," << plane.offset << "\nsubdivisions," << n << "\nroot_tolerance," << tolerance
             << "\ndeformation_scale,1\ntriangle_count," << triangleCount << "\nvertex_count," << vertexCount
             << "\nmaximum_plane_error," << maximumPlaneError << "\nmaximum_sampled_chord_error," << maximumChordError
             << "\nminimum_sampled_density," << minimumDensity << "\nminimum_sampled_enthalpy," << minimumEnthalpy
             << "\nsurface_segments," << segmentCounts[0] << "\nreference_surface_segments," << segmentCounts[1]
             << "\nsurface_coplanar_triangles," << coplanar[0] << "\nreference_surface_coplanar_triangles," << coplanar[1]
             << "\nsurface_maximum_plane_error," << surfaceError[0] << "\nreference_surface_maximum_plane_error," << surfaceError[1]
             << "\nsurface_maximum_sampled_chord_error," << surfaceChordError[0]
             << "\nreference_surface_maximum_sampled_chord_error," << surfaceChordError[1]
             << "\nu_min," << bounds.lower[plane.horizontal] << "\nu_max," << bounds.upper[plane.horizontal]
             << "\nv_min," << bounds.lower[plane.vertical] << "\nv_max," << bounds.upper[plane.vertical]
             << "\nseconds," << std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count() << '\n';
    if (triangleCount && !segmentCounts[0]) throw std::runtime_error("Material slice has no resolved outer boundary; refine or move the plane.");
    if (coplanar[0] || coplanar[1]) throw std::runtime_error("Slice is coplanar with a surface patch; choose a different plane offset.");
}
} // namespace detail

inline void Record(polytrope_validation::PhysicalState& state, const Options& options,
                   const std::filesystem::path& output, const std::string& label) {
    options.Validate();
    int ranks = 0;
    MPI_Comm_size(state.finiteElements.mesh->GetComm(), &ranks);
    if (ranks != 1 || state.finiteElements.mesh->Dimension() != 3 || !state.finiteElements.mesh->Conforming())
        throw std::invalid_argument("Model slices require a conforming 3D mesh on one MPI rank.");
    const mfem::Vector density(state.density), enthalpy(state.enthalpy), displacement(state.displacement);
    const auto directory = output / (label + "_slices");
    if (!std::filesystem::create_directory(directory)) throw std::invalid_argument("Slice output directory already exists.");
    auto checks = detail::Open(directory / "geometry_self_checks.csv");
    checks << "case,observed,expected,tolerance,passed\n";
    for (const auto& check : GeometrySelfChecks())
        checks << check.name << ',' << check.observed << ',' << check.expected << ',' << check.tolerance << ",1\n";
    const std::vector<std::string> planes = options.plane == "both" ? std::vector<std::string>{"xz", "xy"} : std::vector<std::string>{options.plane};
    for (const auto& name : planes) {
        std::cout << label << " physical " << name << " slice, subdivisions=" << options.subdivisions << std::endl;
        detail::ExportPlane(state, options, Plane::Parse(name, options.offset), directory / name);
    }
    if (!detail::Identical(density, state.density) || !detail::Identical(enthalpy, state.enthalpy) || !detail::Identical(displacement, state.displacement))
        throw std::logic_error("Model slice export changed a source field.");
    auto scope = detail::Open(directory / "scope.txt");
    scope << "source_fields_bitwise_identical=1\ndeformation_scale=1\n"
          << "Vertices evaluate the independent FEM density and enthalpy on the actual deformed physical plane.\n"
          << "Triangles remain local to their material element; exterior fields are never sampled.\n"
          << "A coincident internal face uses the positive-normal-side DG trace.\n"
          << "Surface and undeformed reference outlines are separately intersected with the same physical plane.\n"
          << "Interior colors and boundary chords approximate curved cuts at the recorded subdivision count.\n"
          << "Sampled chord errors and bounds are diagnostics, not rigorous geometric enclosures.\n"
          << "Refine subdivisions to assess image stability; subcell tangencies or multiple edge crossings can be missed.\n"
          << "The slice does not certify positivity, equilibrium, virial balance or the absence of unsampled features.\n";
}
} // namespace experiment::model_slices