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

#include "serif/equations/primal_gravity.hpp"
#include "serif/equations/registry.hpp"
#include "serif/tests/equations/fixture.hpp"
#include "serif/tests/test_tags.hpp"

namespace serif::tests::equations {
    // An external-style EOS needs only the advertised relation and derivative.
    struct LinearEOS final {
        using Relations = eos::relations::RelationCatalog<eos::relations::DensityFromSpecificEnthalpy>;
        std::shared_ptr<int> evaluations;

        dimensions::DensityValue evaluate(
            eos::relations::DensityFromSpecificEnthalpy,
            dimensions::SpecificEnthalpyValue h
        ) const {
            ++*evaluations;
            return dimensions::DensityValue{2.0 * h.value()};
        }

        dimensions::PartialDerivative<dimensions::MassDensity, dimensions::SpecificEnthalpy> partial_derivative(
            eos::relations::DensityFromSpecificEnthalpy,
            dimensions::WithRespectTo<dimensions::SpecificEnthalpy>,
            dimensions::SpecificEnthalpyValue
        ) const {
            return dimensions::PartialDerivative<dimensions::MassDensity, dimensions::SpecificEnthalpy>{2.0};
        }
    };

    TEST_CASE("Equation Closure Accepts Compatible EOS Without Backend Specialization", tags::unit & tags::equations) {
        Fixture f;
        const auto calls = std::make_shared<int>(0);
        const auto equation = ::serif::tests::equations::BarotropicClosure{LinearEOS{calls}, 1.0};
        VolumeEquationFixture prepared(
            equation, f.density_binding, f.displacement_binding, f.compactification, f.mapper
        );
        const auto rho = constant(f.density_binding.size(), 1.6), h = constant(f.enthalpy_binding.size(), 0.8);
        const auto u = f.displacement(0.8);
        prepared.prepare(f.state(rho, h), u, f.workspace);
        const int preparation_calls = *calls;
        REQUIRE(preparation_calls > 0);
        ::mfem::Vector residual, action;
        prepared.residual(f.workspace, residual);
        prepared.apply_jacobian(f.state(rho, h), u, f.workspace, action);
        CHECK(residual.Norml2() < 1.0e-13);
        CHECK(action.Norml2() < 1.0e-13);
        CHECK(*calls == preparation_calls);
    }

    TEST_CASE("Equation Closure Local Derivatives Match The Analytic Polytrope", tags::unit & tags::equations) {
        for (const double n : {1.0, 1.5, 3.0}) {
            Closure equation{eos::models::Polytrope{n, 1.5}};
            serif::equations::EquationPoint<Closure::Operands> point;
            point.fields.get<blocks::density::mass>() = 0.55;
            point.fields.get<::serif::tests::equations::Enthalpy>() = 0.9;
            const auto result = equation.evaluate(point);
            const double scale = (n + 1.0) * 1.5;
            CHECK_THAT(result.residual, Catch::Matchers::WithinRel(0.55 - std::pow(0.9 / scale, n), 1.0e-14));
            CHECK(result.field_derivatives.get<blocks::density::mass>() == 1.0);
            CHECK_THAT(
                result.field_derivatives.get<::serif::tests::equations::Enthalpy>(),
                Catch::Matchers::WithinRel(-n / scale * std::pow(0.9 / scale, n - 1.0), 1.0e-14)
            );
            CHECK(result.position_derivative == std::array<double, 3>{});
        }
    }

    TEST_CASE("Equation Registry Defines Rows From Its Equations", tags::unit & tags::equations) {
        const serif::equations::EquationRegistry registry{
            Closure{eos::models::Polytrope{1.0, 0.8}}, serif::equations::PrimalPoissonBalance{1.0}
        };
        int count = 0;
        registry.for_each([&](const auto& equation) {
            CHECK_FALSE(equation.name.empty());
            ++count;
        });
        CHECK(count == 2);
        CHECK(registry.get<Closure>().name == "barotropic closure");
        CHECK_THROWS_AS(
            (::serif::tests::equations::BarotropicClosure{LinearEOS{std::make_shared<int>(0)}, -1.0}), std::invalid_argument
        );
        CHECK_THROWS_AS(
            (serif::equations::PrimalPoissonBalance{std::numeric_limits<double>::infinity()}),
            std::invalid_argument
        );
    }
} // namespace serif::tests::equations