// Prepared-state contracts ported from MeanField's prepared closure and context tests.

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

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

namespace serif::tests::equations {
    TEST_CASE("Prepared Equations Use Supported Field Dimensions", tags::unit & tags::equations) {
        Fixture f;
        PreparedClosure prepared(
            Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
            f.mapper
        );
        CHECK_FALSE(prepared.is_prepared());
        CHECK(prepared.preparation_count() == 0);
        CHECK(f.density_binding.size() > 0);
        CHECK(f.density_binding.size() < f.density_space.GetTrueVSize());
        CHECK(f.enthalpy_binding.size() < f.enthalpy_space.GetTrueVSize());
        CHECK(f.displacement_binding.size() == f.displacement_space.GetTrueVSize());
        CHECK(prepared.height() == f.density_binding.size());
        const auto rho = f.density(), h = f.enthalpy(), u = f.displacement(0.0);
        prepared.prepare(f.state(rho, h), u, f.workspace);
        CHECK(prepared.width() == rho.Size() + h.Size() + u.Size());
        const auto expanded = f.density_binding.scatter(rho);
        CHECK(relative_error(f.density_binding.gather(expanded), rho) == 0.0);
    }

    TEST_CASE("Prepared Equations Reject EOS Candidates And Recover", tags::unit & tags::equations) {
        Fixture f;
        const auto rho = f.density(), u = f.displacement(0.0);
        auto h = f.enthalpy();
        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);
        for (const double bad_value : {-1.0, 1.0e150}) {
            CAPTURE(bad_value);
            h = bad_value;
            const auto rejected = prepared.try_prepare(f.state(rho, h), u, f.workspace);
            REQUIRE_FALSE(rejected.has_value());
            CHECK(rejected.error().reason == backend::PreparationRejectionReason::EQUATION);
            CHECK(rejected.error().element >= 0);
            CHECK(rejected.error().quadrature_point >= 0);
            CHECK_FALSE(prepared.is_prepared());
            ::mfem::Vector result;
            CHECK_THROWS_AS(prepared.residual(f.workspace, result), std::logic_error);
            CHECK_THROWS_AS(prepared.prepare(f.state(rho, h), u, f.workspace), std::domain_error);
        }
        h = 1.0;
        REQUIRE(prepared.try_prepare(f.state(rho, h), u, f.workspace).has_value());
        CHECK(prepared.preparation_count() == 2);
    }

    TEST_CASE("Prepared Equations Reject Folding Geometry And Recover", tags::unit & tags::equations) {
        Fixture f;
        const auto rho = f.density(), h = f.enthalpy();
        ::mfem::GridFunction folded(&f.displacement_space);
        ::mfem::VectorFunctionCoefficient coefficient(3, [](const ::mfem::Vector& x, ::mfem::Vector& u) {
            u.SetSize(3);
            u = 0.0;
            u[0] = -2.0 * x[0];
        });
        folded.ProjectCoefficient(coefficient);
        auto u = f.displacement_binding.gather(folded);
        PreparedClosure prepared(
            Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
            f.mapper
        );
        const auto rejected = prepared.try_prepare(f.state(rho, h), u, f.workspace);
        REQUIRE_FALSE(rejected.has_value());
        CHECK(rejected.error().reason == backend::PreparationRejectionReason::MAPPING);
        CHECK(rejected.error().mapping_status == mapping::MappingStatus::NON_POSITIVE_DETERMINATE);
        CHECK_FALSE(prepared.is_prepared());
        CHECK_THROWS_AS(prepared.prepare(f.state(rho, h), u, f.workspace), std::domain_error);
        u = 0.0;
        CHECK(prepared.try_prepare(f.state(rho, h), u, f.workspace).has_value());
    }

    TEST_CASE("Prepared Closure Matches Manual Residual And Every Jacobian Column", tags::integration & tags::equations) {
        for (const auto geometry : {::mfem::Element::HEXAHEDRON, ::mfem::Element::TETRAHEDRON}) {
            CAPTURE(geometry);
            Fixture f(geometry);
            const eos::models::Polytrope eos(3.0, 1.5);
            const auto rho = f.density(), h = f.enthalpy(), u = f.displacement(1.0);
            const auto drho = deterministic(rho.Size(), 0.43), dh = deterministic(h.Size(), 0.52),
                       du = f.displacement(0.63);
            const auto zero_rho = constant(rho.Size(), 0.0), zero_h = constant(h.Size(), 0.0),
                       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 residual, density_action, enthalpy_action, geometry_action, complete, packed_action;
            prepared.residual(f.workspace, residual);
            prepared.apply_jacobian(f.state(drho, zero_h), zero_u, f.workspace, density_action);
            prepared.apply_jacobian(f.state(zero_rho, dh), zero_u, f.workspace, enthalpy_action);
            prepared.apply_jacobian(f.state(zero_rho, zero_h), du, f.workspace, geometry_action);
            prepared.apply_jacobian(f.state(drho, dh), du, f.workspace, complete);
            backend::JacobianOperator jacobian(prepared.cache());
            jacobian.Mult(packed(drho, dh, du), packed_action);
            CHECK(relative_error(residual, manual_closure(f, eos, rho, h, u)) < 2.0e-12);
            CHECK(relative_error(density_action, manual_closure(f, eos, rho, h, u, &drho)) < 2.0e-12);
            CHECK(relative_error(enthalpy_action, manual_closure(f, eos, rho, h, u, nullptr, &dh)) < 2.0e-12);
            CHECK(relative_error(geometry_action, manual_closure(f, eos, rho, h, u, nullptr, nullptr, &du)) < 2.0e-12);
            density_action += enthalpy_action;
            density_action += geometry_action;
            CHECK(relative_error(complete, density_action) < 2.0e-12);
            CHECK(relative_error(packed_action, complete) < 2.0e-12);
            CHECK(prepared.preparation_count() == 1);
        }
    }

    TEST_CASE("Prepared Equations Own Frozen Numerical State And Reprepare Every Dependency", tags::unit & tags::equations) {
        Fixture f;
        const eos::models::Polytrope eos(3.0, 1.5);
        auto rho = f.density(), h = f.enthalpy(), u = f.displacement(0.35);
        const auto drho = deterministic(rho.Size(), 0.12), dh = deterministic(h.Size(), 0.53),
                   du = f.displacement(0.63);
        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 frozen, frozen_action;
        prepared.residual(f.workspace, frozen);
        prepared.apply_jacobian(f.state(drho, dh), du, f.workspace, frozen_action);
        for (int dependency = 0; dependency < 3; ++dependency) {
            CAPTURE(dependency);
            if (dependency == 0) {
                rho = f.density(1.0);
            } else if (dependency == 1) {
                h = f.enthalpy(1.0);
            } else {
                u = f.displacement(0.85);
            }
            ::mfem::Vector unchanged, unchanged_action;
            prepared.residual(f.workspace, unchanged);
            prepared.apply_jacobian(f.state(drho, dh), du, f.workspace, unchanged_action);
            CHECK(relative_error(unchanged, frozen) == 0.0);
            CHECK(relative_error(unchanged_action, frozen_action) == 0.0);
            prepared.prepare(f.state(rho, h), u, f.workspace);
            prepared.residual(f.workspace, frozen);
            prepared.apply_jacobian(f.state(drho, dh), du, f.workspace, frozen_action);
            CHECK(relative_error(frozen, manual_closure(f, eos, rho, h, u)) < 2.0e-12);
            CHECK(relative_error(frozen_action, manual_closure(f, eos, rho, h, u, &drho, &dh, &du)) < 2.0e-12);
        }
        CHECK(prepared.preparation_count() == 4);
        prepared.invalidate();
        CHECK_FALSE(prepared.is_prepared());
        CHECK_THROWS_AS(prepared.residual(f.workspace, frozen), std::logic_error);
    }
} // namespace serif::tests::equations