#include <cmath>
#include <limits>

#include "serif/equations/mass_normalization.hpp"
#include "serif/tests/equations/gravity_port.hpp"

namespace {
    using namespace serif;
    using namespace serif::tests::equations;
    using namespace serif::tests::equations::gravity;
    using Equation = serif::equations::MassNormalization;
    using Prepared = VolumeEquationFixture<Equation, Schema>;
    using Vector = ::mfem::Vector;

    auto state(GravityFixture& f, const Vector& density) {
        return backend::make_term_coefficients(backend::bind(f.density_binding, density));
    }

    auto make_operator(GravityFixture& f, double target) {
        return Prepared{Equation{target}, f.constant_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    }

    double residual(GravityFixture& f, Prepared& op, const Vector& density, const Vector& displacement) {
        op.prepare(state(f, density), displacement, f.workspace);
        Vector result;
        op.residual(f.workspace, result);
        REQUIRE(result.Size() == 1);
        return result[0];
    }

    double action(GravityFixture& f, Prepared& op, const Vector& density, const Vector& displacement) {
        Vector result;
        op.apply_jacobian(state(f, density), displacement, f.workspace, result);
        REQUIRE(result.Size() == 1);
        return result[0];
    }

    // MeanField prepared_mass_normalization.cpp:66,98,132. Preserve the original profiles.
    Vector density(GravityFixture& f, double phase, bool direction = false) {
        return f.project(f.density_binding, [=](const Vector& x) {
            if (direction) {
                return 0.19 * std::sin(0.71 * x[0] + phase) - 0.13 * std::cos(0.89 * x[1] - phase) + 0.08 * x[2];
            }
            return 0.91 + 0.07 * std::sin(0.83 * x[0] + phase) + 0.05 * std::cos(0.61 * x[1] - phase) + 0.03 * x[2] * x[2];
        });
    }

    Vector displacement(GravityFixture& f, double scale) {
        return gravity_port::project_vector(f.displacement_binding, [=](const Vector& x, Vector& value) {
            value.SetSize(3);
            value[0] = scale * (0.07 * x[0] + 0.018 * x[1] * x[2]);
            value[1] = scale * (-0.05 * x[1] + 0.013 * x[0] * x[2]);
            value[2] = scale * (0.04 * x[2] - 0.011 * x[0] * x[1]);
        });
    }

    double relative(double actual, double expected) {
        return std::abs(actual - expected) / std::max(std::abs(expected), 100.0 * std::numeric_limits<double>::epsilon());
    }

    // Independent MFEM quadrature of rho det(A) and delta-rho det(A) + rho cof(A):delta-A.
    double reference(GravityFixture& f, const Vector& rho, const Vector& u, const Vector* drho = nullptr, const Vector* du = nullptr) {
        ::mfem::GridFunction density(&f.density_space), displacement(&f.displacement_space);
        ::mfem::GridFunction delta_density(&f.density_space), delta_displacement(&f.displacement_space);
        density = f.density_binding.scatter(rho);
        displacement = f.displacement_binding.scatter(u);
        if (drho) { delta_density = f.density_binding.scatter(*drho); }
        if (du) { delta_displacement = f.displacement_binding.scatter(*du); }
        double result = 0.0;
        for (int element = 0; element < f.mesh.GetNE(); ++element) {
            if (!f.density_binding.supports(element)) { continue; }
            auto& transform = *f.mesh.GetElementTransformation(element);
            const auto& rule = ::mfem::IntRules.Get(f.mesh.GetElementGeometry(element), reference_order);
            for (int point = 0; point < rule.GetNPoints(); ++point) {
                const auto& ip = rule.IntPoint(point);
                transform.SetIntPoint(&ip);
                ::mfem::DenseMatrix A(3), dA(3), inverse(3);
                displacement.GetVectorGradient(transform, A);
                for (int axis = 0; axis < 3; ++axis) { A(axis, axis) += 1.0; }
                double integrand = density.GetValue(transform, ip);
                if (drho || du) {
                    integrand = drho ? delta_density.GetValue(transform, ip) : 0.0;
                    if (du) {
                        delta_displacement.GetVectorGradient(transform, dA);
                        ::mfem::CalcInverse(A, inverse);
                        double trace = 0.0;
                        for (int i = 0; i < 3; ++i) {
                            for (int j = 0; j < 3; ++j) { trace += inverse(i, j) * dA(j, i); }
                        }
                        integrand += density.GetValue(transform, ip) * trace;
                    }
                }
                result += ip.weight * transform.Weight() * A.Det() * integrand;
            }
        }
        return result;
    }

    void check_columns(GravityFixture& f) {
        CAPTURE(f.mesh.GetNE(), f.mesh.GetElementGeometry(0), f.displacement_space.GetOrdering());
        f.set_compactification();
        const auto zero_rho = constant(f.density_binding.size(), 0.0), zero_u = constant(f.displacement_binding.size(), 0.0);
        // Original density and geometry finite differences, with their original steps and tolerances.
        for (int column = 0; column < 2; ++column) {
            const auto rho = density(f, column == 0 ? 0.31 : 0.37);
            const auto u = displacement(f, column == 0 ? 0.43 : 0.51);
            const auto drho = column == 0 ? density(f, 0.67, true) : zero_rho;
            const auto du = column == 0 ? zero_u : displacement(f, -0.79);
            auto op = make_operator(f, column == 0 ? 1.23 : 1.11);
            (void)residual(f, op, rho, u);
            const double actual = action(f, op, drho, du);
            const double step = column == 0 ? 1.0e-3 : 1.0e-6;
            auto plus_rho = rho, minus_rho = rho, plus_u = u, minus_u = u;
            plus_rho.Add(step, drho); minus_rho.Add(-step, drho);
            plus_u.Add(step, du); minus_u.Add(-step, du);
            const double difference = (residual(f, op, plus_rho, plus_u) - residual(f, op, minus_rho, minus_u)) / (2.0 * step);
            CAPTURE(column, actual, difference);
            CHECK(relative(actual, difference) < (column == 0 ? 3.0e-8 : 2.0e-7));
        }

        // Original complete action, independent reference, column sum, and supported local adapter routing.
        const auto rho = density(f, 0.47), u = displacement(f, 0.57);
        const auto drho = density(f, 0.71, true), du = displacement(f, -0.63);
        auto op = make_operator(f, 1.19);
        const auto initial = residual(f, op, rho, u);
        CHECK(relative(initial, reference(f, rho, u) - 1.19) < 1.0e-12);
        const double complete = action(f, op, drho, du);
        CHECK(relative(complete, reference(f, rho, u, &drho, &du)) < 2.0e-11);
        CHECK(relative(complete, action(f, op, drho, zero_u) + action(f, op, zero_rho, du)) < 2.0e-15);
        backend::JacobianOperator adapter(op.cache());
        CHECK(adapter.Height() == 1);
        CHECK(adapter.Width() == drho.Size() + du.Size());
        Vector joined(adapter.Width()), result;
        std::copy_n(drho.GetData(), drho.Size(), joined.GetData());
        std::copy_n(du.GetData(), du.Size(), joined.GetData() + drho.Size());
        adapter.Mult(joined, result);
        CHECK(result[0] == complete);
    }
}

TEST_CASE("Mass normalization retains affine volume scaling and target subtraction", "[equations][mass_normalization][integration]") {
    auto fixture = gravity_port::stellar_fixture();
    auto& f = *fixture;
    const auto rho = constant(f.density_binding.size(), 1.37), zero = constant(f.displacement_binding.size(), 0.0);
    auto op = make_operator(f, 0.73);
    const double undeformed = residual(f, op, rho, zero) + 0.73;
    CHECK(relative(undeformed, reference(f, rho, zero)) < 1.0e-12);
    CHECK(relative(residual(f, op, rho, zero), undeformed - 0.73) < 2.0e-15);
    const double deformed = residual(f, op, rho, gravity_port::affine_displacement(f, 0.086)) + 0.73;
    CHECK(relative(deformed / undeformed, std::pow(1.0 + 0.086, 3)) < 5.0e-7);
}

TEST_CASE("Mass normalization retains density geometry and complete action accuracy", "[equations][mass_normalization][integration]") {
    for (const auto geometry : {::mfem::Element::TETRAHEDRON, ::mfem::Element::HEXAHEDRON}) {
        for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
            GravityFixture f(geometry, ordering);
            check_columns(f);
        }
    }
    auto fixture = gravity_port::stellar_fixture();
    check_columns(*fixture);
}

TEST_CASE("Global mass preparations own their density state", "[equations][mass_normalization][unit]") {
    GravityFixture f;
    auto rho = density(f, 0.47);
    const auto u = displacement(f, 0.57);
    VolumeEquationFixture op{Equation{0.73}, f.constant_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    const double expected = reference(f, rho, u) - 0.73;
    op.prepare(state(f, rho), u, f.workspace);
    Vector result;
    op.residual(f.workspace, result);
    CHECK(relative(result[0], expected) < 2.0e-15);
    rho *= 1.2;
    op.residual(f.workspace, result);
    CHECK(relative(result[0], expected) < 2.0e-15);
    op.prepare(state(f, rho), u, f.workspace);
    op.residual(f.workspace, result);
    CHECK(relative(result[0], reference(f, rho, u) - 0.73) < 2.0e-15);
    const auto drho = density(f, 0.71, true), du = displacement(f, -0.63);
    op.apply_jacobian(state(f, drho), du, f.workspace, result);
    CHECK(relative(result[0], reference(f, rho, u, &drho, &du)) < 2.0e-11);
}

TEST_CASE("Mass normalization rejects invalid targets density and geometry then recovers", "[equations][mass_normalization][unit]") {
    for (double target : {0.0, -1.0, std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()}) {
        CHECK_THROWS_AS(Equation{target}, std::invalid_argument);
    }
    GravityFixture f;
    auto op = make_operator(f, 1.0);
    auto rho = constant(f.density_binding.size(), 1.0);
    auto u = constant(f.displacement_binding.size(), 0.0);
    const double expected = residual(f, op, rho, u);
    rho = std::numeric_limits<double>::quiet_NaN();
    auto rejected = op.try_prepare(state(f, rho), u, f.workspace);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().reason == backend::PreparationRejectionReason::NON_FINITE);
    CHECK_FALSE(op.is_prepared());
    CHECK_THROWS_AS(op.prepare(state(f, rho), u, f.workspace), std::domain_error);
    rho = 1.0;
    u = gravity_port::affine_displacement(f, -2.0);
    rejected = op.try_prepare(state(f, rho), u, f.workspace);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().reason == backend::PreparationRejectionReason::MAPPING);
    u = 0.0;
    CHECK(residual(f, op, rho, u) == expected);
}

// MeanField prepared_mass_normalization.cpp:282: finite coefficients can overflow during interpolation.
TEST_CASE("Mass normalization rejects nonfinite density interpolation and recovers", "[equations][mass_normalization][unit]") {
    auto fixture = gravity_port::stellar_fixture();
    auto& f = *fixture;
    auto rho = constant(f.density_binding.size(), -0.5 * std::numeric_limits<double>::max());
    const auto u = constant(f.displacement_binding.size(), 0.0);
    auto op = make_operator(f, std::numeric_limits<double>::max());
    const auto rejected = op.try_prepare(state(f, rho), u, f.workspace);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().reason == backend::PreparationRejectionReason::NON_FINITE);
    CHECK(rejected.error().element >= 0);
    CHECK_FALSE(op.is_prepared());
    CHECK_THROWS_AS(op.prepare(state(f, rho), u, f.workspace), std::domain_error);
    rho = 1.0;
    auto recovered = make_operator(f, 1.0);
    REQUIRE(recovered.try_prepare(state(f, rho), u, f.workspace).has_value());
    CHECK(recovered.is_prepared());
}

TEST_CASE("Global scalar bindings retain one coefficient and reject stale meshes", "[equations][mass_normalization][unit]") {
    GravityFixture f;
    auto& binding = f.constant_binding;
    const auto coefficient = constant(1, 0.39);
    CHECK(binding.size() == 1);
    CHECK(binding.full_size() == 1);
    CHECK(binding.gather(binding.scatter(coefficient))[0] == coefficient[0]);
    CHECK_THROWS_AS(binding.space(), std::logic_error);
    CHECK_THROWS_AS(binding.validate_values(constant(2, 0.0)), std::invalid_argument);
    for (int element = 0; element < f.mesh.GetNE(); ++element) {
        Vector local;
        binding.element_values(element, coefficient, local);
        REQUIRE(local.Size() == 1);
        CHECK(local[0] == (binding.supports(element) ? coefficient[0] : 0.0));
    }
    f.mesh.GetElement(0)->SetAttribute(3);
    CHECK_THROWS(binding.validate());
}