#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

#include <algorithm>
#include <cmath>
#include <memory>
#include <cassert>

#include <boost/math/quadrature/gauss_kronrod.hpp>
#include <mfem.hpp>
#include <mpi.h>

import mean_field;
import test_helpers;

using namespace mean_field;

namespace {
    struct GravitationalEnergies {
        double binding;
        double virial;
    };

    struct HomogeneousEllipsoidAnalytic {
        double coefficient_x;
        double coefficient_y;
        double coefficient_z;
        double energy_kernel;
    };

    template <typename GravitySolutionType>
    GravitationalEnergies compute_gravitational_energies(fem::FEM& f, const mfem::GridFunction& rho,
                                                         const GravitySolutionType& gravity_solution,
                                                         const int quadrature_order) {
        const int dim = f.mesh->Dimension();
        double local_bind_integral = 0.0;
        double local_virial_integral = 0.0;

        mfem::Vector x_physical(dim);
        mfem::Vector grad_phi_element(dim);
        mfem::Vector grad_phi_physical(dim);
        mfem::DenseMatrix map_jacobian(dim, dim);

        for (int elem_id = 0; elem_id < f.mesh->GetNE(); ++elem_id) {
            if (f.mesh->GetAttribute(elem_id) == 3) {
                continue;
            }

            mfem::ElementTransformation* transformation = f.mesh->GetElementTransformation(elem_id);
            const mfem::IntegrationRule& integration_rule =
                mfem::IntRules.Get(transformation->GetGeometryType(), quadrature_order);

            for (int q = 0; q < integration_rule.GetNPoints(); ++q) {
                const mfem::IntegrationPoint& integration_point = integration_rule.IntPoint(q);
                transformation->SetIntPoint(&integration_point);

                double weight = transformation->Weight() * integration_point.weight;

                if (f.has_mapping()) {
                    const double map_determinant = f.mapping->ComputeDetJ(*transformation, integration_point);
                    MFEM_VERIFY(map_determinant > 0.0, "Domain mapping has a non-positive Jacobian determinant.");

                    weight *= map_determinant;
                    f.mapping->GetPhysicalPoint(*transformation, integration_point, x_physical);
                    gravity_solution.gradPhi.GetVectorValue(elem_id, integration_point, grad_phi_element);
                    f.mapping->ComputeJacobian(*transformation, map_jacobian);
                    map_jacobian.Mult(grad_phi_element, grad_phi_physical);
                    grad_phi_physical /= map_determinant;
                } else {
                    transformation->Transform(integration_point, x_physical);
                    gravity_solution.gradPhi.GetVectorValue(elem_id, integration_point, grad_phi_physical);
                }

                const double rho_value = rho.GetValue(elem_id, integration_point);
                const double phi_value = gravity_solution.phi.GetValue(elem_id, integration_point);
                double radius_dot_gradient = 0.0;

                for (int d = 0; d < dim; ++d) {
                    radius_dot_gradient += (x_physical(d) - f.com(d)) * grad_phi_physical(d);
                }

                local_bind_integral += rho_value * phi_value * weight;
                local_virial_integral += rho_value * radius_dot_gradient * weight;
            }
        }

        const double local_w_bind = 0.5 * local_bind_integral;
        const double local_w_vir = -local_virial_integral;
        double global_w_bind = 0.0;
        double global_w_vir = 0.0;
        MPI_Comm communicator = f.L2_fes->GetComm();

        MPI_Allreduce(&local_w_bind, &global_w_bind, 1, MPI_DOUBLE, MPI_SUM, communicator);
        MPI_Allreduce(&local_w_vir, &global_w_vir, 1, MPI_DOUBLE, MPI_SUM, communicator);

        return {.binding = global_w_bind, .virial = global_w_vir};
    }

    void zero_vacuum_density(const fem::FEM& f, mfem::GridFunction& rho) {
        for (int i = 0; i < f.vacuum_tdof_rho.Size(); ++i) {
            rho(f.vacuum_tdof_rho[i]) = 0.0;
        }
    }

    int get_gravity_quadrature_order(const fem::FEM& f) {
        return 2 * std::max(f.L2_fes->GetMaxElementOrder(), f.RT_fes->GetMaxElementOrder()) + 8;
    }

    double compute_ellipsoid_coefficient(const double normalized_axis_x, const double normalized_axis_y,
                                         const double normalized_axis_z, const double target_axis_squared) {
        auto integrand = [=](const double t) {
            if (t <= 0.0 || t >= 1.0) {
                return 0.0;
            }

            const double one_minus_t = 1.0 - t;
            const double s = t / one_minus_t;
            const double s_squared = s * s;
            const double ds_squared_dt = 2.0 * s / (one_minus_t * one_minus_t);
            const double delta = std::sqrt(
                (normalized_axis_x * normalized_axis_x + s_squared) *
                (normalized_axis_y * normalized_axis_y + s_squared) *
                (normalized_axis_z * normalized_axis_z + s_squared)
            );

            return normalized_axis_x * normalized_axis_y * normalized_axis_z * ds_squared_dt /
                   ((target_axis_squared + s_squared) * delta);
        };

        double integration_error = 0.0;
        return boost::math::quadrature::gauss_kronrod<double, 61>::integrate(
            integrand, 0.0, 1.0, 15, 1.0e-13, &integration_error
        );
    }

    double compute_ellipsoid_energy_kernel(const double normalized_axis_x, const double normalized_axis_y,
                                           const double normalized_axis_z, const double length_scale) {
        auto integrand = [=](const double t) {
            if (t <= 0.0) {
                return 0.0;
            }
            if (t >= 1.0) {
                return 2.0;
            }

            const double one_minus_t = 1.0 - t;
            const double s = t / one_minus_t;
            const double s_squared = s * s;
            const double ds_squared_dt = 2.0 * s / (one_minus_t * one_minus_t);
            const double delta = std::sqrt(
                (normalized_axis_x * normalized_axis_x + s_squared) *
                (normalized_axis_y * normalized_axis_y + s_squared) *
                (normalized_axis_z * normalized_axis_z + s_squared)
            );

            return ds_squared_dt / delta;
        };

        double integration_error = 0.0;
        const double dimensionless_integral = boost::math::quadrature::gauss_kronrod<double, 61>::integrate(
            integrand, 0.0, 1.0, 15, 1.0e-13, &integration_error
        );

        return dimensionless_integral / length_scale;
    }

    HomogeneousEllipsoidAnalytic compute_homogeneous_ellipsoid_analytic(
        const double semi_axis_x, const double semi_axis_y, const double semi_axis_z
    ) {
        const double length_scale = std::cbrt(semi_axis_x * semi_axis_y * semi_axis_z);
        const double normalized_axis_x = semi_axis_x / length_scale;
        const double normalized_axis_y = semi_axis_y / length_scale;
        const double normalized_axis_z = semi_axis_z / length_scale;
        const double coefficient_x = compute_ellipsoid_coefficient(
            normalized_axis_x, normalized_axis_y, normalized_axis_z, normalized_axis_x * normalized_axis_x
        );
        const double coefficient_y = compute_ellipsoid_coefficient(
            normalized_axis_x, normalized_axis_y, normalized_axis_z, normalized_axis_y * normalized_axis_y
        );
        const double coefficient_z = compute_ellipsoid_coefficient(
            normalized_axis_x, normalized_axis_y, normalized_axis_z, normalized_axis_z * normalized_axis_z
        );
        const double energy_kernel = compute_ellipsoid_energy_kernel(
            normalized_axis_x, normalized_axis_y, normalized_axis_z, length_scale
        );

        return {
            .coefficient_x = coefficient_x,
            .coefficient_y = coefficient_y,
            .coefficient_z = coefficient_z,
            .energy_kernel = energy_kernel
        };
    }
}

TEST_CASE("Uniform Potential Matches Analytic", tags::gravity & tags::analytic_comparison) {
    auto args = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);
    f.mapping->ResetDisplacement();
    physics::update_stiffness_matrix(f);

    const double radius = utils::RADIUS;
    const double mass = utils::MASS;
    const double analytic_volume = (4.0 / 3.0) * M_PI * std::pow(radius, 3.0);
    const double density = mass / analytic_volume;

    mfem::GridFunction rho_uniform(f.L2_fes.get());
    rho_uniform = density;
    zero_vacuum_density(f, rho_uniform);
    analysis::conserve_mass(f, rho_uniform, mass);

    f.com = analysis::get_com(f, rho_uniform);
    f.Q = physics::compute_quadrupole_moment_tensor(f, rho_uniform, f.com);

    const auto gravity_solution = physics::grav_potential(f, args, rho_uniform);
    constexpr double potential_tolerance = utils::APPROX_MAX_ACCEPTABLE_POTENTIAL_ERROR_SI_BURNING;
    double local_max_abs_error = 0.0;
    double local_max_rel_error = 0.0;

    const int num_elements_to_test = std::min(30, f.mesh->GetNE());
    for (int elem_id = 0; elem_id < num_elements_to_test; ++elem_id) {
        mfem::ElementTransformation* transformation = f.mesh->GetElementTransformation(elem_id);
        const mfem::IntegrationRule& integration_rule = mfem::IntRules.Get(transformation->GetGeometryType(), 2);
        const mfem::IntegrationPoint& integration_point = integration_rule.IntPoint(0);
        transformation->SetIntPoint(&integration_point);

        mfem::Vector x_physical;
        f.mapping->GetPhysicalPoint(*transformation, integration_point, x_physical);

        const double radial_coordinate = x_physical.Norml2();
        if (radial_coordinate < 1.0e-9) {
            continue;
        }

        const double phi_analytic = -(utils::G * mass / (2.0 * std::pow(radius, 3.0))) *
                                    (3.0 * radius * radius - radial_coordinate * radial_coordinate);
        const double phi_fem = gravity_solution.phi.GetValue(elem_id, integration_point);
        const double absolute_error = std::abs(phi_fem - phi_analytic);
        const double relative_error = absolute_error / std::abs(phi_analytic);

        local_max_abs_error = std::max(local_max_abs_error, absolute_error);
        local_max_rel_error = std::max(local_max_rel_error, relative_error);
        CHECK_THAT(relative_error, Catch::Matchers::WithinAbs(0.0, 0.1 * potential_tolerance));
    }

    double global_max_abs_error = 0.0;
    double global_max_rel_error = 0.0;
    MPI_Comm communicator = f.L2_fes->GetComm();
    MPI_Allreduce(&local_max_abs_error, &global_max_abs_error, 1, MPI_DOUBLE, MPI_MAX, communicator);
    MPI_Allreduce(&local_max_rel_error, &global_max_rel_error, 1, MPI_DOUBLE, MPI_MAX, communicator);

    const int quadrature_order = get_gravity_quadrature_order(f);
    const GravitationalEnergies energies = compute_gravitational_energies(f, rho_uniform, gravity_solution, quadrature_order);
    const double analytic_binding_energy = -(3.0 / 5.0) * utils::G * mass * mass / radius;
    const double relative_binding_error = std::abs(energies.binding - analytic_binding_energy) / std::abs(analytic_binding_energy);
    const double relative_virial_error = std::abs(energies.virial - analytic_binding_energy) / std::abs(analytic_binding_energy);
    const double relative_consistency_error = std::abs(energies.binding - energies.virial) / std::abs(energies.binding);

    INFO("Analytic binding energy = " << analytic_binding_energy);
    INFO("Computed binding energy = " << energies.binding);
    INFO("Computed virial energy = " << energies.virial);
    INFO("Relative virial consistency error = " << relative_consistency_error);

    constexpr double energy_tolerance = 1.0e-5;
    constexpr double consistency_tolerance = 1.0e-6;

    CHECK_THAT(global_max_rel_error, Catch::Matchers::WithinAbs(0.0, 0.1 * potential_tolerance));
    CHECK_THAT(global_max_abs_error, Catch::Matchers::WithinAbs(0.0, 0.1 * potential_tolerance));
    CHECK_THAT(relative_binding_error, Catch::Matchers::WithinAbs(0.0, energy_tolerance));
    CHECK_THAT(relative_virial_error, Catch::Matchers::WithinAbs(0.0, energy_tolerance));
    CHECK_THAT(relative_consistency_error, Catch::Matchers::WithinAbs(0.0, consistency_tolerance));
}

TEST_CASE("Parabolic Density Virial Self-Consistency", tags::gravity & tags::analytic_comparison) {
    auto args = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);
    f.mapping->ResetDisplacement();
    physics::update_stiffness_matrix(f);

    const double radius = utils::RADIUS;
    const double mass = utils::MASS;
    const double central_density = (15.0 * mass) / (8.0 * M_PI * std::pow(radius, 3.0));

    auto parabolic_rho = [central_density, radius](const mfem::Vector& x) {
        const double radial_coordinate = x.Norml2();
        return central_density * (1.0 - radial_coordinate * radial_coordinate / (radius * radius));
    };

    std::unique_ptr<mfem::Coefficient> rho_coeff;
    if (f.has_mapping()) {
        rho_coeff = std::make_unique<mapping::PhysicalPositionFunctionCoefficient>(*f.mapping, parabolic_rho);
    } else {
        rho_coeff = std::make_unique<mfem::FunctionCoefficient>(parabolic_rho);
    }

    mfem::GridFunction rho_grid(f.L2_fes.get());
    rho_grid.ProjectCoefficient(*rho_coeff);
    zero_vacuum_density(f, rho_grid);
    analysis::conserve_mass(f, rho_grid, mass);

    f.com = analysis::get_com(f, rho_grid);
    f.Q = physics::compute_quadrupole_moment_tensor(f, rho_grid, f.com);

    const auto gravity_solution = physics::grav_potential(f, args, rho_grid);
    const int quadrature_order = get_gravity_quadrature_order(f);
    const GravitationalEnergies energies = compute_gravitational_energies(f, rho_grid, gravity_solution, quadrature_order);
    const double analytic_binding_energy = -(5.0 / 7.0) * utils::G * mass * mass / radius;
    const double relative_binding_error = std::abs(energies.binding - analytic_binding_energy) / std::abs(analytic_binding_energy);
    const double relative_virial_error = std::abs(energies.virial - analytic_binding_energy) / std::abs(analytic_binding_energy);
    const double relative_consistency_error = std::abs(energies.binding - energies.virial) / std::abs(energies.binding);

    INFO("Analytic binding energy = " << analytic_binding_energy);
    INFO("Computed binding energy = " << energies.binding);
    INFO("Computed virial energy = " << energies.virial);
    INFO("Relative virial consistency error = " << relative_consistency_error);

    constexpr double analytic_tolerance = 1.0e-5;
    constexpr double consistency_tolerance = 1.0e-6;

    CHECK_THAT(relative_binding_error, Catch::Matchers::WithinAbs(0.0, analytic_tolerance));
    CHECK_THAT(relative_virial_error, Catch::Matchers::WithinAbs(0.0, analytic_tolerance));
    CHECK_THAT(relative_consistency_error, Catch::Matchers::WithinAbs(0.0, consistency_tolerance));
}

TEST_CASE("Rational Density Virial Self-Consistency", tags::gravity & tags::self_consistency) {
    auto args = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);
    f.mapping->ResetDisplacement();
    physics::update_stiffness_matrix(f);

    const double radius = utils::RADIUS;
    const double mass = utils::MASS;

    // Larger values are more centrally concentrated and generally harder for a polynomial to represent.
    // A regression would be considered if this test does not pass for concentrations <= 16.0.
    constexpr double concentration = 16.0;
    const double density_scale = mass / std::pow(radius, 3.0);

    auto rational_rho = [radius, density_scale](const mfem::Vector& x) {
        const double normalized_radius_squared = (x * x) / (radius * radius);
        if (normalized_radius_squared >= 1.0) {
            return 0.0;
        }

        const double denominator = 1.0 + concentration * normalized_radius_squared;
        return density_scale * (1.0 - normalized_radius_squared) / (denominator * denominator);
    };

    std::unique_ptr<mfem::Coefficient> rho_coeff;
    if (f.has_mapping()) {
        rho_coeff = std::make_unique<mapping::PhysicalPositionFunctionCoefficient>(*f.mapping, rational_rho);
    } else {
        rho_coeff = std::make_unique<mfem::FunctionCoefficient>(rational_rho);
    }

    mfem::GridFunction rho_grid(f.L2_fes.get());
    rho_grid.ProjectCoefficient(*rho_coeff);
    zero_vacuum_density(f, rho_grid);
    analysis::conserve_mass(f, rho_grid, mass);

    f.com = analysis::get_com(f, rho_grid);
    f.Q = physics::compute_quadrupole_moment_tensor(f, rho_grid, f.com);

    const auto gravity_solution = physics::grav_potential(f, args, rho_grid);
    const int quadrature_order = get_gravity_quadrature_order(f);
    const GravitationalEnergies energies = compute_gravitational_energies(f, rho_grid, gravity_solution, quadrature_order);

    REQUIRE(energies.binding < 0.0);
    REQUIRE(energies.virial < 0.0);

    const double relative_consistency_error = std::abs(energies.binding - energies.virial) / std::abs(energies.binding);
    INFO("W_bind = " << energies.binding);
    INFO("W_vir = " << energies.virial);
    INFO("Relative virial consistency error = " << relative_consistency_error);

    constexpr double virial_tolerance = 1.0e-5;
    CHECK_THAT(relative_consistency_error, Catch::Matchers::WithinAbs(0.0, virial_tolerance));
}

TEST_CASE("Homogeneous Ellipsoid Analytic Gravity", tags::gravity & tags::analytic_comparison) {
    auto args = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    const double radius = utils::RADIUS;
    const double mass = utils::MASS;
    constexpr double x_scale = 1.15;
    constexpr double y_scale = 0.95;
    constexpr double z_scale = 1.0 / (x_scale * y_scale);
    assert(std::abs(x_scale * y_scale * z_scale - 1.0) < 1.0e-14);

    const double semi_axis_x = x_scale * radius;
    const double semi_axis_y = y_scale * radius;
    const double semi_axis_z = z_scale * radius;

    auto affine_displacement = [](const mfem::Vector& x, mfem::Vector& displacement_value) {
        displacement_value.SetSize(3);
        displacement_value(0) = (x_scale - 1.0) * x(0);
        displacement_value(1) = (y_scale - 1.0) * x(1);
        displacement_value(2) = (z_scale - 1.0) * x(2);
    };

    mfem::VectorFunctionCoefficient displacement_coeff(3, affine_displacement);
    mfem::ParGridFunction displacement(f.Vec_H1_fes.get());
    displacement.ProjectCoefficient(displacement_coeff);
    f.mapping->SetDisplacement(displacement);
    physics::update_stiffness_matrix(f);

    const double analytic_volume = (4.0 / 3.0) * M_PI * semi_axis_x * semi_axis_y * semi_axis_z;
    const double density = mass / analytic_volume;

    mfem::GridFunction rho_grid(f.L2_fes.get());
    rho_grid = density;
    zero_vacuum_density(f, rho_grid);

    const double projected_mass = analysis::domain_integrate_grid_function(f, rho_grid, utils::DOMAINS::STELLAR);
    const double numerical_density = density * mass / projected_mass;
    analysis::conserve_mass(f, rho_grid, mass);

    f.com = analysis::get_com(f, rho_grid);
    f.Q = physics::compute_quadrupole_moment_tensor(f, rho_grid, f.com);

    const HomogeneousEllipsoidAnalytic analytic =
        compute_homogeneous_ellipsoid_analytic(semi_axis_x, semi_axis_y, semi_axis_z);
    const double coefficient_sum = analytic.coefficient_x + analytic.coefficient_y + analytic.coefficient_z;

    INFO("A_x = " << analytic.coefficient_x);
    INFO("A_y = " << analytic.coefficient_y);
    INFO("A_z = " << analytic.coefficient_z);
    INFO("A_x + A_y + A_z = " << coefficient_sum);
    REQUIRE_THAT(coefficient_sum, Catch::Matchers::WithinAbs(2.0, 1.0e-11));

    mfem::DenseMatrix analytic_quadrupole(3, 3);
    analytic_quadrupole = 0.0;
    analytic_quadrupole(0, 0) = (mass / 5.0) * (2.0 * semi_axis_x * semi_axis_x - semi_axis_y * semi_axis_y - semi_axis_z * semi_axis_z);
    analytic_quadrupole(1, 1) = (mass / 5.0) * (2.0 * semi_axis_y * semi_axis_y - semi_axis_x * semi_axis_x - semi_axis_z * semi_axis_z);
    analytic_quadrupole(2, 2) = (mass / 5.0) * (2.0 * semi_axis_z * semi_axis_z - semi_axis_x * semi_axis_x - semi_axis_y * semi_axis_y);

    mfem::DenseMatrix quadrupole_difference(f.Q);
    quadrupole_difference -= analytic_quadrupole;
    const double relative_quadrupole_error = quadrupole_difference.FNorm() / analytic_quadrupole.FNorm();
    INFO("Relative quadrupole error = " << relative_quadrupole_error);

    const auto gravity_solution = physics::grav_potential(f, args, rho_grid);
    const int quadrature_order = get_gravity_quadrature_order(f);
    double local_field_error_squared = 0.0;
    double local_field_norm_squared = 0.0;

    mfem::Vector x_physical(3);
    mfem::Vector grad_phi_element(3);
    mfem::Vector grad_phi_physical(3);
    mfem::Vector grad_phi_analytic(3);
    mfem::Vector grad_phi_difference(3);
    mfem::DenseMatrix map_jacobian(3, 3);

    for (int elem_id = 0; elem_id < f.mesh->GetNE(); ++elem_id) {
        if (f.mesh->GetAttribute(elem_id) == 3) {
            continue;
        }

        mfem::ElementTransformation* transformation = f.mesh->GetElementTransformation(elem_id);
        const mfem::IntegrationRule& integration_rule =
            mfem::IntRules.Get(transformation->GetGeometryType(), quadrature_order);

        for (int q = 0; q < integration_rule.GetNPoints(); ++q) {
            const mfem::IntegrationPoint& integration_point = integration_rule.IntPoint(q);
            transformation->SetIntPoint(&integration_point);

            const double map_determinant = f.mapping->ComputeDetJ(*transformation, integration_point);
            MFEM_VERIFY(map_determinant > 0.0, "Domain mapping has a non-positive Jacobian determinant.");

            const double weight = transformation->Weight() * integration_point.weight * map_determinant;
            f.mapping->GetPhysicalPoint(*transformation, integration_point, x_physical);
            gravity_solution.gradPhi.GetVectorValue(elem_id, integration_point, grad_phi_element);
            f.mapping->ComputeJacobian(*transformation, map_jacobian);
            map_jacobian.Mult(grad_phi_element, grad_phi_physical);
            grad_phi_physical /= map_determinant;

            grad_phi_analytic(0) = 2.0 * M_PI * utils::G * numerical_density * analytic.coefficient_x * x_physical(0);
            grad_phi_analytic(1) = 2.0 * M_PI * utils::G * numerical_density * analytic.coefficient_y * x_physical(1);
            grad_phi_analytic(2) = 2.0 * M_PI * utils::G * numerical_density * analytic.coefficient_z * x_physical(2);

            grad_phi_difference = grad_phi_physical;
            grad_phi_difference -= grad_phi_analytic;
            local_field_error_squared += (grad_phi_difference * grad_phi_difference) * weight;
            local_field_norm_squared += (grad_phi_analytic * grad_phi_analytic) * weight;
        }
    }

    double global_field_error_squared = 0.0;
    double global_field_norm_squared = 0.0;
    MPI_Comm communicator = f.L2_fes->GetComm();
    MPI_Allreduce(&local_field_error_squared, &global_field_error_squared, 1, MPI_DOUBLE, MPI_SUM, communicator);
    MPI_Allreduce(&local_field_norm_squared, &global_field_norm_squared, 1, MPI_DOUBLE, MPI_SUM, communicator);

    const double relative_field_error = std::sqrt(global_field_error_squared / global_field_norm_squared);
    const GravitationalEnergies energies = compute_gravitational_energies(f, rho_grid, gravity_solution, quadrature_order);
    const double analytic_binding_energy = -(3.0 / 10.0) * utils::G * mass * mass * analytic.energy_kernel;
    const double relative_binding_energy_error = std::abs(energies.binding - analytic_binding_energy) / std::abs(analytic_binding_energy);
    const double relative_virial_energy_error = std::abs(energies.virial - analytic_binding_energy) / std::abs(analytic_binding_energy);
    const double relative_consistency_error = std::abs(energies.binding - energies.virial) / std::abs(energies.binding);

    INFO("Analytic binding energy = " << analytic_binding_energy);
    INFO("Computed binding energy = " << energies.binding);
    INFO("Computed virial energy = " << energies.virial);
    INFO("Relative field L2 error = " << relative_field_error);
    INFO("Relative binding energy error = " << relative_binding_energy_error);
    INFO("Relative virial energy error = " << relative_virial_energy_error);
    INFO("Relative virial consistency error = " << relative_consistency_error);

    constexpr double quadrupole_tolerance = 2.0e-4;
    constexpr double field_tolerance = 1.0e-5;
    constexpr double energy_tolerance = 1.0e-5;
    constexpr double consistency_tolerance = 1.0e-5;

    CHECK_THAT(relative_quadrupole_error, Catch::Matchers::WithinAbs(0.0, quadrupole_tolerance));
    CHECK_THAT(relative_field_error, Catch::Matchers::WithinAbs(0.0, field_tolerance));
    CHECK_THAT(relative_binding_energy_error, Catch::Matchers::WithinAbs(0.0, energy_tolerance));
    CHECK_THAT(relative_virial_energy_error, Catch::Matchers::WithinAbs(0.0, energy_tolerance));
    CHECK_THAT(relative_consistency_error, Catch::Matchers::WithinAbs(0.0, consistency_tolerance));
}

TEST_CASE("Deformed Rational Density Virial Self-Consistency", tags::gravity & tags::self_consistency) {
    auto args = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    const double radius = utils::RADIUS;
    const double mass = utils::MASS;
    constexpr double x_scale = 1.15;
    constexpr double y_scale = 0.95;
    constexpr double z_scale = 1.0 / (x_scale * y_scale);
    assert(std::abs(x_scale * y_scale * z_scale - 1.0) < 1.0e-14);

    const double semi_axis_x = x_scale * radius;
    const double semi_axis_y = y_scale * radius;
    const double semi_axis_z = z_scale * radius;

    auto affine_displacement = [](const mfem::Vector& x, mfem::Vector& displacement_value) {
        displacement_value.SetSize(3);
        displacement_value(0) = (x_scale - 1.0) * x(0);
        displacement_value(1) = (y_scale - 1.0) * x(1);
        displacement_value(2) = (z_scale - 1.0) * x(2);
    };

    mfem::VectorFunctionCoefficient displacement_coeff(3, affine_displacement);
    mfem::ParGridFunction displacement(f.Vec_H1_fes.get());
    displacement.ProjectCoefficient(displacement_coeff);
    f.mapping->SetDisplacement(displacement);
    physics::update_stiffness_matrix(f);

    constexpr double concentration = 16.0;
    const double density_scale = mass / std::pow(radius, 3.0);

    auto ellipsoidal_rho = [semi_axis_x, semi_axis_y, semi_axis_z, density_scale](const mfem::Vector& x) {
        const double ellipsoidal_radius_squared =
            x(0) * x(0) / (semi_axis_x * semi_axis_x) +
            x(1) * x(1) / (semi_axis_y * semi_axis_y) +
            x(2) * x(2) / (semi_axis_z * semi_axis_z);

        if (ellipsoidal_radius_squared >= 1.0) {
            return 0.0;
        }

        const double denominator = 1.0 + concentration * ellipsoidal_radius_squared;
        return density_scale * (1.0 - ellipsoidal_radius_squared) / (denominator * denominator);
    };

    std::unique_ptr<mfem::Coefficient> rho_coeff;
    if (f.has_mapping()) {
        rho_coeff = std::make_unique<mapping::PhysicalPositionFunctionCoefficient>(*f.mapping, ellipsoidal_rho);
    } else {
        rho_coeff = std::make_unique<mfem::FunctionCoefficient>(ellipsoidal_rho);
    }

    mfem::GridFunction rho_grid(f.L2_fes.get());
    rho_grid.ProjectCoefficient(*rho_coeff);
    zero_vacuum_density(f, rho_grid);
    analysis::conserve_mass(f, rho_grid, mass);

    f.com = analysis::get_com(f, rho_grid);
    f.Q = physics::compute_quadrupole_moment_tensor(f, rho_grid, f.com);

    const double normalized_quadrupole = f.Q.FNorm() / (mass * radius * radius);
    INFO("Normalized quadrupole = " << normalized_quadrupole);
    REQUIRE(normalized_quadrupole > 1.0e-3);

    const auto gravity_solution = physics::grav_potential(f, args, rho_grid);
    const int quadrature_order = get_gravity_quadrature_order(f);
    const GravitationalEnergies energies = compute_gravitational_energies(f, rho_grid, gravity_solution, quadrature_order);

    REQUIRE(energies.binding < 0.0);
    REQUIRE(energies.virial < 0.0);

    const double relative_consistency_error = std::abs(energies.binding - energies.virial) / std::abs(energies.binding);
    INFO("W_bind = " << energies.binding);
    INFO("W_vir = " << energies.virial);
    INFO("Relative virial consistency error = " << relative_consistency_error);

    constexpr double virial_tolerance = 1.0e-5;
    CHECK_THAT(relative_consistency_error, Catch::Matchers::WithinAbs(0.0, virial_tolerance));
}