#include <cmath>
#include <limits>
#include <memory>

#include "serif/tests/equations/gravity_port.hpp"
#include "serif/equations/mass_normalization.hpp"
#include "serif/equations/angular_momentum.hpp"
#include "serif/equations/primal_gravity.hpp"
#include "serif/equations/barotropic_density.hpp"
#include "serif/backend/mfem/operators/volume_field_block.hpp"
#include "serif/tests/nonconforming.hpp"

namespace {
    using namespace serif;
    using namespace serif::tests::equations;
    using namespace serif::tests::equations::gravity;
    using Vector = ::mfem::Vector;
    using Mode = backend::PreparationMode;

    struct Values {
        Vector rho, h, phi, c, omega, u;
        /// Smooth values covering every migrated equation's unknowns.
        explicit Values(GravityFixture& f) : rho(f.density()), h(f.enthalpy()), phi(f.potential()),
            c(constant(1, 0.7)), omega(constant(1, 0.4)), u(f.displacement(0.04)) {}
        /// Bind each physical term once, independently of which equation consumes it.
        auto state(GravityFixture& f) const {
            return backend::make_term_coefficients(backend::bind(f.density_binding, rho), backend::bind(f.enthalpy_binding, h),
                backend::bind(f.potential_binding, phi),
                backend::bind(f.constant_binding, c), backend::bind(f.angular_velocity_binding, omega));
        }
    };

    /// Use the prepared plan's public column order to construct an independent packed direction.
    template <typename OperatorT, typename StateT>
    Vector pack(const OperatorT& op, const StateT& state, const Vector& displacement) {
        Vector result(op.width());
        int offset = 0;
        OperatorT::Plan::for_each_column([&]<typename OperandT>(OperandT, std::size_t) {
            const auto& values = state.template get<typename OperandT::term_type>().values;
            std::copy_n(values.GetData(), values.Size(), result.GetData() + offset);
            offset += values.Size();
        });
        std::copy_n(displacement.GetData(), displacement.Size(), result.GetData() + offset);
        return result;
    }

    /// Check full field/geometry transposes, dense/tensor agreement, freezing and the MFEM adapter.
    template <typename EquationT, typename BindingT>
    void check_operator(GravityFixture& f, EquationT equation, const BindingT& row) {
        CAPTURE(EquationT::name, f.mesh.GetElementGeometry(0), f.displacement_space.GetOrdering());
        Values values(f), direction(f);
        direction.rho = deterministic(values.rho.Size(), 0.3);
        direction.h = deterministic(values.h.Size(), 0.5);
        direction.phi = deterministic(values.phi.Size(), 0.9);
        direction.c[0] = -0.23;
        direction.omega[0] = 0.17;
        direction.u = f.displacement(-0.07);
        VolumeEquationFixture automatic{equation, row, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
        VolumeEquationFixture dense{equation, row, f.displacement_binding, f.compactification, f.mapper, reference_policy(), backend::ContractionMode::dense};
        automatic.prepare(values.state(f), values.u, f.workspace);
        dense.prepare(values.state(f), values.u, f.workspace);
        CHECK(dense.tensor_element_count() == 0);
        CHECK((automatic.tensor_element_count() > 0) == (f.mesh.GetElementGeometry(0) == ::mfem::Geometry::CUBE));
        Vector residual, expected_residual, action, expected_action, transpose, expected_transpose;
        automatic.residual(f.workspace, residual);
        dense.residual(f.workspace, expected_residual);
        CHECK(relative_error(residual, expected_residual) < 2.0e-12);
        const auto input = pack(automatic, direction.state(f), direction.u);
        automatic.apply_jacobian(input, f.workspace, action);
        dense.apply_jacobian(input, f.workspace, expected_action);
        CHECK(relative_error(action, expected_action) < 2.0e-12);
        auto dual = deterministic(automatic.height(), -0.8);
        automatic.apply_jacobian_transpose(dual, f.workspace, transpose);
        dense.apply_jacobian_transpose(dual, f.workspace, expected_transpose);
        CHECK(relative_error(transpose, expected_transpose) < 2.0e-12);
        // The old mass-adapter transpose pairing threshold is retained for every equation.
        CHECK(gravity_port::pairing_relative(input * transpose, action * dual) < 2.0e-12);
        backend::JacobianOperator adapter{automatic.cache()};
        Vector adapted;
        adapter.MultTranspose(dual, adapted);
        CHECK(relative_error(adapted, transpose) < 2.0e-15);
        Vector zero = constant(dual.Size(), 0.0);
        automatic.apply_jacobian_transpose(zero, f.workspace, adapted);
        CHECK(adapted.Norml2() == 0.0);

        // Probe both ends of each coefficient block and the geometry block independently.
        int offset = 0;
        const auto check_column = [&](int index) {
            Vector basis = constant(automatic.width(), 0.0), column;
            basis[index] = 1.0;
            automatic.apply_jacobian(basis, f.workspace, column);
            CHECK(std::abs(column * dual - transpose[index]) < 2.0e-12 * std::max(1.0, std::abs(transpose[index])));
        };
        decltype(automatic)::Plan::for_each_column([&]<typename OperandT>(OperandT, std::size_t) {
            const auto& input_values = values.state(f).template get<typename OperandT::term_type>().values;
            check_column(offset);
            check_column(offset + input_values.Size() - 1);
            offset += input_values.Size();
        });
        check_column(offset);
        check_column(automatic.width() - 1);
        values.rho = values.h = values.phi = 0.0;
        values.c = values.omega = values.u = 0.0;
        automatic.apply_jacobian_transpose(dual, f.workspace, adapted);
        CHECK(relative_error(adapted, transpose) < 2.0e-15);
        Vector wrong_size(automatic.height() + 1);
        CHECK_THROWS_AS(automatic.apply_jacobian_transpose(wrong_size, f.workspace, adapted), std::invalid_argument);
        dual[0] = std::numeric_limits<double>::quiet_NaN();
        CHECK_THROWS_AS(automatic.apply_jacobian_transpose(dual, f.workspace, adapted), std::domain_error);
    }

    /** @brief Ensure a material-only equation actually encounters weighted input or test coefficients. */
    void require_material_hanging_coefficient(const discretization::backend::mfem::DofMap& binding,
        const discretization::backend::mfem::DofMap& material) {
        const ::mfem::SparseMatrix* prolongation = binding.space().GetConformingProlongation();
        REQUIRE(prolongation != nullptr);
        bool found = false;
        ::mfem::Array<int> dofs;
        for (int element = 0; element < binding.mesh().GetNE(); ++element) {
            if (!material.supports(element)) {
                continue;
            }
            binding.space().GetElementVDofs(element, dofs);
            for (const int encoded : dofs) {
                const int local = ::mfem::FiniteElementSpace::DecodeDof(encoded);
                int contributions = 0;
                for (int entry = prolongation->GetI()[local]; entry < prolongation->GetI()[local + 1]; ++entry) {
                    if (prolongation->GetData()[entry] != 0.0) {
                        ++contributions;
                    }
                }
                found = found || contributions > 1;
            }
        }
        REQUIRE(found);
    }

    /** @brief Isolate the production stellar source without assembling vacuum diffusion. */
    struct StellarDensitySource : serif::equations::PrimalPoissonBalance::DensitySource {
        static constexpr std::string_view name = "stellar density source";
        /** @brief Preserve the primal Poisson source coefficient and sign. */
        explicit StellarDensitySource(const double coefficient) : DensitySource(coefficient) {}
    };

    struct CountingClosure {
        using Test = Closure::Test;
        using Operands = Closure::Operands;
        static constexpr std::string_view name = "counting closure";
        Closure equation{eos::models::Polytrope{1.5, 0.25}};
        std::shared_ptr<int> calls = std::make_shared<int>(0);
        /// Count local evaluations without changing the physical residual or derivative.
        auto evaluate(const serif::equations::EquationPoint<Operands>& point) const { ++*calls; return equation.evaluate(point); }
        /// Use the real closure's quadrature contract.
        template <typename OperandT> int additional_order(OperandT operand, int order) const { return equation.additional_order(operand, order); }
    };
}

TEST_CASE("Every migrated volume equation has a complete serial Jacobian transpose", "[equations][serial_backend][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);
            f.set_compactification();
            const eos::models::Polytrope eos{1.5, 0.25};
            check_operator(f, Closure{eos}, f.density_binding);
            check_operator(f, PotentialRow{four_pi_G}, f.potential_binding);
            check_operator(f, serif::equations::BarotropicDensityBalance{eos}, f.density_binding);
            check_operator(f, serif::equations::MassNormalization{0.8}, f.constant_binding);
            check_operator(f, serif::equations::AngularMomentumConstraint{0.3}, f.angular_velocity_binding);
        }
    }
}

TEST_CASE("Curved compactified geometry preserves transpose pairings", "[equations][serial_backend][integration]") {
    auto fixture = gravity_port::stellar_fixture();
    auto& f = *fixture;
    check_operator(f, PotentialRow{four_pi_G}, f.potential_binding);
    const eos::models::Polytrope eos{1.5, 0.25};
    check_operator(f, serif::equations::BarotropicDensityBalance{eos}, f.density_binding);
}

TEST_CASE("Selective preparation refreshes only changed operand samples and geometry", "[equations][serial_backend][unit]") {
    GravityFixture f;
    f.set_compactification();
    CountingClosure equation;
    VolumeEquationFixture op{equation, f.density_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    Values values(f);
    op.prepare(values.state(f), values.u, f.workspace);
    const auto initial = op.preparation_report();
    REQUIRE(initial.physics_points > 0);
    CHECK(initial.physics_points == static_cast<std::size_t>(*equation.calls));
    const auto initial_calls = *equation.calls;
    Vector before, after;
    op.residual(f.workspace, before);
    op.prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged);
    CHECK_FALSE(op.preparation_report().did_work());
    CHECK(op.preparation_report().reused_elements > 0);
    CHECK(*equation.calls == initial_calls);
    op.residual(f.workspace, after);
    CHECK(relative_error(after, before) == 0.0);

    values.rho += 0.03;
    op.prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged);
    CHECK(op.preparation_report().geometry_points == 0);
    CHECK(op.preparation_report().sampled_operands * 2 == initial.sampled_operands);
    CHECK(op.preparation_report().physics_points == initial.physics_points);
    CHECK(*equation.calls == 2 * initial_calls);
    op.residual(f.workspace, after);
    op.prepare(values.state(f), values.u, f.workspace);
    op.residual(f.workspace, before);
    CHECK(relative_error(after, before) < 2.0e-15);

    values.u *= 1.01;
    op.prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged);
    CHECK(op.preparation_report().geometry_points == initial.geometry_points);
    CHECK(op.preparation_report().sampled_operands == 0);
    f.mesh.NodesUpdated();
    CHECK_THROWS_AS(op.residual(f.workspace, after), std::logic_error);
    op.prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged);
    CHECK(op.preparation_report().geometry_points == initial.geometry_points);
    CHECK(op.preparation_report().sampled_operands == initial.sampled_operands);

    values.h = -1.0;
    CHECK_FALSE(op.try_prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged).has_value());
    CHECK_FALSE(op.is_prepared());
    CHECK_THROWS_AS(op.residual(f.workspace, after), std::logic_error);
    values.h = f.enthalpy();
    op.prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged);
    CHECK(op.preparation_report().geometry_points == initial.geometry_points);
    CHECK(op.preparation_report().sampled_operands == initial.sampled_operands);
    op.invalidate();
    CHECK_FALSE(op.preparation_report().did_work());
    op.prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged);
    CHECK(op.preparation_report().did_work());
}

TEST_CASE("Selective preparation follows contribution dependencies and copied coefficient values", "[equations][serial_backend][unit]") {
    GravityFixture f;
    Values values(f);
    // Only the stellar density source reads density; the potential flux on every domain is reused.
    VolumeEquationFixture op{serif::equations::PrimalPoissonBalance{four_pi_G}, f.potential_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    op.prepare(values.state(f), values.u, f.workspace);
    const auto initial = op.preparation_report();
    auto copied = values;
    op.prepare(copied.state(f), copied.u, f.workspace, Mode::reuse_unchanged);
    CHECK_FALSE(op.preparation_report().did_work());
    const auto elements = op.preparation_report().reused_elements;
    copied.rho *= 1.1;
    op.prepare(copied.state(f), copied.u, f.workspace, Mode::reuse_unchanged);
    CHECK(op.preparation_report().geometry_points == 0);
    CHECK(op.preparation_report().reused_elements > 0);
    CHECK(op.preparation_report().reused_elements < elements);
    CHECK(op.preparation_report().sampled_operands > 0);
    CHECK(op.preparation_report().sampled_operands < initial.sampled_operands);
    CHECK(op.preparation_report().physics_points > 0);
    CHECK(op.preparation_report().physics_points < initial.physics_points);
    Vector incremental, refreshed;
    op.residual(f.workspace, incremental);
    op.prepare(copied.state(f), copied.u, f.workspace);
    op.residual(f.workspace, refreshed);
    CHECK(relative_error(incremental, refreshed) < 2.0e-15);
}

TEST_CASE("Selective preparation refreshes compactification without resampling unchanged operands", "[equations][serial_backend][unit]") {
    GravityFixture f;
    f.set_compactification();
    Values values(f);
    VolumeEquationFixture op{PotentialRow{four_pi_G}, f.potential_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    op.prepare(values.state(f), values.u, f.workspace);
    const auto initial = op.preparation_report();
    Vector frozen, unchanged, incremental, refreshed, action, transpose;
    op.residual(f.workspace, frozen);
    f.compactification *= 0.97;
    op.residual(f.workspace, unchanged);
    CHECK(relative_error(unchanged, frozen) == 0.0);

    op.prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged);
    CHECK(op.preparation_report().geometry_points == initial.geometry_points);
    CHECK(op.preparation_report().sampled_operands == 0);
    op.residual(f.workspace, incremental);
    CHECK(relative_error(incremental, frozen) > 1.0e-8);
    const auto direction = pack(op, values.state(f), values.u);
    const auto dual = deterministic(op.height(), 0.31);
    op.apply_jacobian(direction, f.workspace, action);
    op.apply_jacobian_transpose(dual, f.workspace, transpose);

    op.prepare(values.state(f), values.u, f.workspace);
    op.residual(f.workspace, refreshed);
    CHECK(relative_error(incremental, refreshed) < 2.0e-12);
    op.apply_jacobian(direction, f.workspace, refreshed);
    CHECK(relative_error(action, refreshed) < 2.0e-12);
    op.apply_jacobian_transpose(dual, f.workspace, refreshed);
    CHECK(relative_error(transpose, refreshed) < 2.0e-12);
}

TEST_CASE("Dependency preparation detects coefficient binding replacement and ignores unused state roles", "[equations][serial_backend][unit]") {
    GravityFixture f;
    Values values(f);
    CountingClosure equation;
    VolumeEquationFixture op{equation, f.density_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    op.prepare(values.state(f), values.u, f.workspace);
    const auto initial = op.preparation_report();
    const auto calls = *equation.calls;

    values.omega[0] *= 1.2; // Rotation is present in the state but is not an input to closure.
    op.prepare(values.state(f), values.u, f.workspace, Mode::reuse_unchanged);
    CHECK_FALSE(op.preparation_report().did_work());
    CHECK(*equation.calls == calls);

    discretization::backend::mfem::TermBinding<Density, Schema> replacement{f.density_space};
    const auto rebound = backend::make_term_coefficients(backend::bind(replacement, values.rho), backend::bind(f.enthalpy_binding, values.h));
    op.prepare(rebound, values.u, f.workspace, Mode::reuse_unchanged);
    CHECK(op.preparation_report().geometry_points == initial.geometry_points);
    CHECK(op.preparation_report().sampled_operands == initial.sampled_operands);
    CHECK(*equation.calls == 2 * calls);

    Vector incremental, rebuilt;
    op.residual(f.workspace, incremental);
    op.prepare(rebound, values.u, f.workspace);
    op.residual(f.workspace, rebuilt);
    CHECK(relative_error(incremental, rebuilt) < 2.0e-15);
}

TEST_CASE("STROID hanging nodes preserve complete volume actions and sparse field blocks", "[integration][serial_backend][nonconforming]") {
    const auto generated = serif::tests::nonconforming::mesh(1, 2);
    for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
        GravityFixture fixture(::mfem::Mesh(*generated.mesh), ordering);
        ::mfem::GridFunctionCoefficient coordinate(generated.exterior_coordinate->values.get());
        fixture.compactification.ProjectCoefficient(coordinate);
        REQUIRE(fixture.potential_space.GetTrueVSize() < fixture.potential_space.GetVSize());
        check_operator(fixture, PotentialRow{four_pi_G}, fixture.potential_binding);
        Values values(fixture);
        VolumeEquationFixture equation{PotentialRow{four_pi_G}, fixture.potential_binding, fixture.displacement_binding,
            fixture.compactification, fixture.mapper, reference_policy()};
        equation.prepare(values.state(fixture), values.u, fixture.workspace);
        const auto matrix = backend::assemble_volume_field<Potential>(equation.data());
        const auto diagonal = backend::volume_field_diagonal<Potential>(equation.data());
        const Vector input = deterministic(fixture.potential_binding.size(), 0.7);
        Vector actual, expected(matrix->Height()), native_diagonal;
        auto direction = gravity_port::zeros(fixture);
        direction.potential = input;
        equation.apply_jacobian(fixture.state(direction.potential, direction.density), direction.displacement, fixture.workspace, actual);
        matrix->Mult(input, expected);
        CHECK(relative_error(actual, expected) < 3.0e-12);
        matrix->GetDiag(native_diagonal);
        for (int index = 0; index < native_diagonal.Size(); ++index) {
            CHECK(diagonal[index] == native_diagonal[index]);
        }
    }
}

TEST_CASE("STROID material hanging nodes preserve nonlinear density-balance and density-source actions", "[integration][serial_backend][nonconforming]") {
    const auto generated = serif::tests::nonconforming::mesh(0, 3);
    GravityFixture fixture(::mfem::Mesh(*generated.mesh), ::mfem::Ordering::byVDIM);
    ::mfem::GridFunctionCoefficient coordinate(generated.exterior_coordinate->values.get());
    fixture.compactification.ProjectCoefficient(coordinate);
    require_material_hanging_coefficient(fixture.density_binding, fixture.density_binding);
    require_material_hanging_coefficient(fixture.displacement_binding, fixture.density_binding);

    const serif::equations::BarotropicDensityBalance density_balance{eos::models::Polytrope{1.5, 0.25}};
    check_operator(fixture, density_balance, fixture.density_binding);
    check_operator(fixture, StellarDensitySource{four_pi_G}, fixture.potential_binding);

    // Positive interior density keeps this nonlinear tangent away from the vacuum singularity.
    Values values(fixture), direction(fixture);
    direction.rho = deterministic(values.rho.Size(), 0.3);
    direction.phi = 0.0;
    direction.c = 0.0;
    direction.omega = 0.0;
    direction.u = fixture.displacement(0.001);
    VolumeEquationFixture density_equation{density_balance, fixture.density_binding, fixture.displacement_binding,
        fixture.compactification, fixture.mapper, reference_policy()};
    density_equation.prepare(values.state(fixture), values.u, fixture.workspace);
    const auto density_matrix = backend::assemble_volume_field<blocks::density::mass>(density_equation.data());
    const Vector zero_displacement = constant(values.u.Size(), 0.0);
    Vector actual, expected(density_matrix->Height());
    density_equation.apply_jacobian(direction.state(fixture), zero_displacement, fixture.workspace, actual);
    density_matrix->Mult(direction.rho, expected);
    CHECK(relative_error(actual, expected) < 3.0e-12);
    const Vector density_dual = deterministic(density_matrix->Height(), -0.8);
    Vector complete_transpose, sparse_transpose(density_matrix->Width());
    density_equation.apply_jacobian_transpose(density_dual, fixture.workspace, complete_transpose);
    density_matrix->MultTranspose(density_dual, sparse_transpose);
    const Vector field_transpose(complete_transpose.GetData(), direction.rho.Size());
    CHECK(relative_error(sparse_transpose, field_transpose) < 3.0e-12);

    // The nonlinear residual independently checks the constrained field and geometry tangent together.
    density_equation.apply_jacobian(direction.state(fixture), direction.u, fixture.workspace, actual);
    constexpr double step = 1.0e-6;
    Values plus(values), minus(values);
    plus.rho.Add(step, direction.rho);
    minus.rho.Add(-step, direction.rho);
    plus.u.Add(step, direction.u);
    minus.u.Add(-step, direction.u);
    Vector positive, negative;
    density_equation.prepare(plus.state(fixture), plus.u, fixture.workspace);
    density_equation.residual(fixture.workspace, positive);
    density_equation.prepare(minus.state(fixture), minus.u, fixture.workspace);
    density_equation.residual(fixture.workspace, negative);
    positive -= negative;
    positive /= 2.0 * step;
    CHECK(relative_error(actual, positive) < 4.0e-8);

    VolumeEquationFixture source_equation{StellarDensitySource{four_pi_G}, fixture.potential_binding,
        fixture.displacement_binding, fixture.compactification, fixture.mapper, reference_policy()};
    source_equation.prepare(values.state(fixture), values.u, fixture.workspace);
    const auto source_matrix = backend::assemble_volume_field<blocks::density::mass>(source_equation.data());
    ReferenceGeometry geometry{fixture, fixture.mapper, values.u};
    ReferenceVolume volume{geometry};
    ::mfem::MixedBilinearForm native_source(&fixture.density_space, &fixture.potential_space);
    auto* integrator = new ::mfem::MixedScalarMassIntegrator(volume);
    const auto& rule = ::mfem::IntRules.Get(fixture.mesh.GetElementBaseGeometry(0), reference_order);
    integrator->SetIntRule(&rule);
    ::mfem::Array<int> stellar{1, 1, 0};
    native_source.AddDomainIntegrator(integrator, stellar);
    native_source.Assemble();
    native_source.Finalize();

    // Native local assembly and native P transpose retain all hanging-node cross terms.
    Vector native_full(fixture.potential_binding.full_size()), native_true(fixture.potential_binding.true_size());
    native_source.Mult(fixture.density_binding.scatter(values.rho), native_full);
    native_full *= four_pi_G;
    fixture.potential_space.GetConformingProlongation()->MultTranspose(native_full, native_true);
    const Vector native_residual = fixture.potential_binding.gather_true(native_true);
    source_equation.residual(fixture.workspace, actual);
    CHECK(relative_error(actual, native_residual) < 3.0e-12);
    expected.SetSize(source_matrix->Height());
    source_matrix->Mult(values.rho, expected);
    CHECK(relative_error(expected, native_residual) < 3.0e-12);

    const Vector source_dual = deterministic(source_matrix->Height(), 0.7);
    native_full.SetSize(fixture.density_binding.full_size());
    native_true.SetSize(fixture.density_binding.true_size());
    native_source.MultTranspose(fixture.potential_binding.scatter(source_dual), native_full);
    native_full *= four_pi_G;
    fixture.density_space.GetConformingProlongation()->MultTranspose(native_full, native_true);
    const Vector native_transpose = fixture.density_binding.gather_true(native_true);
    sparse_transpose.SetSize(source_matrix->Width());
    source_matrix->MultTranspose(source_dual, sparse_transpose);
    CHECK(relative_error(sparse_transpose, native_transpose) < 3.0e-12);
}