#include <algorithm>
#include <cmath>
#include <memory>
#include <numeric>
#include <span>
#include <vector>
#include <catch2/catch_test_macros.hpp>
#include "serif/geometry/deformation/radial_extension.hpp"
#include "serif/geometry/deformation/surface_embedding.hpp"
#include "serif/tests/deformation/mesh.hpp"

namespace {
    namespace deformation = serif::geometry::deformation;
    namespace backend = deformation::backend::mfem;
    using serif::tests::deformation::PairedMesh;

    /** @brief Preserve the old physical Roche equipotential and its asymmetric companion potential. */
    struct BinaryRochePotential {
        double primary_mass;
        double companion_mass;
        double separation;

        /** @brief Evaluate gravitational and centrifugal contributions in units with G equal to one. */
        [[nodiscard]] double operator()(double x, double y, double z) const {
            const double primary_distance = std::sqrt(x * x + y * y + z * z);
            const double companion_offset = x - separation;
            const double companion_distance = std::sqrt(companion_offset * companion_offset + y * y + z * z);
            const double total_mass = primary_mass + companion_mass;
            const double center_of_mass = separation * companion_mass / total_mass;
            const double angular_speed_squared = total_mass / (separation * separation * separation);
            const double rotation_radius_squared = (x - center_of_mass) * (x - center_of_mass) + y * y;
            return -primary_mass / primary_distance - companion_mass / companion_distance - 0.5 * angular_speed_squared * rotation_radius_squared;
        }

        /** @brief Preserve the first closed-envelope radial crossing and the original root tolerances. */
        [[nodiscard]] double radial_intersection(const std::array<double, 3>& direction, double target_potential, double reference_scale) const {
            double lower_radius = 1.0e-8 * reference_scale;
            const double lower_value = (*this)(lower_radius * direction[0], lower_radius * direction[1], lower_radius * direction[2]) - target_potential;
            REQUIRE(lower_value < 0.0);
            double upper_radius = lower_radius;
            double upper_value = lower_value;
            constexpr int bracket_samples = 512;
            for (int sample = 1; sample <= bracket_samples && upper_value <= 0.0; ++sample) {
                upper_radius = 0.45 * separation * static_cast<double>(sample) / bracket_samples;
                upper_value = (*this)(upper_radius * direction[0], upper_radius * direction[1], upper_radius * direction[2]) - target_potential;
            }
            REQUIRE(upper_value > 0.0);
            for (int iteration = 0; iteration < 80; ++iteration) {
                const double middle_radius = 0.5 * (lower_radius + upper_radius);
                const double middle_value = (*this)(middle_radius * direction[0], middle_radius * direction[1], middle_radius * direction[2]) - target_potential;
                if (middle_value > 0.0) {
                    upper_radius = middle_radius;
                } else {
                    lower_radius = middle_radius;
                }
            }
            return 0.5 * (lower_radius + upper_radius);
        }
    };

    /**
     * @brief Sample the actual finite-element trace using its moved tangents for physical area weights.
     * Only trace basis functions affect this restriction; no ideal continuous radial surface is substituted.
     * The independent boundary calculation avoids requiring an interior volume extension to test a surface representation.
     */
    template <typename PointError>
    [[nodiscard]] std::array<double, 2> sample_surface_error(PairedMesh& mesh, const deformation::ReferenceSurface& reference, std::span<const double> displacement, const PointError& point_error) {
        std::vector<int> surface_indices(mesh.scalar_space.GetNDofs(), -1);
        for (std::size_t node = 0; node < reference.scalar_dofs.size(); ++node) {
            surface_indices[reference.scalar_dofs[node]] = static_cast<int>(node);
        }
        double maximum_error = 0.0, squared_error = 0.0, surface_area = 0.0;
        int samples = 0;
        ::mfem::Array<int> boundary_dofs;
        ::mfem::Vector shape, position(3);
        ::mfem::DenseMatrix derivatives;
        for (int boundary = 0; boundary < mesh.generated.mesh->GetNBE(); ++boundary) {
            if (mesh.generated.mesh->GetBdrAttribute(boundary) != PairedMesh::surface_attribute) {
                continue;
            }
            ::mfem::ElementTransformation& transformation = *mesh.generated.mesh->GetBdrElementTransformation(boundary);
            const ::mfem::FiniteElement& element = *mesh.scalar_space.GetBE(boundary);
            const ::mfem::IntegrationRule& rule = ::mfem::IntRules.Get(transformation.GetGeometryType(), 2 * element.GetOrder());
            mesh.scalar_space.GetBdrElementDofs(boundary, boundary_dofs);
            shape.SetSize(boundary_dofs.Size());
            derivatives.SetSize(boundary_dofs.Size(), 2);
            for (int point = 0; point < rule.GetNPoints(); ++point) {
                const ::mfem::IntegrationPoint& integration_point = rule.IntPoint(point);
                transformation.Transform(integration_point, position);
                transformation.SetIntPoint(&integration_point);
                ::mfem::DenseMatrix tangents(transformation.Jacobian());
                element.CalcShape(integration_point, shape);
                element.CalcDShape(integration_point, derivatives);
                for (int node = 0; node < boundary_dofs.Size(); ++node) {
                    const int surface_node = surface_indices[::mfem::FiniteElementSpace::DecodeDof(boundary_dofs[node])];
                    REQUIRE(surface_node >= 0);
                    for (int component = 0; component < 3; ++component) {
                        const double value = displacement[3 * surface_node + component];
                        position[component] += shape[node] * value;
                        for (int tangent = 0; tangent < 2; ++tangent) {
                            tangents(component, tangent) += derivatives(node, tangent) * value;
                        }
                    }
                }
                double area_squared = 0.0;
                for (int component = 0; component < 3; ++component) {
                    const int second = (component + 1) % 3;
                    const int third = (component + 2) % 3;
                    const double normal = tangents(second, 0) * tangents(third, 1) - tangents(third, 0) * tangents(second, 1);
                    area_squared += normal * normal;
                }
                const double weight = integration_point.weight * std::sqrt(area_squared);
                const double error = point_error(position);
                REQUIRE(std::isfinite(error));
                REQUIRE(weight > 0.0);
                maximum_error = std::max(maximum_error, error);
                squared_error += error * error * weight;
                surface_area += weight;
                ++samples;
            }
        }
        REQUIRE(samples > 0);
        REQUIRE(surface_area > 0.0);
        return {maximum_error, std::sqrt(squared_error / surface_area)};
    }
}

TEST_CASE("Nodal radial displacement resolves oblate and binary Roche envelopes between surface nodes", "[unit][deformation]") {
    // MeanField's baseline sandbox mesh uses two refinements, order-four geometry and order-three displacement.
    PairedMesh mesh(2, 0.0, 3, 4, 5.0);
    const auto reference = backend::sample_reference_surface(mesh.scalar_space, PairedMesh::surface_attribute);
    REQUIRE(reference.has_value());
    const auto surface = deformation::build_surface_embedding(deformation::NodalRadialSurface{}, reference->positions);
    REQUIRE(surface.has_value());
    const double reference_scale = std::accumulate(surface->reference_radii().begin(), surface->reference_radii().end(), 0.0) / surface->surface_node_count();
    const double equatorial_radius = 1.08 * reference_scale;
    const double polar_radius = 0.94 * reference_scale;
    std::vector<double> oblate_parameters(surface->input_size()), roche_parameters(surface->input_size());
    const BinaryRochePotential potential{.primary_mass = 1.0, .companion_mass = 0.7, .separation = 4.0 * reference_scale};
    const double target_potential = potential(0.0, 0.0, reference_scale);
    for (std::size_t node = 0; node < surface->surface_node_count(); ++node) {
        const std::array<double, 3>& direction = surface->radial_directions()[node];
        const double oblate_radius = 1.0 / std::sqrt((direction[0] * direction[0] + direction[1] * direction[1]) / (equatorial_radius * equatorial_radius) + direction[2] * direction[2] / (polar_radius * polar_radius));
        oblate_parameters[node] = oblate_radius - surface->reference_radii()[node];
        roche_parameters[node] = potential.radial_intersection(direction, target_potential, reference_scale) - surface->reference_radii()[node];
    }
    std::vector<double> oblate_displacement(surface->output_size()), roche_displacement(surface->output_size());
    REQUIRE(surface->apply(oblate_parameters, oblate_displacement));
    REQUIRE(surface->apply(roche_parameters, roche_displacement));
    for (std::size_t node = 0; node < surface->surface_node_count(); ++node) {
        std::array<double, 3> oblate_position{}, roche_position{};
        for (std::size_t component = 0; component < 3; ++component) {
            oblate_position[component] = reference->positions[node][component] + oblate_displacement[3 * node + component];
            roche_position[component] = reference->positions[node][component] + roche_displacement[3 * node + component];
        }
        const double level = (oblate_position[0] * oblate_position[0] + oblate_position[1] * oblate_position[1]) / (equatorial_radius * equatorial_radius) + oblate_position[2] * oblate_position[2] / (polar_radius * polar_radius);
        CHECK(std::abs(level - 1.0) <= 3.0e-12);
        CHECK(std::abs(potential(roche_position[0], roche_position[1], roche_position[2]) - target_potential) <= 2.0e-11 * std::max(1.0, std::abs(target_potential)));
    }
    const auto oblate_error = sample_surface_error(mesh, *reference, oblate_displacement, [equatorial_radius, polar_radius](const ::mfem::Vector& position) {
        return std::abs((position[0] * position[0] + position[1] * position[1]) / (equatorial_radius * equatorial_radius) + position[2] * position[2] / (polar_radius * polar_radius) - 1.0);
    });
    const auto roche_error = sample_surface_error(mesh, *reference, roche_displacement, [&potential, target_potential](const ::mfem::Vector& position) {
        return std::abs(potential(position[0], position[1], position[2]) - target_potential) / std::max(1.0, std::abs(target_potential));
    });
    CAPTURE(oblate_error, roche_error);
    CHECK(oblate_error[0] < 5.0e-3);
    CHECK(oblate_error[1] < 1.0e-3);
    CHECK(roche_error[0] < 2.0e-2);
    CHECK(roche_error[1] < 5.0e-3);
    const double companion_facing = potential.radial_intersection({1.0, 0.0, 0.0}, target_potential, reference_scale);
    const double companion_opposing = potential.radial_intersection({-1.0, 0.0, 0.0}, target_potential, reference_scale);
    CHECK(std::abs(companion_facing - companion_opposing) > 1.0e-3 * reference_scale);
}