// Numerical contracts ported from MeanField tests/operators/kernels/barotropic_closure_kernels.cpp.

#include <memory>
#include <catch2/catch_test_macros.hpp>

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

namespace serif::tests::equations {
    TEST_CASE("Equation Closure Vanishes For A Representable Constant State", tags::unit & tags::equations) {
        Fixture f;
        const eos::models::Polytrope eos(3.0, 1.5);
        const double h_value = 0.8;
        const auto h = constant(f.enthalpy_binding.size(), h_value);
        const auto rho = constant(f.density_binding.size(), std::pow(h_value / 6.0, 3.0));
        const auto u = f.displacement(0.0);
        PreparedClosure prepared(Closure{eos}, f.density_binding, f.displacement_binding, f.compactification, f.mapper);
        prepared.prepare(f.state(rho, h), u, f.workspace);
        ::mfem::Vector residual, scale;
        prepared.residual(f.workspace, residual);
        const auto zero_h = constant(h.Size(), 0.0);
        prepared.apply_jacobian(f.state(rho, zero_h), u, f.workspace, scale);
        REQUIRE(scale.Norml2() > 0.0);
        CHECK(residual.Norml2() / scale.Norml2() < 5.0e-12);
    }

    TEST_CASE("Equation Closure Density Action Matches The Stellar Mass Matrix", tags::unit & tags::equations) {
        for (const auto geometry : {::mfem::Element::HEXAHEDRON, ::mfem::Element::TETRAHEDRON}) {
            CAPTURE(geometry);
            Fixture f(geometry);
            const auto rho = f.density();
            const auto h = f.enthalpy();
            const auto zero_h = constant(h.Size(), 0.0);
            const auto u = f.displacement(0.0);
            const auto drho = deterministic(rho.Size(), 0.37);
            PreparedClosure prepared(
                Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
                f.mapper
            );
            prepared.prepare(f.state(rho, h), u, f.workspace);
            ::mfem::Vector action;
            prepared.apply_jacobian(f.state(drho, zero_h), u, f.workspace, action);
            ::mfem::Array<int> marker(3);
            marker = 0;
            marker[0] = marker[1] = 1;
            ::mfem::BilinearForm mass(&f.density_space);
            mass.AddDomainIntegrator(new ::mfem::MassIntegrator(), marker);
            mass.Assemble();
            mass.Finalize();
            const auto full_direction = f.density_binding.scatter(drho);
            ::mfem::Vector full_action(full_direction.Size());
            mass.Mult(full_direction, full_action);
            CHECK(relative_error(action, f.density_binding.gather(full_action)) < 5.0e-12);
        }
    }

    TEST_CASE("Equation Closure Thermodynamic Jacobian Matches Centered Differences", tags::unit & tags::equations) {
        Fixture f;
        const eos::models::Polytrope eos(3.0, 1.5);
        const auto rho = f.density();
        const auto h = f.enthalpy();
        const auto u = f.displacement(1.0);
        const auto drho = deterministic(rho.Size(), 0.63);
        const auto dh = f.project(f.enthalpy_binding, [](const ::mfem::Vector& x) {
            return 0.07 + 0.015 * x[0] + 0.008 * x[2];
        });
        const auto zero_u = constant(u.Size(), 0.0);
        PreparedClosure prepared(Closure{eos}, f.density_binding, f.displacement_binding, f.compactification, f.mapper);
        prepared.prepare(f.state(rho, h), u, f.workspace);
        ::mfem::Vector analytic;
        prepared.apply_jacobian(f.state(drho, dh), zero_u, f.workspace, analytic);
        constexpr double step = 1.0e-6;
        auto plus_rho = rho, minus_rho = rho, plus_h = h, minus_h = h;
        plus_rho.Add(step, drho);
        minus_rho.Add(-step, drho);
        plus_h.Add(step, dh);
        minus_h.Add(-step, dh);
        auto difference = manual_closure(f, eos, plus_rho, plus_h, u);
        difference -= manual_closure(f, eos, minus_rho, minus_h, u);
        difference /= 2.0 * step;
        CHECK(relative_error(analytic, difference) < 2.0e-8);
    }

    TEST_CASE("Equation Closure Excludes Vacuum And Uses Mapped Volume", tags::unit & tags::equations) {
        Fixture f;
        const auto rho = f.density(), h = f.enthalpy();
        const auto zero_h = constant(h.Size(), 0.0);
        const auto zero_u = f.displacement(0.0);
        const auto ones = constant(rho.Size(), 1.0);
        ::mfem::Vector vacuum(f.density_space.GetVSize());
        vacuum = 1.0;
        vacuum -= f.density_binding.scatter(ones);
        const auto vacuum_direction = f.density_binding.gather(vacuum);
        CHECK(vacuum.Norml2() > 0.0);
        ::mfem::Vector identity, deformed, vacuum_action;
        PreparedClosure prepared(
            Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
            f.mapper
        );
        prepared.prepare(f.state(rho, h), zero_u, f.workspace);
        prepared.apply_jacobian(f.state(ones, zero_h), zero_u, f.workspace, identity);
        prepared.apply_jacobian(f.state(vacuum_direction, zero_h), zero_u, f.workspace, vacuum_action);
        prepared.prepare(f.state(rho, h), f.displacement(1.0), f.workspace);
        prepared.apply_jacobian(f.state(ones, zero_h), zero_u, f.workspace, deformed);
        REQUIRE(identity.Norml2() > 0.0);
        CHECK(vacuum_action.Norml2() <= 1.0e-13 * identity.Norml2());
        CHECK(relative_error(deformed, identity) > 1.0e-5);
    }

    TEST_CASE("Equation Closure Geometry Jacobian Matches Centered Differences", tags::integration & tags::equations) {
        for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
            Fixture f(::mfem::Element::HEXAHEDRON, ordering);
            const eos::models::Polytrope eos(3.0, 1.5);
            const auto rho = f.density(), h = f.enthalpy(), du = f.displacement(0.65);
            const auto zero_rho = constant(rho.Size(), 0.0), zero_h = constant(h.Size(), 0.0);
            for (const double deformation : {0.0, 1.0}) {
                CAPTURE(ordering, deformation);
                const auto u = f.displacement(deformation);
                PreparedClosure prepared(
                    Closure{eos}, f.density_binding, f.displacement_binding, f.compactification, f.mapper
                );
                prepared.prepare(f.state(rho, h), u, f.workspace);
                ::mfem::Vector analytic;
                prepared.apply_jacobian(f.state(zero_rho, zero_h), du, f.workspace, analytic);
                auto plus = u, minus = u;
                constexpr double step = 1.0e-5;
                plus.Add(step, du);
                minus.Add(-step, du);
                auto difference = manual_closure(f, eos, rho, h, plus);
                difference -= manual_closure(f, eos, rho, h, minus);
                difference /= 2.0 * step;
                REQUIRE(analytic.Norml2() > 1.0e-12);
                CHECK(relative_error(analytic, difference) < 5.0e-8);
            }
        }
    }

    TEST_CASE("Equation Closure Geometry Action Is Linear And Zero Outside Stellar Support", tags::unit & tags::equations) {
        Fixture f;
        const auto rho = f.density(), h = f.enthalpy(), u = f.displacement(0.8);
        const auto zero_rho = constant(rho.Size(), 0.0), zero_h = constant(h.Size(), 0.0),
                   zero_u = constant(u.Size(), 0.0);
        const auto first = f.displacement(0.4), second = deterministic(u.Size(), 0.91);
        auto combined = first;
        combined *= 1.7;
        combined.Add(-0.43, second);
        PreparedClosure prepared(
            Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
            f.mapper
        );
        prepared.prepare(f.state(rho, h), u, f.workspace);
        ::mfem::Vector a, b, sum, zero;
        prepared.apply_jacobian(f.state(zero_rho, zero_h), first, f.workspace, a);
        prepared.apply_jacobian(f.state(zero_rho, zero_h), second, f.workspace, b);
        prepared.apply_jacobian(f.state(zero_rho, zero_h), combined, f.workspace, sum);
        prepared.apply_jacobian(f.state(zero_rho, zero_h), zero_u, f.workspace, zero);
        a *= 1.7;
        a.Add(-0.43, b);
        CHECK(relative_error(sum, a) < 5.0e-12);
        CHECK(zero.Norml2() == 0.0);
        auto vacuum_direction = second;
        ::mfem::Array<int> dofs;
        for (int element = 0; element < f.mesh.GetNE(); ++element) {
            if (f.mesh.GetAttribute(element) != 3) {
                f.displacement_space.GetElementVDofs(element, dofs);
                for (int dof : dofs) {
                    vacuum_direction[dof < 0 ? -1 - dof : dof] = 0.0;
                }
            }
        }
        REQUIRE(vacuum_direction.Norml2() > 0.0);
        prepared.apply_jacobian(f.state(zero_rho, zero_h), vacuum_direction, f.workspace, zero);
        CHECK(zero.Norml2() == 0.0);
    }
} // namespace serif::tests::equations