#include <algorithm>
#include <array>
#include <bit>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <numbers>
#include <string>
#include <vector>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

#include "serif/backend/mfem/discretization/field_values.hpp"
#include "serif/backend/mfem/serialize/resources.hpp"
#include "serif/backend/mfem/serialize/top.hpp"
#include "serif/backend/mfem/solver/seed/state.hpp"
#include "serif/model/finite_elements.hpp"
#include "serif/solver/seed/polytrope.hpp"
#include "serif/tests/solver/stellar_fixture.hpp"
#include "../sandbox/stellar_rayfem.hpp"

namespace {
    namespace fixture = serif::tests::solver;
    namespace serialization = serif::utils::serialize;
    namespace native_archive = serif::backend::mfem::serialize;
    namespace seed = serif::solver::seed;

    /** @brief Supply the exact zero displacement and accepted shape used by a manufactured display state. */
    struct StationaryDisplayCache {
        struct ShapeSnapshot { std::span<const double> parameters; } linearization;
        ::mfem::Vector displacement;

        /** @brief Expose only the two accepted-state records consumed by the display adapter. */
        template <typename Key>
        [[nodiscard]] const auto& get() const {
            if constexpr (std::same_as<Key, serif::geometry::deformation::keys::Linearization>) {
                return linearization;
            } else {
                static_assert(std::same_as<Key, serif::geometry::deformation::backend::mfem::volume_displacement_keys::Coefficients>);
                return displacement;
            }
        }
    };

    /** @brief Keep successful exports and rejected partial files within a test's lifetime. */
    struct ExportDirectory {
        std::filesystem::path path = std::filesystem::temp_directory_path() /
            ("serif-top-mfem-" + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()));
        ~ExportDirectory() { std::error_code error; std::filesystem::remove_all(path, error); }
    };

    /** @brief Small sampling grids test the adapter; their coarse P2 density is not a seismic accuracy benchmark. */
    serialization::TopExportOptions small_export() {
        return {.radial_points = 12, .maximum_degree = 4, .angular_points = 12, .azimuthal_points = 4,
            .relative_symmetry_tolerance = 3.0e-2, .relative_surface_tolerance = 1.0e-6,
            .relative_surface_projection_tolerance = 1.0e-3};
    }

    /** @brief Evaluate the original FE field at the center independently of the exporter and seed scales. */
    template <typename Discretization, typename State>
    double central_density(const Discretization& finite_elements, const State& state) {
        const auto& binding = finite_elements.template binding<fixture::Density>();
        const auto values = state.template get<fixture::Density>();
        ::mfem::Vector reduced(static_cast<int>(values.size()));
        std::ranges::copy(values, reduced.GetData());
        ::mfem::GridFunction density(&binding.space());
        density = binding.scatter(reduced);
        ::mfem::Vector center(3);
        center = 0.0;
        for (int element = 0; element < binding.mesh().GetNE(); ++element) {
            if (!binding.supports(element)) continue;
            ::mfem::InverseElementTransformation inverse(binding.mesh().GetElementTransformation(element));
            ::mfem::IntegrationPoint point;
            if (inverse.Transform(center, point) == ::mfem::InverseElementTransformation::Inside)
                return density.GetValue(element, point);
        }
        FAIL("The independently sampled stellar center must lie in a material element.");
        return 0.0;
    }

    /** @brief Input preservation includes signed zero, as well as all physical and reference mesh coordinates. */
    template <typename State>
    void check_state_unchanged(const State& state, const std::vector<double>& original) {
        REQUIRE(state.values().size() == original.size());
        CHECK(std::ranges::equal(state.values(), original, [](double a, double b) {
            return std::bit_cast<std::uint64_t>(a) == std::bit_cast<std::uint64_t>(b);
        }));
    }

    /** @brief Read the TOP scalar file independently of the serializer's returned report. */
    std::array<double, 9> scalar_file(const std::filesystem::path& directory) {
        std::ifstream input(directory / "lambda_eps");
        REQUIRE(input.good());
        std::array<double, 9> values{};
        for (double& value : values) {
            REQUIRE(static_cast<bool>(input >> value));
            REQUIRE(std::isfinite(value));
        }
        return values;
    }
}

TEST_CASE("TOP exports the current deformed FE star without changing its state or either mesh", "[serialize][top_export][mfem]") {
    constexpr double G = 0.7, K = 0.3, mass = 1.4;
    const auto model = fixture::make_model(1.0, K, mass, 0.03, G);
    const auto radial = seed::make_polytropic_seed(model);
    REQUIRE(radial);
    const auto resources = fixture::make_resources(radial->scales.radius.value());
    const auto finite_elements = fixture::make_discretization(resources, 0);
    auto state = seed::backend::mfem::make_polytropic_state(model, finite_elements, *radial);
    REQUIRE(state);
    // This is a valid manufactured state, not a claimed equilibrium. A nonzero
    // monopole makes using the undeformed source mesh detectably wrong.
    state->template get<fixture::Shape>()[0] = 0.02 * radial->scales.radius.value();
    const double rho_c = central_density(finite_elements, *state);
    REQUIRE(rho_c > 0.0);
    const std::vector<double> original_state(state->values().begin(), state->values().end());
    const auto original_mesh = native_archive::mesh_text(*resources->mesh.mesh);
    const auto original_reference = native_archive::mesh_text(*resources->mesh.reference_mesh);
    const ExportDirectory output;
    const auto result = serialization::export_top(output.path, model, finite_elements, *state, small_export());
    INFO((result ? "Export succeeded" : result.error().path + ": " + result.error().message));
    check_state_unchanged(*state, original_state);
    CHECK(native_archive::mesh_text(*resources->mesh.mesh) == original_mesh);
    CHECK(native_archive::mesh_text(*resources->mesh.reference_mesh) == original_reference);
    REQUIRE(result);
    CHECK(result->equatorial_radius > radial->scales.radius.value() * 1.003);
    CHECK(result->polar_radius == Catch::Approx(result->equatorial_radius).epsilon(5.0e-4));
    CHECK(result->angular_frequency_unit == Catch::Approx(std::sqrt(4.0 * std::numbers::pi * G * rho_c)).epsilon(2.0e-10));
    // At n=1 the density cancels: this tests recovery of both physical G and K.
    CHECK(result->lambda == Catch::Approx(4.0 * std::numbers::pi * G * std::pow(result->equatorial_radius, 2) / (2.0 * K)).epsilon(2.0e-12));
    const auto written = scalar_file(output.path);
    CHECK(written[0] == result->lambda);
    CHECK(written[2] == result->alpha);
    CHECK(written[4] == 1.0);
    CHECK(written[8] == small_export().radial_points);
    CHECK(written[7] / std::sqrt(3.0 * written[2]) ==
        Catch::Approx(state->template get<fixture::Omega>()[0] / result->angular_frequency_unit).epsilon(2.0e-12));
    CHECK(std::filesystem::file_size(output.path / "enthalpy_et_der") > 0);
    CHECK(std::filesystem::file_size(output.path / "domains_boundaries") > 0);
    // The EOS and current density determine the export. A different unsatisfied
    // mass prescription must not change TOP's background normalization.
    const auto other_mass = fixture::make_model(1.0, K, 3.0 * mass, 0.03, G);
    const ExportDirectory other_output;
    const auto other = serialization::export_top(other_output.path, other_mass, finite_elements, *state, small_export());
    INFO((other ? "Changed-prescription export succeeded" : other.error().path + ": " + other.error().message));
    REQUIRE(other);
    CHECK(other->alpha == result->alpha);
    CHECK(other->rotation_over_kepler == result->rotation_over_kepler);
    check_state_unchanged(*state, original_state);
    CHECK(native_archive::mesh_text(*resources->mesh.mesh) == original_mesh);
    CHECK(native_archive::mesh_text(*resources->mesh.reference_mesh) == original_reference);
}

TEST_CASE("TOP sampling honors a tilted rotation axis on a surface and equator refined true-DOF mesh", "[serialize][top_export][mfem][nonconforming]") {
    namespace model = serif::model;
    namespace deformation = serif::geometry::deformation;
    const auto base = fixture::make_model();
    const auto star = model::make_stellar_model(base.equation_of_state(), base.equations(),
        model::FixedTotalMass{serif::dimensions::MassValue{1.0}},
        model::FixedAngularMomentum{model::FixedAngularMomentum::Parameters{
            .angular_momentum = serif::dimensions::AngularMomentumValue{0.04}, .axis = {1.0, 2.0, 3.0}}});
    const auto& axis = star.specification<model::FixedAngularMomentum>().axis();
    const auto radial = seed::make_polytropic_seed(star);
    REQUIRE(radial);
    const auto finite_elements = model::make_stellar_discretization(star,
        {.radius = radial->scales.radius, .geometry_order = 6,
            .surface_refinement = {.levels = 1, .equatorial_levels = 1, .axis = axis}},
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface>{
            .surface = {.maximum_degree = 0}, .stellar_interior = {}, .vacuum = {}});
    const auto& binding = finite_elements.template binding<fixture::Density>();
    REQUIRE(binding.space().GetTrueVSize() < binding.space().GetVSize());
    REQUIRE(binding.constrained());
    auto state = seed::backend::mfem::make_polytropic_state(star, finite_elements, *radial);
    REQUIRE(state);
    // A positive quadrupolar perturbation makes the chosen axis observable:
    // treating this field as axisymmetric about z must not silently pass.
    const double radius = radial->scales.radius.value();
    ::mfem::FunctionCoefficient prescribed_density([&](const ::mfem::Vector& position) {
        const double r = position.Norml2();
        if (r >= radius) return 0.0;
        const double x = std::numbers::pi * r / radius;
        const double profile = (x == 0.0 ? 1.0 : std::sin(x) / x) * radial->scales.central_density.value();
        double axial = 0.0;
        for (int component = 0; component < 3; ++component) axial += axis[component] * position[component];
        return profile * (1.0 + 0.6 * (3.0 * axial * axial - r * r) / (radius * radius));
    });
    ::mfem::GridFunction density(&binding.space());
    density.ProjectCoefficient(prescribed_density);
    const auto reduced = binding.gather(density);
    std::copy_n(reduced.GetData(), reduced.Size(), state->template get<fixture::Density>().begin());
    const auto trace = serif::discretization::backend::mfem::select_boundary_dofs<
        serif::discretization::domain::StellarSurfaceBoundary>(binding);
    for (const int dof : trace.reduced_dofs()) state->template get<fixture::Density>()[dof] = 0.0;
    const std::vector<double> original_state(state->values().begin(), state->values().end());
    const auto original_mesh = native_archive::mesh_text(binding.mesh());
    const ExportDirectory output;
    const auto result = serialization::export_top(output.path, star, finite_elements, *state, small_export());
    INFO((result ? "Export succeeded" : result.error().path + ": " + result.error().message));
    check_state_unchanged(*state, original_state);
    CHECK(native_archive::mesh_text(binding.mesh()) == original_mesh);
    REQUIRE(result);
    CHECK(result->polar_radius == Catch::Approx(radial->scales.radius.value()).epsilon(5.0e-4));
    CHECK(result->equatorial_radius == Catch::Approx(radial->scales.radius.value()).epsilon(5.0e-4));
    CHECK(result->angular_frequency_unit == Catch::Approx(std::sqrt(4.0 * std::numbers::pi * central_density(finite_elements, *state))).epsilon(2.0e-10));
    const auto written = scalar_file(output.path);
    CHECK(written[7] / std::sqrt(3.0 * written[2]) ==
        Catch::Approx(state->template get<fixture::Omega>()[0] / result->angular_frequency_unit).epsilon(2.0e-12));
    const auto wrong_axis = fixture::make_model(1.0, 2.0 / std::numbers::pi, 1.0, 0.04);
    const ExportDirectory wrong_output;
    const auto rejected = serialization::export_top(wrong_output.path, wrong_axis, finite_elements, *state, small_export());
    REQUIRE_FALSE(rejected);
    CHECK_FALSE(std::filesystem::exists(wrong_output.path / "lambda_eps"));
}

TEST_CASE("TOP rejects a finite pressure stellar boundary before publishing vacuum polytrope files", "[serialize][top_export][mfem]") {
    const auto vacuum = fixture::make_model();
    const auto radial = seed::make_polytropic_seed(vacuum);
    REQUIRE(radial);
    const auto resources = fixture::make_resources(radial->scales.radius.value());
    const auto finite_elements = fixture::make_discretization(resources, 0);
    const auto state = seed::backend::mfem::make_polytropic_state(vacuum, finite_elements, *radial);
    REQUIRE(state);
    const auto pressured = fixture::make_model(1.0, 2.0 / std::numbers::pi, 1.0, 0.0, 1.0, 0.01);
    const std::vector<double> original_state(state->values().begin(), state->values().end());
    const auto original_mesh = native_archive::mesh_text(*resources->mesh.mesh);
    const ExportDirectory output;
    const auto result = serialization::export_top(output.path, pressured, finite_elements, *state, small_export());
    REQUIRE_FALSE(result);
    CHECK(result.error().message.find("pressure") != std::string::npos);
    CHECK_FALSE(std::filesystem::exists(output.path / "lambda_eps"));
    check_state_unchanged(*state, original_state);
    CHECK(native_archive::mesh_text(*resources->mesh.mesh) == original_mesh);
}

TEST_CASE("FieldValues samples reconstructed density and its correctly mapped gradient", "[serialize][mfem][reconstruction]") {
    namespace discretization = serif::discretization;
    namespace backend = discretization::backend::mfem;
    using Density = discretization::blocks::density::mass;
    using Schema = discretization::domain::schema::CoreEnvelopeVacuumDomainSchema;
    constexpr double exponent = 4.0 / 3.0;

    ::mfem::Mesh mesh = ::mfem::Mesh::MakeCartesian3D(1, 1, 1, ::mfem::Element::HEXAHEDRON, 2.0, 3.0, 4.0);
    for (int face = 0; face < mesh.GetNBE(); ++face) {
        auto& transformation = *mesh.GetBdrElementTransformation(face);
        ::mfem::Vector center(3);
        transformation.Transform(::mfem::Geometries.GetCenter(transformation.GetGeometryType()), center);
        mesh.GetBdrElement(face)->SetAttribute(center[0] == 0.0 ? 1 : 2);
    }
    ::mfem::H1_FECollection collection(discretization::blocks::SpatialTermTraits<Density>::order, 3);
    ::mfem::FiniteElementSpace space(&mesh, &collection);
    backend::TermBinding<Density, Schema> ordinary(space);
    auto plan = std::make_shared<backend::ReconstructionPlan>(ordinary, 1,
        std::make_shared<discretization::reconstruction::SurfacePowerProfile>(exponent));
    backend::TermBinding<Density, Schema> reconstructed(space, plan);
    ::mfem::FunctionCoefficient profile([](const ::mfem::Vector& position) {
        return std::pow(position[0] / 2.0, exponent) * (1.0 + position[1]) * (2.0 - position[2]);
    });
    ::mfem::GridFunction nodal(&space);
    nodal.ProjectCoefficient(profile);
    const ::mfem::Vector coefficients = reconstructed.gather(nodal);
    backend::FieldValues<Density> values(reconstructed, coefficients);
    backend::FieldValues<Density> polynomial(ordinary, coefficients);
    CHECK_THROWS(backend::FieldValues<discretization::blocks::displacement::vector>(ordinary, coefficients));
    ::mfem::IntegrationPoint point;
    point.Set3(0.2, 0.3, 0.4);
    const double expected = std::pow(point.x, exponent) * 1.9 * 0.4;
    CHECK(values.value(0, point) == Catch::Approx(expected).margin(2.0e-13));
    CHECK(std::abs(polynomial.value(0, point) - expected) > 1.0e-4);
    CHECK(polynomial.value(0, point) == Catch::Approx(nodal.GetValue(0, point)).margin(2.0e-13));
    const auto gradient = values.gradient(0, point);
    CHECK(gradient[0] == Catch::Approx(exponent * std::pow(point.x, exponent - 1.0) * 1.9 * 0.4 / 2.0).margin(2.0e-13));
    CHECK(gradient[1] == Catch::Approx(std::pow(point.x, exponent) * 0.4).margin(2.0e-13));
    CHECK(gradient[2] == Catch::Approx(-std::pow(point.x, exponent) * 1.9).margin(2.0e-13));
    serif::mapping::MappingPointContext dilation;
    dilation.inverse_mapping_jacobian.SetSize(3);
    dilation.inverse_mapping_jacobian = 0.0;
    for (int axis = 0; axis < 3; ++axis) dilation.inverse_mapping_jacobian(axis, axis) = 0.5;
    const auto physical_gradient = values.gradient(0, point, dilation);
    for (int axis = 0; axis < 3; ++axis) CHECK(physical_gradient[axis] == Catch::Approx(0.5 * gradient[axis]).margin(2.0e-13));
    CHECK(values.value(0, point, dilation) == Catch::Approx(expected).margin(2.0e-13));
    const auto& points = ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 8);
    ::mfem::Vector batch;
    values.values(0, points, batch);
    REQUIRE(batch.Size() == points.GetNPoints());
    for (int sample = 0; sample < points.GetNPoints(); ++sample) {
        const auto& at = points.IntPoint(sample);
        CHECK(batch[sample] == Catch::Approx(std::pow(at.x, exponent) * (1.0 + 3.0 * at.y) * (2.0 - 4.0 * at.z)).margin(2.0e-12));
    }
    point.x = 0.0;
    CHECK(values.value(0, point) == Catch::Approx(0.0).margin(2.0e-13));
    for (double component : values.gradient(0, point)) CHECK(component == Catch::Approx(0.0).margin(2.0e-13));
}

TEST_CASE("TOP samples the reconstructed field between its density nodes", "[serialize][top_export][mfem][reconstruction]") {
    namespace model = serif::model;
    namespace deformation = serif::geometry::deformation;
    namespace discretization = serif::discretization;
    namespace backend = discretization::backend::mfem;
    const auto star = fixture::make_model(4.0 / 3.0);
    const auto radial = seed::make_polytropic_seed(star);
    REQUIRE(radial);
    const auto finite_elements = model::make_stellar_discretization(star,
        {.radius = radial->scales.radius, .geometry_order = 6},
        deformation::DomainDeformation<deformation::SphericalHarmonicSurface>{
            .surface = {.maximum_degree = 0}, .stellar_interior = {}, .vacuum = {}},
        discretization::reconstruction::SurfacePower<fixture::Density>{.density_exponent = 4.0 / 3.0});
    const auto state = seed::backend::mfem::make_polytropic_state(star, finite_elements, *radial);
    REQUIRE(state);
    const auto& binding = finite_elements.template binding<fixture::Density>();
    const auto density = state->template get<fixture::Density>();
    ::mfem::Vector coefficients(static_cast<int>(density.size()));
    std::ranges::copy(density, coefficients.GetData());
    backend::FieldValues<fixture::Density> values(binding, coefficients);
    ::mfem::GridFunction polynomial(&binding.space());
    polynomial = binding.scatter(coefficients);
    const std::vector<double> displacement(finite_elements.deformation.reference.displacement_space.GetVSize(), 0.0);
    using Binding = std::remove_cvref_t<decltype(binding)>;
    auto temporary_binding = std::make_unique<Binding>(binding.space(), binding.reconstruction());
    serif::backend::mfem::serialize::detail::TopModelSampler sampler(*temporary_binding, density,
        finite_elements.deformation.reference, displacement, {0.0, 0.0, 1.0}, {0.0, 0.0, 0.0});
    temporary_binding.reset();
    int selected = -1;
    for (int element = 0; element < binding.mesh().GetNE(); ++element) {
        if (binding.reconstruction_element(element) != nullptr) { selected = element; break; }
    }
    REQUIRE(selected >= 0);
    ::mfem::IntegrationPoint point;
    point.Set3(0.27, 0.41, 0.69);
    ::mfem::Vector position(3);
    binding.mesh().GetElementTransformation(selected)->Transform(point, position);
    const double radius = position.Norml2();
    const double expected = values.value(selected, point);
    REQUIRE(std::abs(expected - polynomial.GetValue(selected, point)) > 1.0e-8 * radial->scales.central_density.value());
    const auto exported = sampler.density(radius, position[2] / radius, std::atan2(position[1], position[0]));
    INFO((exported ? "Sample succeeded" : exported.error().message));
    REQUIRE(exported);
    CHECK(*exported == Catch::Approx(expected).epsilon(2.0e-10));

    // Rayfem receives a display approximation sampled from the same physical field.
    // This checks the parent/submesh coordinate correspondence without opening a viewer.
    serif::visualization::backend::mfem::StellarFrame frame(finite_elements);
    StationaryDisplayCache cache{.linearization = {state->template get<fixture::Shape>()},
        .displacement = ::mfem::Vector(finite_elements.template binding<fixture::Displacement>().size())};
    cache.displacement = 0.0;
    frame.update(*state, cache);
    const auto fields = frame.fields();
    REQUIRE(fields.front().name == "density (sampled)");
    const auto& display = *fields.front().gf;
    CHECK(display.FESpace()->GetMaxElementOrder() == binding.space().GetMaxElementOrder() + 3);
    const auto& parent_elements = static_cast<const ::mfem::SubMesh&>(frame.mesh()).GetParentElementIDMap();
    bool displayed = false;
    for (int element = 0; element < parent_elements.Size(); ++element) {
        if (parent_elements[element] != selected) continue;
        displayed = true;
        const auto& nodes = display.FESpace()->GetFE(element)->GetNodes();
        for (int node = 0; node < nodes.GetNPoints(); ++node) {
            const auto& at = nodes.IntPoint(node);
            CHECK(display.GetValue(element, at) == Catch::Approx(values.value(selected, at)).margin(3.0e-12 * radial->scales.central_density.value()));
        }
    }
    REQUIRE(displayed);
}