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

#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <sstream>

#include <mfem.hpp>
#include <mpi.h>

import mean_field;
import test_helpers;

namespace {
constexpr std::array<double, 6> shell_boundaries{0.0,  0.25, 0.50,
                                                 0.75, 0.90, 1.0};
constexpr int shell_count = static_cast<int>(shell_boundaries.size()) - 1;

struct ShellAccumulator {
  long long points{0};
  double weight{0.0};
  double minimum_radius{std::numeric_limits<double>::infinity()};
  double maximum_radius{0.0};
  double potential_error_squared{0.0};
  double radial_error_squared{0.0};
  double tangential_squared{0.0};
};

struct ShellMetrics {
  long long points{0};
  double minimum_radius{0.0};
  double maximum_radius{0.0};
  double potential_rms_error{0.0};
  double radial_rms_error{0.0};
  double tangential_rms{0.0};
};

struct ShellMeasurement {
  std::array<ShellMetrics, shell_count> shells{};
  long long invalid_points{0};
};

constexpr int mapping_status_count = 8;

struct GravitationalEnergies {
  double binding{0.0};
  double virial{0.0};
  double minimum_mapping_determinant{std::numeric_limits<double>::infinity()};
  double maximum_mapping_determinant{-std::numeric_limits<double>::infinity()};
  long long invalid_points{0};
  std::array<long long, mapping_status_count> mapping_status_counts{};
  int first_invalid_element{-1};
  int first_invalid_attribute{-1};
  int first_invalid_quadrature_point{-1};
  double first_invalid_determinant{std::numeric_limits<double>::quiet_NaN()};
};

constexpr int
mapping_status_index(const mean_field::mapping::MappingStatus status) {
  return static_cast<int>(status);
}

void zero_vacuum_density(const mean_field::fem::FEM &f,
                         mfem::GridFunction &density) {
  using DomainSchema =
      mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

  const mean_field::field::FieldDofMap densityMap =
      mean_field::field::make_field_dof_map<mean_field::field::Density,
                                            DomainSchema>(*f.densityFes);

  mfem::Vector densityTrue;
  density.GetTrueDofs(densityTrue);

  const mfem::Vector supportedDensity = densityMap.gather(densityTrue);
  densityMap.scatter(supportedDensity, densityTrue);
  density.SetFromTrueDofs(densityTrue);
}

double global_norm(const mfem::Vector &vector, MPI_Comm communicator) {
  const double local_norm_squared = vector * vector;
  double global_norm_squared = 0.0;
  MPI_Allreduce(&local_norm_squared, &global_norm_squared, 1, MPI_DOUBLE,
                MPI_SUM, communicator);
  return std::sqrt(global_norm_squared);
}

double global_dot(const mfem::Vector &lhs, const mfem::Vector &rhs,
                  MPI_Comm communicator) {
  const double local_dot = lhs * rhs;
  double result = 0.0;
  MPI_Allreduce(&local_dot, &result, 1, MPI_DOUBLE, MPI_SUM, communicator);
  return result;
}

double global_relative_error(const mfem::Vector &computed,
                             const mfem::Vector &reference,
                             MPI_Comm communicator) {
  REQUIRE(computed.Size() == reference.Size());

  mfem::Vector difference(computed);
  difference -= reference;

  return global_norm(difference, communicator) /
         std::max(global_norm(reference, communicator),
                  std::numeric_limits<double>::epsilon());
}

int get_shell(const double coordinate) {
  const double clamped = std::clamp(coordinate, 0.0, std::nextafter(1.0, 0.0));

  for (int shell = 0; shell < shell_count; ++shell) {
    if (clamped < shell_boundaries[shell + 1]) {
      return shell;
    }
  }

  return shell_count - 1;
}

bool retryable_infinity_status(
    const mean_field::mapping::MappingStatus status) {
  return status ==
             mean_field::mapping::MappingStatus::at_compactified_infinity ||
         status ==
             mean_field::mapping::MappingStatus::outside_reference_domain ||
         status == mean_field::mapping::MappingStatus::non_finite_result ||
         status == mean_field::mapping::MappingStatus::non_positive_determinant;
}

class ProjectionGeometry {
public:
  ProjectionGeometry(const mean_field::fem::FEM &f,
                     const mean_field::mapping::DomainMapper &mapper,
                     const mfem::GridFunction &displacement)
      : m_fem(f), m_mapper(mapper), m_displacement(displacement),
        m_workspace(f.mesh->Dimension()) {}

  mean_field::mapping::MappingStatus
  Evaluate(mfem::ElementTransformation &transformation,
           const mfem::IntegrationPoint &integration_point,
           mean_field::mapping::MappingPointContext &context,
           const bool permit_infinity_limit) {
    m_used_infinity_limit = false;

    const int element_id = transformation.ElementNo;
    MFEM_VERIFY(element_id >= 0,
                "Projection coefficient received an invalid element number.");

    const mfem::FiniteElement &displacement_element =
        *m_fem.displacementFes->GetFE(element_id);
    const mfem::FiniteElement &compactification_element =
        *m_fem.compactificationFes->GetFE(element_id);

    mfem::Array<int> displacement_dofs;
    mfem::Array<int> compactification_dofs;

    mfem::DofTransformation *displacement_transform =
        m_fem.displacementFes->GetElementVDofs(element_id, displacement_dofs);

    mfem::DofTransformation *compactification_transform =
        m_fem.compactificationFes->GetElementDofs(element_id,
                                                  compactification_dofs);

    mfem::Vector element_displacement;
    mfem::Vector element_compactification;

    m_displacement.GetSubVector(displacement_dofs, element_displacement);
    m_fem.compactificationCoordinate->GetSubVector(compactification_dofs,
                                                   element_compactification);

    if (displacement_transform != nullptr) {
      displacement_transform->InvTransformPrimal(element_displacement);
    }

    if (compactification_transform != nullptr) {
      compactification_transform->InvTransformPrimal(element_compactification);
    }

    const mean_field::mapping::ElementDisplacementData displacement_data =
        mean_field::mapping::ElementDisplacementDataFromElementVDofs(
            displacement_element, element_displacement);

    const mean_field::mapping::ElementCompactificationData
        compactification_data(compactification_element,
                              element_compactification);

    const mean_field::mapping::ElementMappingData mapping_data{
        .displacement = displacement_data,
        .compactification = compactification_data};

    mfem::Vector compactification_shape(compactification_element.GetDof());
    compactification_element.CalcShape(integration_point,
                                       compactification_shape);

    const double coordinate = element_compactification * compactification_shape;

    mean_field::mapping::MappingStatus status = m_mapper.EvaluatePoint(
        mapping_data, transformation, integration_point, m_workspace, context);

    if (status == mean_field::mapping::MappingStatus::valid) {
      transformation.SetIntPoint(&integration_point);
      return status;
    }

    const bool infinity_request =
        m_mapper.IsCompactifiedElement(transformation) &&
        coordinate >= 1.0 - 1.0e-10;

    if (!permit_infinity_limit || !infinity_request ||
        !retryable_infinity_status(status)) {
      transformation.SetIntPoint(&integration_point);
      return status;
    }

    const mfem::IntegrationPoint &center =
        mfem::Geometries.GetCenter(transformation.GetGeometryType());

    constexpr std::array<double, 11> inward_fractions{
        1.0e-12, 1.0e-11, 1.0e-10, 1.0e-9, 1.0e-8, 1.0e-7,
        1.0e-6,  1.0e-5,  1.0e-4,  1.0e-3, 1.0e-2};

    for (const double fraction : inward_fractions) {
      mfem::IntegrationPoint inward;
      inward.x = (1.0 - fraction) * integration_point.x + fraction * center.x;
      inward.y = (1.0 - fraction) * integration_point.y + fraction * center.y;
      inward.z = (1.0 - fraction) * integration_point.z + fraction * center.z;
      inward.weight = integration_point.weight;

      status = m_mapper.EvaluatePoint(mapping_data, transformation, inward,
                                      m_workspace, context);

      if (status == mean_field::mapping::MappingStatus::valid) {
        m_used_infinity_limit = true;
        transformation.SetIntPoint(&integration_point);
        return status;
      }

      if (!retryable_infinity_status(status)) {
        break;
      }
    }

    transformation.SetIntPoint(&integration_point);
    return status;
  }

  [[nodiscard]] bool UsedInfinityLimit() const noexcept {
    return m_used_infinity_limit;
  }

private:
  const mean_field::fem::FEM &m_fem;
  const mean_field::mapping::DomainMapper &m_mapper;
  const mfem::GridFunction &m_displacement;
  mean_field::mapping::DomainMapper::Workspace m_workspace;
  bool m_used_infinity_limit{false};
};

class MonopolePotentialCoefficient final : public mfem::Coefficient {
public:
  MonopolePotentialCoefficient(
      const mean_field::fem::FEM &f,
      const mean_field::mapping::DomainMapper &mapper,
      const mfem::GridFunction &displacement, const double mass,
      const double radius)
      : m_geometry(f, mapper, displacement),
        m_vacuum_attribute(field_dof_test_utils::vacuum_material_attribute),
        m_mass(mass), m_radius(radius) {}

  double Eval(mfem::ElementTransformation &transformation,
              const mfem::IntegrationPoint &integration_point) override {
    mean_field::mapping::MappingPointContext context;

    const mean_field::mapping::MappingStatus status =
        m_geometry.Evaluate(transformation, integration_point, context, true);

    MFEM_VERIFY(status == mean_field::mapping::MappingStatus::valid,
                "Stateless monopole-potential projection failed with status "
                    << static_cast<int>(status) << " on element "
                    << transformation.ElementNo << '.');

    if (m_geometry.UsedInfinityLimit()) {
      return 0.0;
    }

    const double radius = context.physical_position.Norml2();

    MFEM_VERIFY(std::isfinite(radius) && radius > 0.0,
                "Invalid monopole projection radius.");

    if (transformation.Attribute == m_vacuum_attribute) {
      return -mean_field::utils::G * m_mass / radius;
    }

    return -mean_field::utils::G * m_mass *
           (3.0 * m_radius * m_radius - radius * radius) /
           (2.0 * m_radius * m_radius * m_radius);
  }

private:
  ProjectionGeometry m_geometry;
  int m_vacuum_attribute;
  double m_mass;
  double m_radius;
};

void local_to_true(const mfem::ParFiniteElementSpace &space,
                   const mfem::Vector &local, mfem::Vector &true_vector) {
  true_vector.SetSize(space.GetTrueVSize());
  true_vector = 0.0;

  const mfem::Operator *prolongation = space.GetProlongationMatrix();

  if (prolongation != nullptr) {
    prolongation->MultTranspose(local, true_vector);
  } else {
    true_vector = local;
  }
}

mfem::Vector assemble_monopole_projection_rhs(
    mean_field::fem::FEM &f, const mfem::GridFunction &displacement,
    const double mass, const double stellar_radius) {
  mfem::Vector local_rhs(f.gravityFluxFes->GetVSize());
  local_rhs = 0.0;

  mean_field::mapping::DomainMapper::Workspace workspace(
      f.mesh->Dimension());

  const int vacuum_attribute = field_dof_test_utils::vacuum_material_attribute;
  const int quadrature_order = 2 * f.gravityFluxFes->GetMaxElementOrder() + 8;

  for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
    const mfem::FiniteElement &gravity_element =
        *f.gravityFluxFes->GetFE(element_id);
    const mfem::FiniteElement &displacement_element =
        *f.displacementFes->GetFE(element_id);
    const mfem::FiniteElement &compactification_element =
        *f.compactificationFes->GetFE(element_id);

    mfem::ElementTransformation *transformation =
        f.mesh->GetElementTransformation(element_id);

    mfem::Array<int> gravity_dofs;
    mfem::Array<int> displacement_dofs;
    mfem::Array<int> compactification_dofs;

    mfem::DofTransformation *gravity_transform =
        f.gravityFluxFes->GetElementVDofs(element_id, gravity_dofs);
    mfem::DofTransformation *displacement_transform =
        f.displacementFes->GetElementVDofs(element_id, displacement_dofs);
    mfem::DofTransformation *compactification_transform =
        f.compactificationFes->GetElementDofs(element_id,
                                              compactification_dofs);

    mfem::Vector element_displacement;
    mfem::Vector element_compactification;

    displacement.GetSubVector(displacement_dofs, element_displacement);
    f.compactificationCoordinate->GetSubVector(compactification_dofs,
                                               element_compactification);

    if (displacement_transform != nullptr) {
      displacement_transform->InvTransformPrimal(element_displacement);
    }

    if (compactification_transform != nullptr) {
      compactification_transform->InvTransformPrimal(element_compactification);
    }

    const mean_field::mapping::ElementDisplacementData displacement_data =
        mean_field::mapping::ElementDisplacementDataFromElementVDofs(
            displacement_element, element_displacement);

    const mean_field::mapping::ElementCompactificationData
        compactification_data(compactification_element,
                              element_compactification);

    const mean_field::mapping::ElementMappingData mapping_data{
        .displacement = displacement_data,
        .compactification = compactification_data};

    const int dof_count = gravity_element.GetDof();
    const int dimension = transformation->GetSpaceDim();

    mfem::Vector element_rhs(dof_count);
    mfem::Vector analytic_field(dimension);
    mfem::Vector pulled_rhs_field(dimension);
    mfem::DenseMatrix vector_shape(dof_count, dimension);

    element_rhs = 0.0;

    const mfem::IntegrationRule &rule =
        mfem::IntRules.Get(transformation->GetGeometryType(), quadrature_order);

    for (int q = 0; q < rule.GetNPoints(); ++q) {
      const mfem::IntegrationPoint &point = rule.IntPoint(q);

      mean_field::mapping::VolumeMappingContext context;

      const mean_field::mapping::MappingStatus status =
          f.domainMapperStateless->EvaluateVolume(mapping_data, *transformation,
                                                  point, workspace, context);

      MFEM_VERIFY(status == mean_field::mapping::MappingStatus::valid,
                  "Mapped monopole projection RHS failed with status "
                      << static_cast<int>(status) << " on element "
                      << element_id << ", quadrature point " << q << '.');

      analytic_field = context.mapping.physical_position;

      const double radius = analytic_field.Norml2();

      MFEM_VERIFY(std::isfinite(radius) && radius > 0.0,
                  "Invalid physical radius in projection RHS.");

      if (transformation->Attribute == vacuum_attribute) {
        analytic_field *=
            mean_field::utils::G * mass / (radius * radius * radius);
      } else {
        analytic_field *= mean_field::utils::G * mass /
                          (stellar_radius * stellar_radius * stellar_radius);
      }

      context.mapping.mapping_jacobian.MultTranspose(analytic_field,
                                                     pulled_rhs_field);

      transformation->SetIntPoint(&point);
      gravity_element.CalcVShape(*transformation, vector_shape);

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

      for (int i = 0; i < dof_count; ++i) {
        for (int component = 0; component < dimension; ++component) {
          element_rhs(i) +=
              weight * vector_shape(i, component) * pulled_rhs_field(component);
        }
      }
    }

    if (gravity_transform != nullptr) {
      gravity_transform->TransformDual(element_rhs);
    }

    local_rhs.AddElementVector(gravity_dofs, element_rhs);
  }

  mfem::Vector true_rhs;
  local_to_true(*f.gravityFluxFes, local_rhs, true_rhs);

  return true_rhs;
}

mfem::Vector
project_monopole_gradient(mean_field::fem::FEM &f,
                          const mfem::ParGridFunction &displacement,
                          const double mass, const double stellar_radius) {
  mfem::Vector displacement_true;
  displacement.GetTrueDofs(displacement_true);

  const mfem::Vector projection_rhs =
      assemble_monopole_projection_rhs(f, displacement, mass, stellar_radius);

  mean_field::operators::PreparedMappedHDivMassOperator mass_operator(
      f, *f.domainMapperStateless);
  mass_operator.Prepare(
      mass_operator.GetDisplacementMap().gather(displacement_true));
  const mfem::Vector reduced_projection_rhs =
      mass_operator.GetFluxMap().gather(projection_rhs);

  mfem::Vector projected_gradient(mass_operator.Width());
  projected_gradient = 0.0;

  mfem::CGSolver solver(f.gravityFluxFes->GetComm());
  solver.SetOperator(mass_operator);
  solver.SetRelTol(1.0e-9);
  solver.SetAbsTol(1.0e-12);
  solver.SetMaxIter(2000);
  solver.SetPrintLevel(0);
  solver.Mult(reduced_projection_rhs, projected_gradient);

  mfem::Vector projection_residual;
  mass_operator.Mult(projected_gradient, projection_residual);
  projection_residual -= reduced_projection_rhs;

  const double source_norm =
      global_norm(reduced_projection_rhs, f.gravityFluxFes->GetComm());
  const double residual_norm =
      global_norm(projection_residual, f.gravityFluxFes->GetComm());
  const double relative_residual =
      residual_norm /
      std::max(source_norm, std::numeric_limits<double>::epsilon());

  INFO("Mapped H(div) projection converged = " << solver.GetConverged());
  INFO("Mapped H(div) projection iterations = " << solver.GetNumIterations());
  INFO("Mapped H(div) projection reported final norm = "
       << solver.GetFinalNorm());
  INFO("Mapped H(div) projection direct residual norm = " << residual_norm);
  INFO("Mapped H(div) projection direct relative residual = "
       << relative_residual);

  REQUIRE(std::isfinite(relative_residual));
  REQUIRE(relative_residual < 1.0e-8);
  return mass_operator.GetFluxMap().scatter(projected_gradient);
}

double mapped_hdiv_relative_error(mean_field::fem::FEM &f,
                                  const mfem::ParGridFunction &displacement,
                                  const mfem::Vector &computed,
                                  const mfem::Vector &reference) {
  REQUIRE(computed.Size() == reference.Size());

  mfem::Vector displacement_true;
  displacement.GetTrueDofs(displacement_true);

  mean_field::operators::PreparedMappedHDivMassOperator mass_operator(
      f, *f.domainMapperStateless);
  mass_operator.Prepare(
      mass_operator.GetDisplacementMap().gather(displacement_true));

  mfem::Vector difference(computed);
  difference -= reference;

  mfem::Vector difference_action;
  mfem::Vector reference_action;

  const mfem::Vector reduced_difference =
      mass_operator.GetFluxMap().gather(difference);
  const mfem::Vector reduced_reference =
      mass_operator.GetFluxMap().gather(reference);
  mass_operator.Mult(reduced_difference, difference_action);
  mass_operator.Mult(reduced_reference, reference_action);

  MPI_Comm communicator = f.gravityFluxFes->GetComm();

  const double difference_energy =
      global_dot(reduced_difference, difference_action, communicator);
  const double reference_energy =
      global_dot(reduced_reference, reference_action, communicator);

  REQUIRE(difference_energy >= -1.0e-12 * std::abs(reference_energy));
  REQUIRE(reference_energy > 0.0);

  return std::sqrt(std::max(0.0, difference_energy) / reference_energy);
}

ShellMeasurement
measure_exterior_shells(mean_field::fem::FEM &f,
                        const mean_field::physics::GravitySolution &solution,
                        const mfem::GridFunction &displacement,
                        const double mass) {
  std::array<ShellAccumulator, shell_count> local{};
  long long local_invalid_points = 0;

  mean_field::mapping::DomainMapper::Workspace workspace(
      f.mesh->Dimension());

  const int vacuum_attribute = field_dof_test_utils::vacuum_material_attribute;
  const int quadrature_order =
      2 * std::max(f.gravityPotentialFes->GetMaxElementOrder(),
                   f.gravityFluxFes->GetMaxElementOrder()) +
      8;

  for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
    mfem::ElementTransformation *transformation =
        f.mesh->GetElementTransformation(element_id);

    if (transformation->Attribute != vacuum_attribute) {
      continue;
    }

    const mfem::FiniteElement &displacement_element =
        *f.displacementFes->GetFE(element_id);
    const mfem::FiniteElement &compactification_element =
        *f.compactificationFes->GetFE(element_id);

    mfem::Array<int> displacement_dofs;
    mfem::Array<int> compactification_dofs;

    mfem::DofTransformation *displacement_transform =
        f.displacementFes->GetElementVDofs(element_id, displacement_dofs);
    mfem::DofTransformation *compactification_transform =
        f.compactificationFes->GetElementDofs(element_id,
                                              compactification_dofs);

    mfem::Vector element_displacement;
    mfem::Vector element_compactification;

    displacement.GetSubVector(displacement_dofs, element_displacement);
    f.compactificationCoordinate->GetSubVector(compactification_dofs,
                                               element_compactification);

    if (displacement_transform != nullptr) {
      displacement_transform->InvTransformPrimal(element_displacement);
    }

    if (compactification_transform != nullptr) {
      compactification_transform->InvTransformPrimal(element_compactification);
    }

    const mean_field::mapping::ElementDisplacementData displacement_data =
        mean_field::mapping::ElementDisplacementDataFromElementVDofs(
            displacement_element, element_displacement);

    const mean_field::mapping::ElementCompactificationData
        compactification_data(compactification_element,
                              element_compactification);

    const mean_field::mapping::ElementMappingData mapping_data{
        .displacement = displacement_data,
        .compactification = compactification_data};

    mfem::Vector compactification_shape(compactification_element.GetDof());

    const mfem::IntegrationRule &rule =
        mfem::IntRules.Get(transformation->GetGeometryType(), quadrature_order);

    for (int q = 0; q < rule.GetNPoints(); ++q) {
      const mfem::IntegrationPoint &point = rule.IntPoint(q);

      transformation->SetIntPoint(&point);
      compactification_element.CalcShape(point, compactification_shape);

      const double coordinate =
          element_compactification * compactification_shape;
      const int shell = get_shell(coordinate);

      mfem::Vector reference_field(3);
      mfem::Vector physical_field(3);
      mfem::Vector physical_position(3);

      solution.gradPhi.GetVectorValue(element_id, point, reference_field);

      mean_field::mapping::VolumeMappingContext context;

      const mean_field::mapping::MappingStatus status =
          f.domainMapperStateless->EvaluateVolume(mapping_data, *transformation,
                                                  point, workspace, context);

      if (status != mean_field::mapping::MappingStatus::valid) {
        ++local_invalid_points;
        continue;
      }

      physical_position = context.mapping.physical_position;
      mean_field::mapping::MapHDivFluxToPhysical(
          context.mapping, reference_field, physical_field);

      const double radius = physical_position.Norml2();

      if (!std::isfinite(radius) || radius <= 0.0) {
        ++local_invalid_points;
        continue;
      }

      mfem::Vector radial_unit(physical_position);
      radial_unit /= radius;

      const double radial_field = physical_field * radial_unit;

      mfem::Vector tangential_field(physical_field);
      tangential_field.Add(-radial_field, radial_unit);

      const double potential = solution.phi.GetValue(element_id, point);

      const double scaled_potential =
          -radius * potential / (mean_field::utils::G * mass);
      const double scaled_radial_field =
          radius * radius * radial_field / (mean_field::utils::G * mass);
      const double scaled_tangential_field = radius * radius *
                                             tangential_field.Norml2() /
                                             (mean_field::utils::G * mass);

      if (!std::isfinite(scaled_potential) ||
          !std::isfinite(scaled_radial_field) ||
          !std::isfinite(scaled_tangential_field)) {
        ++local_invalid_points;
        continue;
      }

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

      ShellAccumulator &accumulator = local[shell];

      ++accumulator.points;
      accumulator.weight += weight;

      accumulator.minimum_radius = std::min(accumulator.minimum_radius, radius);
      accumulator.maximum_radius = std::max(accumulator.maximum_radius, radius);

      accumulator.potential_error_squared +=
          weight * (scaled_potential - 1.0) * (scaled_potential - 1.0);
      accumulator.radial_error_squared +=
          weight * (scaled_radial_field - 1.0) * (scaled_radial_field - 1.0);
      accumulator.tangential_squared +=
          weight * scaled_tangential_field * scaled_tangential_field;
    }
  }

  MPI_Comm communicator = f.gravityFluxFes->GetComm();

  ShellMeasurement measurement;

  MPI_Allreduce(&local_invalid_points, &measurement.invalid_points, 1,
                MPI_LONG_LONG, MPI_SUM, communicator);

  for (int shell = 0; shell < shell_count; ++shell) {
    long long points = 0;

    MPI_Allreduce(&local[shell].points, &points, 1, MPI_LONG_LONG, MPI_SUM,
                  communicator);

    const double local_sums[4]{
        local[shell].weight, local[shell].potential_error_squared,
        local[shell].radial_error_squared, local[shell].tangential_squared};

    double sums[4]{};

    MPI_Allreduce(local_sums, sums, 4, MPI_DOUBLE, MPI_SUM, communicator);

    double minimum_radius = 0.0;
    double maximum_radius = 0.0;

    MPI_Allreduce(&local[shell].minimum_radius, &minimum_radius, 1, MPI_DOUBLE,
                  MPI_MIN, communicator);
    MPI_Allreduce(&local[shell].maximum_radius, &maximum_radius, 1, MPI_DOUBLE,
                  MPI_MAX, communicator);

    measurement.shells[shell] = {
        .points = points,
        .minimum_radius = minimum_radius,
        .maximum_radius = maximum_radius,
        .potential_rms_error = sums[0] > 0.0
                                   ? std::sqrt(sums[1] / sums[0])
                                   : std::numeric_limits<double>::infinity(),
        .radial_rms_error = sums[0] > 0.0
                                ? std::sqrt(sums[2] / sums[0])
                                : std::numeric_limits<double>::infinity(),
        .tangential_rms = sums[0] > 0.0
                              ? std::sqrt(sums[3] / sums[0])
                              : std::numeric_limits<double>::infinity()};
  }

  return measurement;
}

GravitationalEnergies
compute_stellar_energies(mean_field::fem::FEM &f,
                         const mfem::GridFunction &density,
                         const mean_field::physics::GravitySolution &solution,
                         const mfem::GridFunction &displacement) {
  mean_field::mapping::DomainMapper::Workspace workspace(
      f.mesh->Dimension());

  double local_binding = 0.0;
  double local_virial = 0.0;
  long long local_invalid_points = 0;
  double local_minimum_determinant = std::numeric_limits<double>::infinity();
  double local_maximum_determinant = -std::numeric_limits<double>::infinity();

  const int vacuum_attribute = field_dof_test_utils::vacuum_material_attribute;

  const int order = 2 * std::max(f.gravityPotentialFes->GetMaxElementOrder(),
                                 f.gravityFluxFes->GetMaxElementOrder()) +
                    8;
  std::array<long long, mapping_status_count> local_status_counts{};

  for (int element_id = 0; element_id < f.mesh->GetNE(); ++element_id) {
    mfem::ElementTransformation *transformation =
        f.mesh->GetElementTransformation(element_id);
    if (transformation->Attribute == vacuum_attribute) {
      continue;
    }

    const mfem::FiniteElement &displacement_element =
        *f.displacementFes->GetFE(element_id);
    const mfem::FiniteElement &compactification_element =
        *f.compactificationFes->GetFE(element_id);

    mfem::Array<int> displacement_dofs;
    mfem::Array<int> compactification_dofs;

    mfem::DofTransformation *displacement_transform =
        f.displacementFes->GetElementVDofs(element_id, displacement_dofs);
    mfem::DofTransformation *compactification_transform =
        f.compactificationFes->GetElementDofs(element_id,
                                              compactification_dofs);

    mfem::Vector element_displacement;
    mfem::Vector element_compactification;

    displacement.GetSubVector(displacement_dofs, element_displacement);
    f.compactificationCoordinate->GetSubVector(compactification_dofs,
                                               element_compactification);

    if (displacement_transform != nullptr) {
      displacement_transform->InvTransformPrimal(element_displacement);
    }

    if (compactification_transform != nullptr) {
      compactification_transform->InvTransformPrimal(element_compactification);
    }

    const mean_field::mapping::ElementDisplacementData displacement_data =
        mean_field::mapping::ElementDisplacementDataFromElementVDofs(
            displacement_element, element_displacement);

    const mean_field::mapping::ElementCompactificationData
        compactification_data(compactification_element,
                              element_compactification);

    const mean_field::mapping::ElementMappingData mapping_data{
        .displacement = displacement_data,
        .compactification = compactification_data};

    const mfem::IntegrationRule &rule =
        mfem::IntRules.Get(transformation->GetGeometryType(), order);

    for (int q = 0; q < rule.GetNPoints(); ++q) {
      const mfem::IntegrationPoint &point = rule.IntPoint(q);

      mean_field::mapping::VolumeMappingContext context;

      const mean_field::mapping::MappingStatus status =
          f.domainMapperStateless->EvaluateVolume(mapping_data, *transformation,
                                                  point, workspace, context);

      const double mapping_determinant = context.mapping.mapping_determinant;
      if (std::isfinite(mapping_determinant)) {
        local_minimum_determinant =
            std::min(local_minimum_determinant, mapping_determinant);
        local_maximum_determinant =
            std::max(local_maximum_determinant, mapping_determinant);
      }

      const int status_index = mapping_status_index(status);
      MFEM_VERIFY(status_index >= 0 && status_index < mapping_status_count,
                  "Unexpected mapping status.");
      ++local_status_counts[status_index];

      if (status != mean_field::mapping::MappingStatus::valid) {
        ++local_invalid_points;
        continue;
      }

      if (status != mean_field::mapping::MappingStatus::valid) {
        ++local_invalid_points;
        continue;
      }

      mfem::Vector reference_field(3);
      mfem::Vector physical_field(3);

      solution.gradPhi.GetVectorValue(element_id, point, reference_field);

      mean_field::mapping::MapHDivFluxToPhysical(
          context.mapping, reference_field, physical_field);

      const double rho = density.GetValue(element_id, point);
      const double phi = solution.phi.GetValue(element_id, point);

      local_binding += 0.5 * rho * phi * context.quadrature.weight;
      local_virial -= rho *
                      (context.mapping.physical_position * physical_field) *
                      context.quadrature.weight;
    }
  }

  GravitationalEnergies energies;

  MPI_Comm communicator = f.densityFes->GetComm();

  MPI_Allreduce(&local_binding, &energies.binding, 1, MPI_DOUBLE, MPI_SUM,
                communicator);
  MPI_Allreduce(&local_virial, &energies.virial, 1, MPI_DOUBLE, MPI_SUM,
                communicator);
  MPI_Allreduce(&local_invalid_points, &energies.invalid_points, 1,
                MPI_LONG_LONG, MPI_SUM, communicator);
  MPI_Allreduce(local_status_counts.data(),
                energies.mapping_status_counts.data(), mapping_status_count,
                MPI_LONG_LONG, MPI_SUM, communicator);
  MPI_Allreduce(&local_minimum_determinant,
                &energies.minimum_mapping_determinant, 1, MPI_DOUBLE, MPI_MIN,
                communicator);
  MPI_Allreduce(&local_maximum_determinant,
                &energies.maximum_mapping_determinant, 1, MPI_DOUBLE, MPI_MAX,
                communicator);

  return energies;
}
} // namespace

TEST_CASE("Gravity Field Monopole Accuracy And Projection Floor",
          tags::gravity_analytic_accuracy) {
  auto args = test_utils::setup_args();
  args.p.rtol = 1.0e-13;
  args.p.max_iters = std::max(args.p.max_iters, 1000);

  mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

  REQUIRE(f.domainMapperStateless != nullptr);
  REQUIRE(f.domainMapperStateless != nullptr);
  REQUIRE(f.compactificationCoordinate != nullptr);

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

  const double density_value =
      mass / ((4.0 / 3.0) * M_PI * radius * radius * radius);

  mfem::ParGridFunction displacement(f.displacementFes.get());
  displacement = 0.0;

  *f.displacement = 0.0;

  mfem::GridFunction density(f.densityFes.get());
  density = density_value;

  zero_vacuum_density(f, density);

  mean_field::analysis::conserve_mass(f, density, mass);

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

  const mean_field::physics::GravitySolution numerical_solution =
      mean_field::physics::solve_gravity_field(f, args, density, displacement);

  MonopolePotentialCoefficient potential_coefficient(
      f, *f.domainMapperStateless, displacement, mass, radius);

  mean_field::physics::GravitySolution projected_solution(f);
  projected_solution.phi = 0.0;
  projected_solution.gradPhi = 0.0;

  projected_solution.phi.ProjectCoefficient(potential_coefficient);

  const mfem::Vector projected_gradient_true =
      project_monopole_gradient(f, displacement, mass, radius);

  projected_solution.gradPhi.SetFromTrueDofs(projected_gradient_true);

  const ShellMeasurement numerical =
      measure_exterior_shells(f, numerical_solution, displacement, mass);

  const ShellMeasurement projected =
      measure_exterior_shells(f, projected_solution, displacement, mass);

  const GravitationalEnergies energies =
      compute_stellar_energies(f, density, numerical_solution, displacement);

  mfem::Vector numerical_gradient;
  mfem::Vector numerical_potential;
  mfem::Vector projected_gradient;
  mfem::Vector projected_potential;

  numerical_solution.gradPhi.GetTrueDofs(numerical_gradient);
  numerical_solution.phi.GetTrueDofs(numerical_potential);

  projected_solution.gradPhi.GetTrueDofs(projected_gradient);
  projected_solution.phi.GetTrueDofs(projected_potential);

  MPI_Comm communicator = f.gravityFluxFes->GetComm();

  const double gradient_projection_gap = mapped_hdiv_relative_error(
      f, displacement, numerical_gradient, projected_gradient);

  const double potential_projection_gap = global_relative_error(
      numerical_potential, projected_potential, communicator);

  double maximum_numerical_potential_error = 0.0;
  double maximum_projected_potential_error = 0.0;
  double maximum_numerical_radial_error = 0.0;
  double maximum_projected_radial_error = 0.0;
  double maximum_numerical_tangential = 0.0;
  double maximum_projected_tangential = 0.0;

  std::ostringstream report;

  for (int shell = 0; shell < shell_count; ++shell) {
    const ShellMetrics &numerical_shell = numerical.shells[shell];
    const ShellMetrics &projected_shell = projected.shells[shell];

    maximum_numerical_potential_error = std::max(
        maximum_numerical_potential_error, numerical_shell.potential_rms_error);

    maximum_projected_potential_error = std::max(
        maximum_projected_potential_error, projected_shell.potential_rms_error);

    maximum_numerical_radial_error = std::max(maximum_numerical_radial_error,
                                              numerical_shell.radial_rms_error);

    maximum_projected_radial_error = std::max(maximum_projected_radial_error,
                                              projected_shell.radial_rms_error);

    maximum_numerical_tangential =
        std::max(maximum_numerical_tangential, numerical_shell.tangential_rms);

    maximum_projected_tangential =
        std::max(maximum_projected_tangential, projected_shell.tangential_rms);

    report << "shell " << shell << " xi=[" << shell_boundaries[shell] << ", "
           << shell_boundaries[shell + 1] << ")\n"
           << "  radius=[" << numerical_shell.minimum_radius << ", "
           << numerical_shell.maximum_radius << "]\n"
           << "  potential error: solved="
           << numerical_shell.potential_rms_error
           << ", projection=" << projected_shell.potential_rms_error << '\n'
           << "  radial error: solved=" << numerical_shell.radial_rms_error
           << ", projection=" << projected_shell.radial_rms_error << '\n'
           << "  tangential amplitude: solved="
           << numerical_shell.tangential_rms
           << ", projection=" << projected_shell.tangential_rms << '\n';
  }

  const double analytic_energy =
      -3.0 * mean_field::utils::G * mass * mass / (5.0 * radius);

  const double binding_error =
      std::abs(energies.binding - analytic_energy) / std::abs(analytic_energy);
  const double virial_error =
      std::abs(energies.virial - analytic_energy) / std::abs(analytic_energy);
  const double consistency_error =
      std::abs(energies.binding - energies.virial) / std::abs(energies.binding);

  INFO(report.str());

  INFO("Gradient solution/projection mapped H(div) gap = "
       << gradient_projection_gap);
  INFO("Potential solution/projection DOF gap = " << potential_projection_gap);
  INFO("Analytic energy = " << analytic_energy);
  INFO("Computed binding energy = " << energies.binding);
  INFO("Computed virial energy = " << energies.virial);
  INFO("Relative binding error = " << binding_error);
  INFO("Relative virial error = " << virial_error);
  INFO("Relative virial consistency error = " << consistency_error);

  REQUIRE(numerical.invalid_points == 0);
  REQUIRE(projected.invalid_points == 0);
  REQUIRE(energies.invalid_points == 0);

  for (int shell = 0; shell < shell_count; ++shell) {
    REQUIRE(numerical.shells[shell].points > 0);
    REQUIRE(projected.shells[shell].points > 0);
  }

  CHECK(maximum_numerical_potential_error < 5.0e-2);
  CHECK(maximum_projected_potential_error < 5.0e-2);
  CHECK(maximum_numerical_radial_error < 5.0e-3);
  CHECK(maximum_projected_radial_error < 5.0e-3);
  CHECK(maximum_numerical_tangential < 5.0e-3);
  CHECK(maximum_projected_tangential < 5.0e-3);
  CHECK(gradient_projection_gap < 5.0e-3);
  CHECK(potential_projection_gap < maximum_numerical_potential_error);

  constexpr double virial_target = 1.0e-5;

  CHECK(binding_error < virial_target);
  CHECK(virial_error < virial_target);
  CHECK(consistency_error < virial_target);
}

TEST_CASE(
    "Gravity Field Virial Consistency Across Volume Preserving Deformation",
    tags::gravity_consistency_accuracy) {
  auto args = test_utils::setup_args();
  args.p.rtol = 1.0e-13;
  args.p.max_iters = std::max(args.p.max_iters, 1000);

  mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

  REQUIRE(f.domainMapperStateless != nullptr);
  REQUIRE(f.domainMapperStateless != nullptr);

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

  const double central_density =
      15.0 * mass / (8.0 * M_PI * radius * radius * radius);

  auto density_function = [central_density,
                           radius](const mfem::Vector &position) {
    const double normalized_radius_squared =
        (position * position) / (radius * radius);
    return central_density * std::max(0.0, 1.0 - normalized_radius_squared);
  };

  mfem::FunctionCoefficient density_coefficient(density_function);
  mfem::GridFunction density(f.densityFes.get());
  density.ProjectCoefficient(density_coefficient);

  zero_vacuum_density(f, density);

  mean_field::analysis::conserve_mass(f, density, mass);

  constexpr std::array<double, 7> amplitudes{0.0, 0.02, 0.05, 0.1,
                                             0.2, 0.5,  1.0};

  std::array<double, amplitudes.size()> consistency_errors{};
  std::array<double, amplitudes.size()> normalized_quadrupoles{};
  std::array<double, amplitudes.size()> binding_energies{};
  std::array<double, amplitudes.size()> virial_energies{};

  std::ostringstream report;

  for (std::size_t index = 0; index < amplitudes.size(); ++index) {
    const double amplitude = amplitudes[index];

    const double x_scale = 1.0 + 0.15 * amplitude;
    const double y_scale = 1.0 - 0.05 * amplitude;
    const double z_scale = 1.0 / (x_scale * y_scale);

    REQUIRE_THAT(x_scale * y_scale * z_scale,
                 Catch::Matchers::WithinAbs(1.0, 1.0e-14));

    auto displacement_function = [x_scale, y_scale,
                                  z_scale](const mfem::Vector &position,
                                           mfem::Vector &value) {
      value.SetSize(3);

      value(0) = (x_scale - 1.0) * position(0);
      value(1) = (y_scale - 1.0) * position(1);
      value(2) = (z_scale - 1.0) * position(2);
    };

    mfem::VectorFunctionCoefficient displacement_coefficient(
        3, displacement_function);
    mfem::ParGridFunction displacement(f.displacementFes.get());
    displacement.ProjectCoefficient(displacement_coefficient);

    *f.displacement = displacement;

    f.com = mean_field::analysis::get_com(f, density);
    f.Q = mean_field::physics::compute_quadrupole_moment_tensor(f, density,
                                                                f.com);
    const mfem::FiniteElementSpace *nodal_space = f.mesh->GetNodalFESpace();

    const mean_field::physics::GravitySolution solution =
        mean_field::physics::solve_gravity_field(f, args, density,
                                                 displacement);
    const GravitationalEnergies energies =
        compute_stellar_energies(f, density, solution, displacement);

    CAPTURE(amplitude, x_scale, y_scale, z_scale);
    INFO("Mapping status valid = "
         << energies.mapping_status_counts[mapping_status_index(
                mean_field::mapping::MappingStatus::valid)]);
    INFO("Mapping status invalid_dimension = "
         << energies.mapping_status_counts[mapping_status_index(
                mean_field::mapping::MappingStatus::invalid_dimension)]);
    INFO("Mapping status non_finite_input = "
         << energies.mapping_status_counts[mapping_status_index(
                mean_field::mapping::MappingStatus::non_finite_input)]);
    INFO("Mapping status invalid_reference_radius = "
         << energies.mapping_status_counts[mapping_status_index(
                mean_field::mapping::MappingStatus::invalid_reference_radius)]);
    INFO("Mapping status at_compactified_infinity = "
         << energies.mapping_status_counts[mapping_status_index(
                mean_field::mapping::MappingStatus::at_compactified_infinity)]);
    INFO("Mapping status outside_reference_domain = "
         << energies.mapping_status_counts[mapping_status_index(
                mean_field::mapping::MappingStatus::outside_reference_domain)]);
    INFO("Mapping status non_finite_result = "
         << energies.mapping_status_counts[mapping_status_index(
                mean_field::mapping::MappingStatus::non_finite_result)]);
    INFO("Mapping status non_positive_determinant = "
         << energies.mapping_status_counts[mapping_status_index(
                mean_field::mapping::MappingStatus::non_positive_determinant)]);
    INFO("Total invalid mapping points = " << energies.invalid_points);
    INFO("Mesh nodal order = "
         << (nodal_space != nullptr ? nodal_space->GetMaxElementOrder() : -1));
    INFO("Displacement order = " << f.displacementFes->GetMaxElementOrder());
    INFO("Minimum discrete mapping determinant = "
         << energies.minimum_mapping_determinant);
    INFO("Maximum discrete mapping determinant = "
         << energies.maximum_mapping_determinant);

    REQUIRE(energies.invalid_points == 0);
    REQUIRE(std::isfinite(energies.binding));
    REQUIRE(std::isfinite(energies.virial));
    REQUIRE(energies.binding < 0.0);
    REQUIRE(energies.virial < 0.0);

    binding_energies[index] = energies.binding;
    virial_energies[index] = energies.virial;

    consistency_errors[index] = std::abs(energies.binding - energies.virial) /
                                std::abs(energies.binding);

    normalized_quadrupoles[index] = f.Q.FNorm() / (mass * radius * radius);

    report << "amplitude=" << amplitude << ", scales=(" << x_scale << ", "
           << y_scale << ", " << z_scale
           << "), normalized quadrupole=" << normalized_quadrupoles[index]
           << ", binding=" << binding_energies[index]
           << ", virial=" << virial_energies[index]
           << ", consistency error=" << consistency_errors[index] << '\n';
  }

  INFO(report.str());

  for (std::size_t index = 1; index < amplitudes.size(); ++index) {
    CHECK(normalized_quadrupoles[index] > normalized_quadrupoles[index - 1]);
  }

  CHECK(consistency_errors[0] < 1.0e-5);
  for (std::size_t index = 1; index < amplitudes.size(); ++index) {
    CHECK(consistency_errors[index] < 5.0e-5);
  }
}