// EOS/geometry conditions and tolerances from MeanField's prepared closure tests.

#include <algorithm>
#include <array>
#include <cmath>
#include <catch2/catch_test_macros.hpp>

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

namespace serif::tests::equations {
    struct ClosureCondition final {
        double n, K, enthalpy_offset, enthalpy_gradient, density_factor, density_offset, density_gradient;
        double deformation, direction_phase;
    };

    constexpr std::array<ClosureCondition, 3> closure_conditions{
        {{1.0, 0.8, 0.65, 0.06, 0.80, 0.015, 0.004, 0.0, 0.31}, {1.5, 1.2, 0.90, 0.09, 1.15, -0.003, 0.003, 0.45, 0.53},
            {3.0, 1.5, 1.20, 0.12, 1.40, 0.006, 0.002, 1.0, 0.79}}
    };

    TEST_CASE("Equation Closure Covers Multiple EOS And Geometry Conditions", tags::integration & tags::equations) {
        Fixture f;
        for (const auto& condition : closure_conditions) {
            CAPTURE(condition.n, condition.K, condition.deformation);
            const eos::models::Polytrope eos(condition.n, condition.K);
            const auto enthalpy_function = [&](const ::mfem::Vector& x) {
                return condition.enthalpy_offset +
                       condition.enthalpy_gradient * (0.50 * x[0] - 0.30 * x[1] + 0.20 * x[2]);
            };
            const auto h = f.project(f.enthalpy_binding, enthalpy_function);
            const auto rho = f.project(f.density_binding, [&](const ::mfem::Vector& x) {
                return condition.density_factor * std::pow(enthalpy_function(x) / eos.enthalpy_scale(), condition.n) +
                       condition.density_offset +
                       condition.density_gradient * (0.40 * x[0] + 0.25 * x[1] - 0.15 * x[2]);
            });
            const auto u = f.displacement(condition.deformation);
            auto drho = deterministic(rho.Size(), condition.direction_phase);
            auto dh = f.project(f.enthalpy_binding, [](const ::mfem::Vector& x) {
                return 0.07 + 0.015 * x[0] + 0.008 * x[2];
            });
            auto du = f.displacement(0.63);
            auto density_action = manual_closure(f, eos, rho, h, u, &drho);
            auto enthalpy_action = manual_closure(f, eos, rho, h, u, nullptr, &dh);
            auto geometry_action = manual_closure(f, eos, rho, h, u, nullptr, nullptr, &du);
            const double target =
                std::min({density_action.Norml2(), enthalpy_action.Norml2(), geometry_action.Norml2()});
            REQUIRE(target > 1.0e-12);
            drho *= target / density_action.Norml2();
            dh *= target / enthalpy_action.Norml2();
            du *= target / geometry_action.Norml2();
            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, action;
            prepared.residual(f.workspace, residual);
            prepared.apply_jacobian(f.state(drho, dh), du, f.workspace, action);
            CHECK(relative_error(residual, manual_closure(f, eos, rho, h, u)) < 5.0e-12);
            CHECK(relative_error(action, manual_closure(f, eos, rho, h, u, &drho, &dh, &du)) < 5.0e-12);
            const int expected_order = std::max(4, 5 + static_cast<int>(std::ceil((condition.n - 1.0) * 3))) + 6;
            CHECK(prepared.integration_order(0) >= expected_order);
            constexpr double step = 1.0e-5;
            auto plus_rho = rho, minus_rho = rho, plus_h = h, minus_h = h, plus_u = u, minus_u = u;
            plus_rho.Add(step, drho);
            minus_rho.Add(-step, drho);
            plus_h.Add(step, dh);
            minus_h.Add(-step, dh);
            plus_u.Add(step, du);
            minus_u.Add(-step, du);
            auto difference = manual_closure(f, eos, plus_rho, plus_h, plus_u);
            difference -= manual_closure(f, eos, minus_rho, minus_h, minus_u);
            difference /= 2.0 * step;
            difference -= action;
            CHECK(difference.Norml2() / (3.0 * target) < 2.0e-7);
        }
    }

    TEST_CASE("Exact Equation Closure Remains Zero Under Deformation", tags::integration & tags::equations) {
        Fixture f;
        const auto h = constant(f.enthalpy_binding.size(), 1.20);
        const auto rho = constant(f.density_binding.size(), std::pow(1.20 / 6.0, 3.0));
        auto reference_rho = rho;
        reference_rho += 1.0;
        const auto zero_rho = constant(rho.Size(), 0.0), zero_h = constant(h.Size(), 0.0), du = f.displacement(0.67);
        for (const double deformation : {0.0, 0.5, 1.0}) {
            CAPTURE(deformation);
            const auto u = f.displacement(deformation);
            PreparedClosure prepared(
                Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
                f.mapper
            );
            ::mfem::Vector exact, exact_geometry, reference, reference_geometry;
            prepared.prepare(f.state(rho, h), u, f.workspace);
            prepared.residual(f.workspace, exact);
            prepared.apply_jacobian(f.state(zero_rho, zero_h), du, f.workspace, exact_geometry);
            prepared.prepare(f.state(reference_rho, h), u, f.workspace);
            prepared.residual(f.workspace, reference);
            prepared.apply_jacobian(f.state(zero_rho, zero_h), du, f.workspace, reference_geometry);
            REQUIRE(reference.Norml2() > 1.0e-12);
            REQUIRE(reference_geometry.Norml2() > 1.0e-14);
            CHECK(exact.Norml2() <= 5.0e-12 * reference.Norml2());
            CHECK(exact_geometry.Norml2() <= 5.0e-12 * reference_geometry.Norml2());
        }
    }

    TEST_CASE("Equation Quadrature Resolves Both Closure Forms Through The Existing Policy", tags::unit & tags::equations) {
        Fixture f;
        const auto rho = f.density(), h = f.enthalpy(), u = f.displacement(0.0);
        using Mass = discretization::forms::mass<blocks::density::mass, blocks::density::mass>;
        using Source = discretization::forms::source<::serif::tests::equations::Enthalpy, blocks::density::mass>;
        discretization::quadrature::RuleSet rules;
        rules.set_form_control<Mass>({.fixed_order = 18});
        rules.set_form_control<Source>({.fixed_order = 8});
        PreparedClosure prepared(
            Closure{eos::models::Polytrope{1.0, 0.8}}, f.density_binding, f.displacement_binding, f.compactification,
            f.mapper, discretization::quadrature::Policy{rules}
        );
        prepared.prepare(f.state(rho, h), u, f.workspace);
        CHECK(prepared.integration_order(0) == ::mfem::IntRules.Get(::mfem::Geometry::CUBE, 18).GetOrder());
        CHECK_THROWS_AS(prepared.integration_order(2), std::out_of_range);
    }
} // namespace serif::tests::equations