#include <cmath>

#include <catch2/catch_test_macros.hpp>

#include "serif/equations/barotropic_density.hpp"
#include "serif/tests/equations/barotropic_eos.hpp"
#include "serif/tests/equations/gravity_fixture.hpp"

namespace {
    using Balance = serif::equations::BarotropicDensityBalance<serif::tests::equations::MixedPowerBarotrope>;
}

TEST_CASE("Barotropic density assembles a nonpolytropic value balance and complete ALE tangent", "[equations][barotropic_density][unit]") {
    using namespace serif::tests::equations;

    for (const auto geometry : {::mfem::Element::HEXAHEDRON, ::mfem::Element::TETRAHEDRON}) {
        gravity::GravityFixture fixture(geometry);
        auto density = fixture.density(), potential = fixture.potential();
        auto surface_level = constant(1, -0.8), omega = constant(1, 0.2), displacement = fixture.displacement(0.3);
        /** @brief Bind exactly the four physical inputs; the EOS supplies every thermodynamic value. */
        const auto fields = [&](const auto& rho, const auto& phi, const auto& level, const auto& rotation) {
            return backend::make_term_coefficients(backend::bind(fixture.density_binding, rho),
                backend::bind(fixture.potential_binding, phi), backend::bind(fixture.constant_binding, level),
                backend::bind(fixture.angular_velocity_binding, rotation));
        };
        const Balance balance{MixedPowerBarotrope{.gauge = 7.5}, serif::dimensions::PressureValue{0.3125},
            {1.0, 2.0, 3.0}, {0.2, -0.1, 0.3}};
        VolumeEquationFixture equation{balance, fixture.density_binding, fixture.displacement_binding,
            fixture.compactification, fixture.mapper, gravity::reference_policy()};
        equation.prepare(fields(density, potential, surface_level, omega), constant(displacement.Size(), 0.0), fixture.workspace);
        ::mfem::Vector residual;
        equation.residual(fixture.workspace, residual);

        /** @brief Integrate the analytic physical residual independently of the equation's evaluator. */
        ::mfem::FunctionCoefficient expected_balance([](const ::mfem::Vector& position) {
            const double rho = 0.55 + 0.025 * position[0] - 0.010 * position[1] + 0.006 * position[2];
            const double phi = -0.5 + 0.1 * position[0] * position[1] - 0.07 * position[2] * position[2];
            const double x = position[0] - 0.2, y = position[1] + 0.1, z = position[2] - 0.3;
            const double axial_projection = x + 2.0 * y + 3.0 * z;
            const double centrifugal = 0.5 * (x * x + y * y + z * z - axial_projection * axial_projection / 14.0);
            // P_surface = 5/16 fixes rho_surface = 1/2 and A_surface - gauge = 7/6.
            return 2.0 * rho + (4.0 / 3.0) * rho * rho * rho - 7.0 / 6.0 + phi - 0.04 * centrifugal + 0.8;
        });
        ::mfem::LinearForm reference(&fixture.density_space);
        ::mfem::Array<int> stellar{1, 1, 0};
        const auto& rule = ::mfem::IntRules.Get(fixture.mesh.GetElementBaseGeometry(0), gravity::reference_order);
        auto* integrator = new ::mfem::DomainLFIntegrator(expected_balance);
        integrator->SetIntRule(&rule);
        reference.AddDomainIntegrator(integrator, stellar);
        reference.Assemble();
        CHECK(relative_error(residual, fixture.density_binding.gather(reference)) < 3.0e-12);

        const auto density_direction = deterministic(density.Size(), 0.1), potential_direction = deterministic(potential.Size(), 0.3);
        const auto level_direction = constant(1, -0.05), omega_direction = constant(1, 0.07), displacement_direction = fixture.displacement(0.2);
        equation.prepare(fields(density, potential, surface_level, omega), displacement, fixture.workspace);
        ::mfem::Vector action, transpose;
        equation.apply_jacobian(fields(density_direction, potential_direction, level_direction, omega_direction),
            displacement_direction, fixture.workspace, action);
        const auto dual = deterministic(residual.Size(), 0.7);
        equation.apply_jacobian_transpose(dual, fixture.workspace, transpose);
        const auto direction = packed(packed(density_direction, potential_direction, level_direction), omega_direction, displacement_direction);
        CHECK(std::abs(dual * action - direction * transpose) < 4.0e-11);

        constexpr double step = 1.0e-6;
        auto plus_density = density, minus_density = density, plus_potential = potential, minus_potential = potential;
        auto plus_level = surface_level, minus_level = surface_level, plus_omega = omega, minus_omega = omega;
        auto plus_displacement = displacement, minus_displacement = displacement;
        plus_density.Add(step, density_direction);
        minus_density.Add(-step, density_direction);
        plus_potential.Add(step, potential_direction);
        minus_potential.Add(-step, potential_direction);
        plus_level.Add(step, level_direction);
        minus_level.Add(-step, level_direction);
        plus_omega.Add(step, omega_direction);
        minus_omega.Add(-step, omega_direction);
        plus_displacement.Add(step, displacement_direction);
        minus_displacement.Add(-step, displacement_direction);
        ::mfem::Vector positive, negative;
        equation.prepare(fields(plus_density, plus_potential, plus_level, plus_omega), plus_displacement, fixture.workspace);
        equation.residual(fixture.workspace, positive);
        equation.prepare(fields(minus_density, minus_potential, minus_level, minus_omega), minus_displacement, fixture.workspace);
        equation.residual(fixture.workspace, negative);
        positive -= negative;
        positive /= 2.0 * step;
        CHECK(relative_error(action, positive) < 4.0e-8);
    }
}

TEST_CASE("Reconstructed density retains the field and geometry Jacobian including its transpose", "[equations][unit][reconstruction]") {
    using namespace serif::tests::equations;
    namespace reconstruction = serif::discretization::reconstruction;
    namespace discretization_backend = serif::discretization::backend::mfem;
    using Density = blocks::density::mass;
    using PolytropicBalance = serif::equations::BarotropicDensityBalance<serif::eos::models::Polytrope>;

    gravity::GravityFixture fixture(::mfem::Mesh::MakeCartesian3D(2, 1, 1, ::mfem::Element::HEXAHEDRON));
    const double exponent = 5.0 / 3.0;
    const auto reconstruction_plan = discretization_backend::make_reconstruction(
        reconstruction::SurfacePower<Density>{.density_exponent = exponent}, fixture.density_binding);
    discretization_backend::TermBinding<Density, Schema> density_binding(fixture.density_space, reconstruction_plan);
    auto density = fixture.project(density_binding, [exponent](const ::mfem::Vector& point) {
        return std::pow(point[2], exponent) * (1.0 + 0.1 * point[0]);
    });
    const auto density_direction = fixture.project(density_binding, [exponent](const ::mfem::Vector& point) {
        return 0.1 * std::pow(point[2], exponent) * (point[0] - 0.3 * point[1]);
    });
    const auto potential = fixture.potential(), level = constant(1, -0.8), omega = constant(1, 0.2);
    const auto displacement = fixture.displacement(0.3), displacement_direction = fixture.displacement(0.2);
    const auto fields = [&](const auto& rho) {
        return backend::make_term_coefficients(backend::bind(density_binding, rho),
            backend::bind(fixture.potential_binding, potential), backend::bind(fixture.constant_binding, level),
            backend::bind(fixture.angular_velocity_binding, omega));
    };
    VolumeEquationFixture equation{PolytropicBalance{serif::eos::models::Polytrope{exponent, 0.7}, serif::dimensions::PressureValue{0.0}},
        density_binding, fixture.displacement_binding, fixture.compactification, fixture.mapper, gravity::reference_policy()};
    equation.prepare(fields(density), displacement, fixture.workspace);
    CHECK(equation.integration_order(0) == ::mfem::IntRules.Get(::mfem::Geometry::CUBE, gravity::reference_order).GetOrder());
    ::mfem::Vector action, transpose;
    const auto potential_direction = constant(potential.Size(), 0.0), level_direction = constant(1, 0.0), omega_direction = constant(1, 0.0);
    equation.apply_jacobian(backend::make_term_coefficients(backend::bind(density_binding, density_direction),
        backend::bind(fixture.potential_binding, potential_direction), backend::bind(fixture.constant_binding, level_direction),
        backend::bind(fixture.angular_velocity_binding, omega_direction)), displacement_direction, fixture.workspace, action);
    const auto dual = deterministic(density_binding.size(), 0.7);
    equation.apply_jacobian_transpose(dual, fixture.workspace, transpose);
    const auto direction = packed(packed(density_direction, potential_direction, level_direction), omega_direction, displacement_direction);
    CHECK(std::abs(dual * action - direction * transpose) < 4e-11);
    for (const double step : {1e-5, 1e-6}) {
        auto plus_density = density, minus_density = density;
        auto plus_displacement = displacement, minus_displacement = displacement;
        plus_density.Add(step, density_direction);
        minus_density.Add(-step, density_direction);
        plus_displacement.Add(step, displacement_direction);
        minus_displacement.Add(-step, displacement_direction);
        ::mfem::Vector plus, minus;
        equation.prepare(fields(plus_density), plus_displacement, fixture.workspace);
        equation.residual(fixture.workspace, plus);
        equation.prepare(fields(minus_density), minus_displacement, fixture.workspace);
        equation.residual(fixture.workspace, minus);
        plus -= minus;
        plus /= 2.0 * step;
        CHECK(relative_error(action, plus) < 4e-8);
    }
}