#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <utility>

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

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace polytropic_eos_test_utils {
template <typename Function>
double centered_derivative(Function &&function, const double position,
                           const double step) {
  return (function(position + step) - function(position - step)) / (2.0 * step);
}

template <typename Integrand>
double integrate_cube(const mfem::IntegrationRule &integrationRule,
                      Integrand &&integrand) {
  double integral = 0.0;

  for (int pointIndex = 0; pointIndex < integrationRule.GetNPoints();
       ++pointIndex) {
    const mfem::IntegrationPoint &integrationPoint =
        integrationRule.IntPoint(pointIndex);

    integral += integrationPoint.weight * integrand(integrationPoint);
  }

  return integral;
}
} // namespace polytropic_eos_test_utils

TEST_CASE("Polytropic EOS Satisfies Its Thermodynamic Identities",
          tags::barotrope_eos_unit) {
  constexpr std::array<double, 3> polytropicIndices{1.0, 1.5, 3.0};

  constexpr std::array<double, 4> densities{1.0e-4, 0.02, 0.37, 2.4};

  constexpr double polytropicConstant = 0.73;

  for (const double polytropicIndex : polytropicIndices) {
    DYNAMIC_SECTION("polytropic index n = " << polytropicIndex) {
      const mean_field::eos::Polytrope barotrope(polytropicIndex,
                                                 polytropicConstant);

      const double expectedEnthalpyScale =
          (polytropicIndex + 1.0) * polytropicConstant;

      CHECK(barotrope.polytropic_index() == polytropicIndex);

      CHECK(barotrope.polytropic_constant() == polytropicConstant);

      CHECK(barotrope.enthalpy_scale() == expectedEnthalpyScale);

      for (const double density : densities) {
        CAPTURE(polytropicIndex, polytropicConstant, density);

        const double expectedPressure =
            polytropicConstant * std::pow(density, 1.0 + 1.0 / polytropicIndex);

        const double expectedEnthalpy =
            expectedEnthalpyScale * std::pow(density, 1.0 / polytropicIndex);

        const double pressureFromDensity =
            barotrope.pressure_from_density(density);

        const double enthalpyFromDensity =
            barotrope.enthalpy_from_density(density);

        const double recoveredDensity =
            barotrope.density_from_enthalpy(enthalpyFromDensity);

        const double pressureFromEnthalpy =
            barotrope.pressure_from_enthalpy(enthalpyFromDensity);

        CHECK_THAT(pressureFromDensity,
                   Catch::Matchers::WithinRel(expectedPressure, 2.0e-13));

        CHECK_THAT(enthalpyFromDensity,
                   Catch::Matchers::WithinRel(expectedEnthalpy, 2.0e-13));

        CHECK_THAT(recoveredDensity,
                   Catch::Matchers::WithinRel(density, 5.0e-13));

        CHECK_THAT(pressureFromEnthalpy,
                   Catch::Matchers::WithinRel(expectedPressure, 5.0e-13));

        /*
         * Polytropic identity:
         *
         *     P = rho h / (n + 1).
         */
        CHECK_THAT(pressureFromEnthalpy,
                   Catch::Matchers::WithinRel(density * enthalpyFromDensity /
                                                  (polytropicIndex + 1.0),
                                              5.0e-13));

        /*
         * Polytropic identity:
         *
         *     dP / dh = rho.
         *
         * The implementation should return the same
         * value as density_from_enthalpy().
         */
        CHECK(
            barotrope.pressure_derivative_from_enthalpy(enthalpyFromDensity) ==
            barotrope.density_from_enthalpy(enthalpyFromDensity));

        /*
         * Since
         *
         *     h = (n + 1) K rho^(1/n),
         *
         * it follows that
         *
         *     dP / d rho = h / n.
         */
        CHECK_THAT(barotrope.pressure_derivative_from_density(density),
                   Catch::Matchers::WithinRel(
                       enthalpyFromDensity / polytropicIndex, 5.0e-13));
      }
    }
  }
}

TEST_CASE("Polytropic EOS Pressure Derivatives Match Centered Differences",
          tags::barotrope_eos_jacobian) {
  constexpr std::array<double, 3> polytropicIndices{1.0, 1.5, 3.0};

  constexpr std::array<double, 3> positiveValues{0.2, 0.73, 1.8};

  constexpr double polytropicConstant = 0.61;

  for (const double polytropicIndex : polytropicIndices) {
    const mean_field::eos::Polytrope barotrope(polytropicIndex,
                                               polytropicConstant);

    DYNAMIC_SECTION("polytropic index n = " << polytropicIndex) {
      for (const double enthalpy : positiveValues) {
        const double step = 2.0e-6 * std::max(1.0, std::abs(enthalpy));

        const double numericalDerivative =
            polytropic_eos_test_utils::centered_derivative(
                [&barotrope](const double perturbedEnthalpy) {
                  return barotrope.pressure_from_enthalpy(perturbedEnthalpy);
                },
                enthalpy, step);

        const double analyticDerivative =
            barotrope.pressure_derivative_from_enthalpy(enthalpy);

        CAPTURE(polytropicIndex, enthalpy, step, numericalDerivative,
                analyticDerivative);

        CHECK_THAT(numericalDerivative,
                   Catch::Matchers::WithinRel(analyticDerivative, 5.0e-8));
      }

      for (const double density : positiveValues) {
        const double step = 2.0e-6 * std::max(1.0, std::abs(density));

        const double numericalDerivative =
            polytropic_eos_test_utils::centered_derivative(
                [&barotrope](const double perturbedDensity) {
                  return barotrope.pressure_from_density(perturbedDensity);
                },
                density, step);

        const double analyticDerivative =
            barotrope.pressure_derivative_from_density(density);

        CAPTURE(polytropicIndex, density, step, numericalDerivative,
                analyticDerivative);

        CHECK_THAT(numericalDerivative,
                   Catch::Matchers::WithinRel(analyticDerivative, 5.0e-8));
      }
    }
  }
}

TEST_CASE("Polytropic EOS Density Derivative Matches Centered Differences",
          tags::barotrope_eos_jacobian) {
  constexpr std::array<double, 3> polytropicIndices{1.0, 1.5, 3.0};

  constexpr std::array<double, 3> enthalpies{0.2, 0.73, 1.8};

  constexpr double polytropicConstant = 0.61;

  for (const double polytropicIndex : polytropicIndices) {
    const mean_field::eos::Polytrope barotrope(polytropicIndex,
                                               polytropicConstant);

    DYNAMIC_SECTION("polytropic index n = " << polytropicIndex) {
      for (const double enthalpy : enthalpies) {
        const double step = 2.0e-6 * std::max(1.0, std::abs(enthalpy));

        const double numericalDerivative =
            polytropic_eos_test_utils::centered_derivative(
                [&barotrope](const double perturbedEnthalpy) {
                  return barotrope.density_from_enthalpy(perturbedEnthalpy);
                },
                enthalpy, step);

        const double analyticDerivative =
            barotrope.density_derivative_from_enthalpy(enthalpy);

        CAPTURE(polytropicIndex, enthalpy, step, numericalDerivative,
                analyticDerivative);

        CHECK_THAT(numericalDerivative,
                   Catch::Matchers::WithinRel(analyticDerivative, 5.0e-8));
      }
    }
  }
}

TEST_CASE("Polytropic EOS Defines Consistent Surface And Exterior Behavior",
          tags::barotrope_eos_unit) {
  constexpr std::array<double, 3> polytropicIndices{1.0, 1.5, 3.0};

  constexpr double polytropicConstant = 0.47;
  constexpr double exteriorEnthalpy = -0.3;

  for (const double polytropicIndex : polytropicIndices) {
    const mean_field::eos::Polytrope barotrope(polytropicIndex,
                                               polytropicConstant);

    DYNAMIC_SECTION("polytropic index n = " << polytropicIndex) {
      /*
       * Exact surface values.
       */
      CHECK(barotrope.density_from_enthalpy(0.0) == 0.0);

      CHECK(barotrope.pressure_from_enthalpy(0.0) == 0.0);

      CHECK(barotrope.pressure_derivative_from_enthalpy(0.0) == 0.0);

      CHECK(barotrope.pressure_from_density(0.0) == 0.0);

      CHECK(barotrope.enthalpy_from_density(0.0) == 0.0);

      CHECK(barotrope.pressure_derivative_from_density(0.0) == 0.0);

      /*
       * Positive-part extension into h < 0.
       */
      CHECK(barotrope.density_from_enthalpy(exteriorEnthalpy) == 0.0);

      CHECK(barotrope.pressure_from_enthalpy(exteriorEnthalpy) == 0.0);

      CHECK(barotrope.density_derivative_from_enthalpy(exteriorEnthalpy) ==
            0.0);

      CHECK(barotrope.pressure_derivative_from_enthalpy(exteriorEnthalpy) ==
            0.0);

      /*
       * At h = 0, rho(h) has a nonzero right
       * derivative only for n = 1.
       */
      const double expectedSurfaceDensityDerivative =
          polytropicIndex == 1.0 ? 1.0 / barotrope.enthalpy_scale() : 0.0;

      CHECK(barotrope.density_derivative_from_enthalpy(0.0) ==
            expectedSurfaceDensityDerivative);
    }
  }
}

TEST_CASE("Polytropic EOS Rejects Invalid Physical Inputs",
          tags::barotrope_eos_unit) {
  CHECK_THROWS_AS(mean_field::eos::Polytrope(0.999, 1.0),
                  std::invalid_argument);

  CHECK_THROWS_AS(
      mean_field::eos::Polytrope(std::numeric_limits<double>::infinity(), 1.0),
      std::invalid_argument);

  CHECK_THROWS_AS(mean_field::eos::Polytrope(3.0, 0.0), std::invalid_argument);

  CHECK_THROWS_AS(mean_field::eos::Polytrope(3.0, -1.0), std::invalid_argument);

  const mean_field::eos::Polytrope barotrope(3.0, 0.75);

  CHECK_THROWS_AS(barotrope.pressure_from_density(-0.1), std::domain_error);

  CHECK_THROWS_AS(barotrope.enthalpy_from_density(-0.1), std::domain_error);

  CHECK_THROWS_AS(barotrope.pressure_derivative_from_density(-0.1),
                  std::domain_error);

  constexpr std::array<double, 3> nonfiniteValues{
      std::numeric_limits<double>::infinity(),
      -std::numeric_limits<double>::infinity(),
      std::numeric_limits<double>::quiet_NaN()};

  for (const double nonfiniteValue : nonfiniteValues) {
    CAPTURE(nonfiniteValue);

    CHECK_THROWS_AS(barotrope.density_from_enthalpy(nonfiniteValue),
                    std::domain_error);

    CHECK_THROWS_AS(barotrope.pressure_from_enthalpy(nonfiniteValue),
                    std::domain_error);

    CHECK_THROWS_AS(barotrope.density_derivative_from_enthalpy(nonfiniteValue),
                    std::domain_error);

    CHECK_THROWS_AS(barotrope.pressure_derivative_from_enthalpy(nonfiniteValue),
                    std::domain_error);
  }
}

TEST_CASE("Pressure Force And Pressure Integral Have Distinct Registered Forms",
          tags::barotrope_pressure_quadrature_unit) {
  using EnthalpyField = mean_field::field::Field<mean_field::field::Enthalpy>;

  /*
   * For the registered H1 order p = 3 and n = 3:
   *
   *     h has degree p,
   *     P(h) has degree 4p,
   *
   * so the nonlinear EOS contributes an additional
   *
   *     4p - p = 3p = 9
   *
   * beyond the registered enthalpy operand.
   */
  constexpr int enthalpyOrder =
      mean_field::field::Enthalpy::Scalar::familyOrder;

  constexpr int pressureExtraOrder = 3 * enthalpyOrder;

  constexpr int geometryWeightOrder = 2;

  constexpr mean_field::quadrature::Query pressureIntegralQuery =
      EnthalpyField::make_query<
          mean_field::field::Enthalpy::Form::PressureIntegral>(
          mean_field::quadrature::QuadratureRole::diagnostic,
          geometryWeightOrder, std::array<int, 1>{pressureExtraOrder},
          mean_field::utils::DOMAINS::STELLAR,
          mean_field::quadrature::MappingKind::general);

  constexpr mean_field::quadrature::Query pressureForceQuery =
      EnthalpyField::make_query<
          mean_field::field::Enthalpy::Form::PressureForce>(
          mean_field::quadrature::QuadratureRole::discretization,
          geometryWeightOrder, std::array<int, 1>{pressureExtraOrder},
          mean_field::utils::DOMAINS::STELLAR,
          mean_field::quadrature::MappingKind::general);

  STATIC_CHECK(
      mean_field::field::Enthalpy::Form::PressureIntegral::dynamicOrderCount ==
      1);

  STATIC_CHECK(
      mean_field::field::Enthalpy::Form::PressureForce::dynamicOrderCount == 1);

  STATIC_CHECK(mean_field::field::Enthalpy::Form::PressureIntegral::policyKey !=
               mean_field::field::Enthalpy::Form::PressureForce::policyKey);

  REQUIRE(pressureIntegralQuery.base_order.has_value());

  REQUIRE(pressureForceQuery.base_order.has_value());

  /*
   * Pressure integral:
   *
   *     degree(P) + degree(J)
   *     = 12 + 2
   *     = 14.
   */
  CHECK(*pressureIntegralQuery.base_order == 14);

  /*
   * Pressure force:
   *
   *     degree(P)
   *       + degree(grad w)
   *       + degree(J)
   *
   *     = 12 + 2 + 2
   *     = 16.
   */
  CHECK(*pressureForceQuery.base_order == 16);

  CHECK(pressureIntegralQuery.term ==
        mean_field::quadrature::Term::pressure_integral);

  CHECK(pressureForceQuery.term ==
        mean_field::quadrature::Term::pressure_force);

  CHECK(pressureIntegralQuery.role ==
        mean_field::quadrature::QuadratureRole::diagnostic);

  CHECK(pressureForceQuery.role ==
        mean_field::quadrature::QuadratureRole::discretization);

  CHECK(pressureIntegralQuery.domain == mean_field::utils::DOMAINS::STELLAR);

  CHECK(pressureForceQuery.domain == mean_field::utils::DOMAINS::STELLAR);

  /*
   * Verify that the two terms route to independent policy
   * controls.
   */
  mean_field::quadrature::RuleSet ruleSet =
      mean_field::quadrature::make_rule_set(
          mean_field::quadrature::Mode::production);

  ruleSet.pressure_integral.boost = 3;
  ruleSet.pressure_force.boost = 5;

  const mean_field::quadrature::Policy policy(std::move(ruleSet));

  const mean_field::quadrature::Resolution pressureIntegralResolution =
      policy.resolve(pressureIntegralQuery);

  const mean_field::quadrature::Resolution pressureForceResolution =
      policy.resolve(pressureForceQuery);

  CHECK(pressureIntegralResolution.base_order == 14);

  CHECK(pressureIntegralResolution.boost == 3);

  CHECK(pressureIntegralResolution.order == 17);

  CHECK(pressureForceResolution.base_order == 16);

  CHECK(pressureForceResolution.boost == 5);

  CHECK(pressureForceResolution.order == 21);
}

TEST_CASE("Pressure Quadrature Exactly Integrates An N Three Polynomial",
          tags::barotrope_pressure_quadrature_accuracy) {
  using EnthalpyField = mean_field::field::Field<mean_field::field::Enthalpy>;

  constexpr int enthalpyOrder =
      mean_field::field::Enthalpy::Scalar::familyOrder;

  constexpr int pressureExtraOrder = 3 * enthalpyOrder;

  /*
   * K = 1/4 and n = 3 give
   *
   *     (n + 1) K = 1,
   *     rho(h) = h^3,
   *     P(h) = h^4 / 4.
   */
  const mean_field::eos::Polytrope barotrope(3.0, 0.25);

  constexpr mean_field::quadrature::Query pressureIntegralQuery =
      EnthalpyField::make_query<
          mean_field::field::Enthalpy::Form::PressureIntegral>(
          mean_field::quadrature::QuadratureRole::diagnostic, 0,
          std::array<int, 1>{pressureExtraOrder},
          mean_field::utils::DOMAINS::STELLAR,
          mean_field::quadrature::MappingKind::affine);

  constexpr mean_field::quadrature::Query pressureForceQuery =
      EnthalpyField::make_query<
          mean_field::field::Enthalpy::Form::PressureForce>(
          mean_field::quadrature::QuadratureRole::discretization, 0,
          std::array<int, 1>{pressureExtraOrder},
          mean_field::utils::DOMAINS::STELLAR,
          mean_field::quadrature::MappingKind::affine);

  const mean_field::quadrature::RuleFactory ruleFactory{
      mean_field::quadrature::Policy(mean_field::quadrature::make_rule_set(
          mean_field::quadrature::Mode::production))};

  const mean_field::quadrature::MfemRule pressureIntegralRule =
      ruleFactory.get(pressureIntegralQuery, mfem::Geometry::CUBE);

  const mean_field::quadrature::MfemRule pressureForceRule =
      ruleFactory.get(pressureForceQuery, mfem::Geometry::CUBE);

  /*
   * On the reference cube [0,1]^3 choose
   *
   *     h = x^3 y^3 z^3.
   *
   * This is representable by the order-three H1 space.
   * Then
   *
   *     P = x^12 y^12 z^12 / 4.
   */
  const double numericalPressureIntegral =
      polytropic_eos_test_utils::integrate_cube(
          *pressureIntegralRule.integration_rule,
          [&barotrope](const mfem::IntegrationPoint &integrationPoint) {
            const double coordinateProduct =
                integrationPoint.x * integrationPoint.y * integrationPoint.z;

            const double enthalpy = std::pow(coordinateProduct, 3.0);

            return barotrope.pressure_from_enthalpy(enthalpy);
          });

  const double analyticPressureIntegral = 0.25 / std::pow(13.0, 3.0);

  /*
   * Choose a representable vector test function whose
   * divergence is
   *
   *     div(w) = x^2 y^2 z^2.
   *
   * Therefore
   *
   *     -P div(w)
   *       = -x^14 y^14 z^14 / 4.
   */
  const double numericalPressureForceIntegral =
      polytropic_eos_test_utils::integrate_cube(
          *pressureForceRule.integration_rule,
          [&barotrope](const mfem::IntegrationPoint &integrationPoint) {
            const double coordinateProduct =
                integrationPoint.x * integrationPoint.y * integrationPoint.z;

            const double enthalpy = std::pow(coordinateProduct, 3.0);

            const double pressure = barotrope.pressure_from_enthalpy(enthalpy);

            const double testDivergence =
                integrationPoint.x * integrationPoint.x * integrationPoint.y *
                integrationPoint.y * integrationPoint.z * integrationPoint.z;

            return -pressure * testDivergence;
          });

  const double analyticPressureForceIntegral = -0.25 / std::pow(15.0, 3.0);

  INFO("Pressure-integral quadrature order = "
       << pressureIntegralRule.resolution.order);

  INFO("Pressure-force quadrature order = "
       << pressureForceRule.resolution.order);

  INFO("Numerical pressure integral = " << numericalPressureIntegral);

  INFO("Analytic pressure integral = " << analyticPressureIntegral);

  INFO(
      "Numerical pressure-force integral = " << numericalPressureForceIntegral);

  INFO("Analytic pressure-force integral = " << analyticPressureForceIntegral);

  CHECK(pressureIntegralRule.resolution.base_order == 12);

  CHECK(pressureIntegralRule.resolution.order == 12);

  CHECK(pressureForceRule.resolution.base_order == 14);

  CHECK(pressureForceRule.resolution.order == 14);

  CHECK_THAT(numericalPressureIntegral,
             Catch::Matchers::WithinAbs(analyticPressureIntegral, 5.0e-14));

  CHECK_THAT(
      numericalPressureForceIntegral,
      Catch::Matchers::WithinAbs(analyticPressureForceIntegral, 5.0e-14));
}