#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 eos = mean_field::eos;

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 =
                    eos::evaluate<eos::quantity::Pressure>(barotrope, eos::DensityValue{density}).value();

                const double enthalpyFromDensity =
                    eos::evaluate<eos::quantity::SpecificEnthalpy>(barotrope, eos::DensityValue{density}).value();

                const double recoveredDensity =
                    eos::evaluate<eos::quantity::Density>(barotrope, eos::SpecificEnthalpyValue{enthalpyFromDensity})
                        .value();

                const double pressureFromEnthalpy =
                    eos::evaluate<eos::quantity::Pressure>(barotrope, eos::SpecificEnthalpyValue{enthalpyFromDensity})
                        .value();

                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 the density from specific enthalpy relation.
                 */
                CHECK(
                    (eos::partialDerivative<eos::quantity::Pressure, eos::quantity::SpecificEnthalpy>(
                         barotrope, eos::SpecificEnthalpyValue{enthalpyFromDensity}
                     )
                         .value() == recoveredDensity)
                );

                /*
                 * Since
                 *
                 *     h = (n + 1) K rho^(1/n),
                 *
                 * it follows that
                 *
                 *     dP / d rho = h / n.
                 */
                CHECK_THAT(
                    (eos::partialDerivative<eos::quantity::Pressure, eos::quantity::Density>(
                         barotrope, eos::DensityValue{density}
                    )
                         .value()),
                    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 eos::evaluate<eos::quantity::Pressure>(
                                   barotrope, eos::SpecificEnthalpyValue{perturbedEnthalpy}
                        )
                            .value();
                    },
                    enthalpy, step
                );

                const double analyticDerivative =
                    eos::partialDerivative<eos::quantity::Pressure, eos::quantity::SpecificEnthalpy>(
                        barotrope, eos::SpecificEnthalpyValue{enthalpy}
                    )
                        .value();

                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 eos::evaluate<eos::quantity::Pressure>(barotrope, eos::DensityValue{perturbedDensity})
                            .value();
                    },
                    density, step
                );

                const double analyticDerivative =
                    eos::partialDerivative<eos::quantity::Pressure, eos::quantity::Density>(
                        barotrope, eos::DensityValue{density}
                    )
                        .value();

                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 eos::evaluate<eos::quantity::Density>(
                                   barotrope, eos::SpecificEnthalpyValue{perturbedEnthalpy}
                        )
                            .value();
                    },
                    enthalpy, step
                );

                const double analyticDerivative =
                    eos::partialDerivative<eos::quantity::Density, eos::quantity::SpecificEnthalpy>(
                        barotrope, eos::SpecificEnthalpyValue{enthalpy}
                    )
                        .value();

                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(eos::evaluate<eos::quantity::Density>(barotrope, eos::SpecificEnthalpyValue{0.0}).value() == 0.0);

            CHECK(eos::evaluate<eos::quantity::Pressure>(barotrope, eos::SpecificEnthalpyValue{0.0}).value() == 0.0);

            CHECK(
                (eos::partialDerivative<eos::quantity::Pressure, eos::quantity::SpecificEnthalpy>(
                     barotrope, eos::SpecificEnthalpyValue{0.0}
                 )
                     .value() == 0.0)
            );

            CHECK(eos::evaluate<eos::quantity::Pressure>(barotrope, eos::DensityValue{0.0}).value() == 0.0);

            CHECK(eos::evaluate<eos::quantity::SpecificEnthalpy>(barotrope, eos::DensityValue{0.0}).value() == 0.0);

            CHECK(
                (eos::partialDerivative<eos::quantity::Pressure, eos::quantity::Density>(
                     barotrope, eos::DensityValue{0.0}
                 )
                     .value() == 0.0)
            );

            /*
             * Positive-part extension into h < 0.
             */
            CHECK(
                eos::evaluate<eos::quantity::Density>(barotrope, eos::SpecificEnthalpyValue{exteriorEnthalpy})
                    .value() == 0.0
            );

            CHECK(
                eos::evaluate<eos::quantity::Pressure>(barotrope, eos::SpecificEnthalpyValue{exteriorEnthalpy})
                    .value() == 0.0
            );

            CHECK(
                (eos::partialDerivative<eos::quantity::Density, eos::quantity::SpecificEnthalpy>(
                     barotrope, eos::SpecificEnthalpyValue{exteriorEnthalpy}
                 )
                     .value() == 0.0)
            );

            CHECK(
                (eos::partialDerivative<eos::quantity::Pressure, eos::quantity::SpecificEnthalpy>(
                     barotrope, eos::SpecificEnthalpyValue{exteriorEnthalpy}
                 )
                     .value() == 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(
                (eos::partialDerivative<eos::quantity::Density, eos::quantity::SpecificEnthalpy>(
                     barotrope, eos::SpecificEnthalpyValue{0.0}
                 )
                     .value() == 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(eos::evaluate<eos::quantity::Pressure>(barotrope, eos::DensityValue{-0.1}), std::domain_error);

    CHECK_THROWS_AS(
        eos::evaluate<eos::quantity::SpecificEnthalpy>(barotrope, eos::DensityValue{-0.1}), std::domain_error
    );

    CHECK_THROWS_AS(
        (eos::partialDerivative<eos::quantity::Pressure, eos::quantity::Density>(barotrope, eos::DensityValue{-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(
            eos::evaluate<eos::quantity::Density>(barotrope, eos::SpecificEnthalpyValue{nonfiniteValue}),
            std::domain_error
        );

        CHECK_THROWS_AS(
            eos::evaluate<eos::quantity::Pressure>(barotrope, eos::SpecificEnthalpyValue{nonfiniteValue}),
            std::domain_error
        );

        CHECK_THROWS_AS(
            (eos::partialDerivative<eos::quantity::Density, eos::quantity::SpecificEnthalpy>(
                barotrope, eos::SpecificEnthalpyValue{nonfiniteValue}
            )),
            std::domain_error
        );

        CHECK_THROWS_AS(
            (eos::partialDerivative<eos::quantity::Pressure, eos::quantity::SpecificEnthalpy>(
                barotrope, eos::SpecificEnthalpyValue{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 eos::evaluate<eos::quantity::Pressure>(barotrope, eos::SpecificEnthalpyValue{enthalpy}).value();
        }
    );

    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 =
                eos::evaluate<eos::quantity::Pressure>(barotrope, eos::SpecificEnthalpyValue{enthalpy}).value();

            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));
}