#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <numeric>
#include <tuple>
#include <variant>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include "serif/backend/mfem/geometry/deformation/computations.hpp"
#include "serif/backend/mfem/model/volume_displacement.hpp"
#include "serif/backend/mfem/model/deformation.hpp"
#include "serif/tests/deformation/mesh.hpp"
#include "serif/tests/equations/fixture.hpp"

namespace
{
    namespace deformation = serif::geometry::deformation;
    namespace deformation_backend = deformation::backend::mfem;
    namespace equation_backend = serif::operators::backend::mfem;
    namespace equation_fixture = serif::tests::equations;
    namespace dependencies = serif::utils::dependencies;
    namespace blocks = serif::discretization::blocks;
    namespace stages = equation_backend::volume_keys;
    using Status = dependencies::PreparationStatus;
    using Closure = equation_fixture::Closure;
    using Model = deformation::DomainDeformation<deformation::SphericalHarmonicSurface>;
    using VolumeDiscretization = equation_backend::VolumeDiscretization<Closure::Test, equation_fixture::Schema>;
    using Failure = std::variant<deformation::DeformationFailure, equation_backend::PreparationRejection>;
    using DensityValue = serif::discretization::forms::Operand<blocks::density::mass>;

    /** @brief Exercise displacement support independently of the stellar domain vocabulary. */
    struct SupportedDisplacement : serif::discretization::backend::mfem::DofMap {
        /** @brief A selected element set gives an independent gather/scatter reference. */
        SupportedDisplacement(mfem::FiniteElementSpace& space, const mfem::Array<int>& support) : DofMap(space, support) {}
    };

    /** @brief Keep the collection and coefficient count unchanged while exposing incompatible local numbering. */
    struct ReorderedDisplacementSpace : mfem::FiniteElementSpace
    {
        using mfem::FiniteElementSpace::FiniteElementSpace;

        /** @brief Model an equation companion whose scalar coefficients have different local meanings. */
        void GetElementDofs(int element, mfem::Array<int>& dofs, mfem::DofTransformation& transformation) const override
        {
            mfem::FiniteElementSpace::GetElementDofs(element, dofs, transformation);
            std::swap(dofs[0], dofs[1]);
        }
    };

    /** @brief Reuse the volume equation's existing field bindings and physical mapper. */
    [[nodiscard]] VolumeDiscretization equation_discretization(const equation_fixture::Fixture& fixture)
    {
        return {.residual_binding = fixture.density_binding, .displacement_binding = fixture.displacement_binding, .compactification = fixture.compactification, .mapper = fixture.mapper};
    }

    /** @brief Supply only independent physical inputs; displacement is produced inside the combined graph. */
    [[nodiscard]] auto combined_inputs(const Model& model, const deformation_backend::ReferenceDiscretization& reference, const std::vector<double>& shape, const equation_fixture::Fixture& fixture, const mfem::Vector& density, const mfem::Vector& enthalpy)
    {
        const auto shape_inputs = deformation_backend::make_deformation_inputs(model, reference, shape);
        const auto equation_inputs = equation_backend::make_volume_inputs<Closure>(equation_discretization(fixture), fixture.state(density, enthalpy));
        return dependencies::make_dependency_context(
            shape_inputs.template resource<deformation::keys::ReferenceDiscretization>(),
            shape_inputs.template resource<deformation::keys::SurfaceDiscretization>(),
            shape_inputs.template resource<deformation::keys::ReferenceCoordinates>(),
            shape_inputs.template resource<deformation::keys::SurfaceModel>(),
            shape_inputs.template resource<deformation::keys::InteriorModel>(),
            shape_inputs.template resource<deformation::keys::VacuumModel>(),
            shape_inputs.template resource<deformation::keys::ShapeCoefficients>(),
            equation_inputs.template resource<equation_backend::MeshAndCoefficientBindings>(),
            equation_inputs.template resource<equation_backend::CompactificationCoefficients>(),
            equation_inputs.template resource<blocks::density::mass>(),
            equation_inputs.template resource<::serif::tests::equations::Enthalpy>());
    }

    /** @brief Compare accepted residuals and both linear actions with the established standalone equation path. */
    template <typename Cache>
    void check_against_volume_cache(const Cache& combined, equation_fixture::Fixture& fixture, const serif::eos::models::Polytrope& eos, const mfem::Vector& density, const mfem::Vector& enthalpy)
    {
        const mfem::Vector displacement(combined.template get<equation_backend::DisplacementCoefficients>().values);
        const auto discretization = equation_discretization(fixture);
        const auto inputs = equation_backend::make_volume_inputs<Closure>(discretization, fixture.state(density, enthalpy), displacement);
        auto independent = equation_backend::make_volume_equation_cache(Closure{eos}, discretization, inputs);
        REQUIRE(independent.try_prepare(inputs));
        mfem::Vector actual, expected;
        equation_backend::volume_residual(combined.template get<stages::Linearization>(), actual);
        equation_backend::volume_residual(independent.template get<stages::Linearization>(), expected);
        CHECK(equation_fixture::relative_error(actual, expected) < 2.0e-14);

        const mfem::Vector direction = equation_fixture::packed(fixture.density(0.7), fixture.enthalpy(0.2), fixture.displacement(0.01));
        equation_backend::apply_jacobian(combined.template get<stages::Linearization>(), direction, fixture.workspace, actual);
        equation_backend::apply_jacobian(independent.template get<stages::Linearization>(), direction, fixture.workspace, expected);
        CHECK(equation_fixture::relative_error(actual, expected) < 2.0e-14);
        const mfem::Vector dual = equation_fixture::deterministic(fixture.density_binding.size(), 0.3);
        equation_backend::apply_jacobian_transpose(combined.template get<stages::Linearization>(), dual, fixture.workspace, actual);
        equation_backend::apply_jacobian_transpose(independent.template get<stages::Linearization>(), dual, fixture.workspace, expected);
        CHECK(equation_fixture::relative_error(actual, expected) < 2.0e-14);
    }

    /** @brief Verify the complete harmonic-shape-to-equation tangent and its reverse virtual-work pairing. */
    template <typename Cache>
    void check_shape_equation_linearization(Cache& cache, const Model& model, const deformation_backend::ReferenceDiscretization& reference, std::vector<double>& shape, equation_fixture::Fixture& fixture, const mfem::Vector& density, const mfem::Vector& enthalpy)
    {
        const auto& accepted = cache.template get<deformation::keys::Linearization>();
        auto shape_workspace = deformation::make_deformation_workspace(*accepted.map);
        const std::array shape_direction{0.35, -0.17, 0.22, 0.09};
        std::vector<double> volume_direction(accepted.map->output_size());
        REQUIRE(deformation::apply_jacobian(accepted, shape_direction, volume_direction, shape_workspace));

        // This fixture uses the same complete component-block ordering on both sides of the interface.
        REQUIRE(fixture.displacement_space.GetOrdering() == mfem::Ordering::byNODES);
        mfem::Vector displacement_direction(fixture.displacement_binding.size());
        REQUIRE(volume_direction.size() == static_cast<std::size_t>(displacement_direction.Size()));
        std::ranges::copy(volume_direction, displacement_direction.GetData());
        const mfem::Vector zero_density = equation_fixture::constant(fixture.density_binding.size(), 0.0);
        const mfem::Vector zero_enthalpy = equation_fixture::constant(fixture.enthalpy_binding.size(), 0.0);
        const mfem::Vector direction = equation_fixture::packed(zero_density, zero_enthalpy, displacement_direction);
        mfem::Vector analytic, equation_transpose;
        equation_backend::apply_jacobian(cache.template get<stages::Linearization>(), direction, fixture.workspace, analytic);
        const mfem::Vector equation_dual = equation_fixture::deterministic(fixture.density_binding.size(), 0.3);
        equation_backend::apply_jacobian_transpose(cache.template get<stages::Linearization>(), equation_dual, fixture.workspace, equation_transpose);

        const int displacement_offset = zero_density.Size() + zero_enthalpy.Size();
        REQUIRE(equation_transpose.Size() == displacement_offset + displacement_direction.Size());
        const std::span<const double> displacement_force(equation_transpose.GetData() + displacement_offset, displacement_direction.Size());
        std::array<double, 4> shape_force{};
        REQUIRE(deformation::apply_jacobian_transpose(accepted, displacement_force, shape_force, shape_workspace));
        const double equation_work = analytic * equation_dual;
        const double shape_work = std::inner_product(shape_direction.begin(), shape_direction.end(), shape_force.begin(), 0.0);
        CHECK(std::abs(equation_work - shape_work) <= 2.0e-12 * std::max({1.0, std::abs(equation_work), std::abs(shape_work)}));

        const std::vector<double> accepted_shape = shape;
        constexpr double step = 1.0e-6;
        for (std::size_t parameter = 0; parameter < shape.size(); ++parameter)
        {
            shape[parameter] = accepted_shape[parameter] + step * shape_direction[parameter];
        }
        REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
        mfem::Vector plus_residual, minus_residual;
        equation_backend::volume_residual(cache.template get<stages::Linearization>(), plus_residual);
        for (std::size_t parameter = 0; parameter < shape.size(); ++parameter)
        {
            shape[parameter] = accepted_shape[parameter] - step * shape_direction[parameter];
        }
        REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
        equation_backend::volume_residual(cache.template get<stages::Linearization>(), minus_residual);
        plus_residual -= minus_residual;
        plus_residual /= 2.0 * step;
        CHECK(equation_fixture::relative_error(analytic, plus_residual) < 2.0e-8);

        shape = accepted_shape;
        REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
    }
} // namespace

TEST_CASE("Deformation and volume equations share one dependency graph with atomic preparation", "[unit][deformation][equation_integration]")
{
    auto generated = serif::tests::deformation::make_mesh();
    equation_fixture::Fixture fixture(std::move(*generated.mesh));
    const auto reference = deformation_backend::deformation_discretization(fixture.enthalpy_space, fixture.displacement_space, *generated.reference_mesh);
    Model model{.surface = {.maximum_degree = 1}, .stellar_interior = {}, .vacuum = {}};
    std::vector<double> shape{0.004, -0.002, 0.001, 0.003};
    mfem::Vector density = fixture.density();
    mfem::Vector enthalpy = fixture.enthalpy();
    const serif::eos::models::Polytrope eos{1.5, 0.25};
    const auto inputs = combined_inputs(model, reference, shape, fixture, density, enthalpy);
    const auto calculations = std::tuple_cat(
        deformation_backend::make_deformation_computations<deformation::SphericalHarmonicSurface>(),
        deformation_backend::make_volume_displacement_computations(fixture.displacement_binding),
        equation_backend::make_volume_equation_computations(Closure{eos}, equation_discretization(fixture)));
    auto cache = std::apply([&](auto... calculation)
    {
        return dependencies::make_computation_cache<Failure>(inputs, std::move(calculation)...);
    },
                            calculations);
    CHECK_FALSE(cache.ready());
    REQUIRE(cache.try_prepare(inputs));
    check_against_volume_cache(cache, fixture, eos, density, enthalpy);
    check_shape_equation_linearization(cache, model, reference, shape, fixture, density, enthalpy);

    const auto& backend_vector = cache.get<deformation_backend::volume_displacement_keys::Coefficients>();
    CHECK(&cache.get<equation_backend::DisplacementCoefficients>().values == &backend_vector);
    REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
    CHECK(cache.report().computed_count() == 0);

    shape[1] += 0.001;
    REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
    CHECK(cache.report().status<deformation::keys::SurfaceMap>() == Status::reused);
    CHECK(cache.report().status<deformation::keys::RadialInterpolation>() == Status::reused);
    CHECK(cache.report().status<deformation::keys::VolumeDisplacement>() == Status::computed);
    CHECK(cache.report().status<deformation_backend::volume_displacement_keys::CheckedBinding>() == Status::reused);
    CHECK(cache.report().status<stages::Basis<0>>() == Status::reused);
    CHECK(cache.report().status<stages::Geometry<0>>() == Status::computed);
    CHECK(cache.report().status<stages::ReferenceValues<0, DensityValue>>() == Status::reused);
    check_against_volume_cache(cache, fixture, eos, density, enthalpy);

    density += 0.02;
    REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
    CHECK(cache.report().status<deformation::keys::VolumeDisplacement>() == Status::reused);
    CHECK(cache.report().status<stages::Geometry<0>>() == Status::reused);
    CHECK(cache.report().status<stages::ReferenceValues<0, DensityValue>>() == Status::computed);
    check_against_volume_cache(cache, fixture, eos, density, enthalpy);

    // A different interior response changes derivative information even when a spherical zero-displacement state has equal values.
    std::ranges::fill(shape, 0.0);
    REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
    model.stellar_interior.radial_power = 3.0;
    REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
    CHECK(cache.report().status<deformation::keys::VolumeExtension>() == Status::computed);
    CHECK(cache.report().status<deformation::keys::RadialInterpolation>() == Status::reused);
    CHECK(cache.report().status<equation_backend::DisplacementCoefficients>() == Status::computed);
    CHECK(cache.report().status<stages::Geometry<0>>() == Status::computed);
    CHECK(cache.get<equation_backend::DisplacementCoefficients>().values.Norml2() == 0.0);

    shape[0] = std::numeric_limits<double>::quiet_NaN();
    const auto invalid_shape = cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy));
    REQUIRE_FALSE(invalid_shape);
    CHECK(std::holds_alternative<deformation::DeformationFailure>(invalid_shape.error().error));
    CHECK_FALSE(cache.ready());
    CHECK_FALSE(cache.contains<stages::Linearization>());
    CHECK_FALSE(cache.contains<deformation::keys::Linearization>());

    shape[0] = 0.003;
    enthalpy = -1.0;
    const auto invalid_equation = cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy));
    REQUIRE_FALSE(invalid_equation);
    CHECK(std::holds_alternative<equation_backend::PreparationRejection>(invalid_equation.error().error));
    CHECK_FALSE(cache.ready());
    CHECK_FALSE(cache.contains<stages::Linearization>());
    CHECK_FALSE(cache.contains<deformation::keys::Linearization>());

    enthalpy = fixture.enthalpy();
    REQUIRE(cache.try_prepare(combined_inputs(model, reference, shape, fixture, density, enthalpy)));
    check_against_volume_cache(cache, fixture, eos, density, enthalpy);
}

TEST_CASE("The deformation displacement connection respects both global MFEM orderings", "[unit][deformation][equation_integration]")
{
    for (const mfem::Ordering::Type ordering : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM})
    {
        equation_fixture::Fixture fixture(mfem::Element::HEXAHEDRON, ordering);
        const mfem::Vector density = fixture.density();
        const mfem::Vector enthalpy = fixture.enthalpy();
        const mfem::Vector reference_displacement = fixture.displacement(0.03);
        const int scalar_nodes = fixture.displacement_space.GetNDofs();
        std::vector<double> displacement(fixture.displacement_binding.full_size());
        for (int component = 0; component < 3; ++component)
        {
            for (int node = 0; node < scalar_nodes; ++node)
            {
                const int full = fixture.displacement_space.DofToVDof(node, component);
                const int reduced = *fixture.displacement_binding.reduced_dof(full);
                displacement[node + component * scalar_nodes] = reference_displacement[reduced];
            }
        }
        const auto equation_inputs = equation_backend::make_volume_inputs<Closure>(equation_discretization(fixture), fixture.state(density, enthalpy));
        const auto reference = deformation_backend::deformation_discretization(fixture.enthalpy_space, fixture.displacement_space, fixture.mesh);
        const auto inputs = dependencies::make_dependency_context(
            dependencies::borrow<deformation::keys::VolumeDisplacement>(displacement, deformation_backend::DeformationValueTracking{}),
            dependencies::own<deformation::keys::ReferenceDiscretization>(reference, deformation_backend::ReferenceDiscretizationTracking{}),
            equation_inputs.resource<equation_backend::MeshAndCoefficientBindings>(),
            equation_inputs.resource<equation_backend::CompactificationCoefficients>(),
            equation_inputs.resource<blocks::density::mass>(),
            equation_inputs.resource<::serif::tests::equations::Enthalpy>());
        const serif::eos::models::Polytrope eos{1.5, 0.25};
        const auto calculations = std::tuple_cat(deformation_backend::make_volume_displacement_computations(fixture.displacement_binding), equation_backend::make_volume_equation_computations(Closure{eos}, equation_discretization(fixture)));
        auto cache = std::apply([&](auto... calculation)
        {
            return dependencies::make_computation_cache<Failure>(inputs, std::move(calculation)...);
        },
                                calculations);
        REQUIRE(cache.try_prepare(inputs));
        const mfem::Vector& actual = cache.get<equation_backend::DisplacementCoefficients>().values;
        CHECK(equation_fixture::relative_error(actual, reference_displacement) == 0.0);
        check_against_volume_cache(cache, fixture, eos, density, enthalpy);

        equation_fixture::Fixture other_fixture(mfem::Element::HEXAHEDRON, ordering);
        const auto other_reference = deformation_backend::deformation_discretization(other_fixture.enthalpy_space, other_fixture.displacement_space, other_fixture.mesh);
        const auto mismatched_inputs = dependencies::make_dependency_context(
            inputs.resource<deformation::keys::VolumeDisplacement>(),
            dependencies::own<deformation::keys::ReferenceDiscretization>(other_reference, deformation_backend::ReferenceDiscretizationTracking{}),
            inputs.resource<equation_backend::MeshAndCoefficientBindings>(),
            inputs.resource<equation_backend::CompactificationCoefficients>(),
            inputs.resource<blocks::density::mass>(),
            inputs.resource<::serif::tests::equations::Enthalpy>());
        const auto mismatch = cache.try_prepare(mismatched_inputs);
        REQUIRE_FALSE(mismatch);
        CHECK(std::holds_alternative<deformation::DeformationFailure>(mismatch.error().error));
        CHECK_FALSE(cache.ready());
        REQUIRE(cache.try_prepare(inputs));

        displacement.pop_back();
        const auto invalid = cache.try_prepare(inputs);
        REQUIRE_FALSE(invalid);
        CHECK(std::holds_alternative<deformation::DeformationFailure>(invalid.error().error));
        CHECK_FALSE(cache.ready());
    }
}

TEST_CASE("Deformation rejects an equation displacement space with matching dimensions but different scalar numbering", "[unit][deformation][equation_integration]")
{
    equation_fixture::Fixture fixture;
    ReorderedDisplacementSpace reordered(&fixture.mesh, &fixture.enthalpy_collection, 3, mfem::Ordering::byNODES);
    using Binding = serif::discretization::backend::mfem::TermBinding<blocks::displacement::vector, equation_fixture::Schema>;
    const Binding binding(reordered);
    REQUIRE(reordered.GetNDofs() == fixture.displacement_space.GetNDofs());
    REQUIRE(reordered.FEColl() == fixture.displacement_space.FEColl());
    REQUIRE(binding.size() == fixture.displacement_binding.size());
    const VolumeDiscretization discretization{
        .residual_binding = fixture.density_binding,
        .displacement_binding = binding,
        .compactification = fixture.compactification,
        .mapper = fixture.mapper};
    const mfem::Vector density = fixture.density();
    const mfem::Vector enthalpy = fixture.enthalpy();
    const auto equation_inputs = equation_backend::make_volume_inputs<Closure>(discretization, fixture.state(density, enthalpy));
    const auto reference = deformation_backend::deformation_discretization(fixture.enthalpy_space, fixture.displacement_space, fixture.mesh);
    const std::vector<double> displacement(binding.full_size(), 0.0);
    const auto inputs = dependencies::make_dependency_context(
        dependencies::borrow<deformation::keys::VolumeDisplacement>(displacement, deformation_backend::DeformationValueTracking{}),
        dependencies::own<deformation::keys::ReferenceDiscretization>(reference, deformation_backend::ReferenceDiscretizationTracking{}),
        equation_inputs.resource<equation_backend::MeshAndCoefficientBindings>());
    auto cache = std::apply([&](auto... calculation)
    {
        return dependencies::make_computation_cache<deformation::DeformationFailure>(inputs, std::move(calculation)...);
    },
                            deformation_backend::make_volume_displacement_computations(binding));
    const auto outcome = cache.try_prepare(inputs);
    REQUIRE_FALSE(outcome);
    CHECK(outcome.error().error.reason == deformation::DeformationFailureReason::incompatible_discretization);
    CHECK(outcome.error().error.element == 0);
    CHECK_FALSE(cache.ready());
    CHECK_FALSE(cache.contains<equation_backend::DisplacementCoefficients>());
}

TEST_CASE("Displacement values and sparse derivatives share supported MFEM coordinates", "[unit][deformation][displacement_mapping]") {
    for (const mfem::Ordering::Type ordering : {mfem::Ordering::byNODES, mfem::Ordering::byVDIM}) {
        for (const bool restrict_support : {false, true}) {
            mfem::Vector component_values, full_values, expected, dual, expected_transpose;
            const auto projection = [&]() {
                mfem::Mesh mesh = mfem::Mesh::MakeCartesian3D(2, 1, 1, mfem::Element::HEXAHEDRON);
                mfem::H1_FECollection collection(1, 3);
                mfem::FiniteElementSpace space(&mesh, &collection, 3, ordering);
                mfem::Array<int> support(mesh.GetNE());
                support = 1;
                if (restrict_support) {
                    support[1] = 0;
                }
                const SupportedDisplacement binding(space, support);
                if (restrict_support) {
                    REQUIRE(binding.size() < binding.full_size());
                }
                component_values.SetSize(binding.full_size());
                full_values.SetSize(binding.full_size());
                for (int component = 0; component < space.GetVDim(); ++component) {
                    for (int node = 0; node < space.GetNDofs(); ++node) {
                        const int column = node + component * space.GetNDofs();
                        component_values[column] = 0.2 + 0.03 * column;
                        full_values[space.DofToVDof(node, component)] = component_values[column];
                    }
                }
                expected = binding.gather(full_values);
                mfem::Vector converted(binding.size());
                std::vector<int> visits(static_cast<std::size_t>(binding.size()), 0);
                deformation_backend::for_each_displacement_coefficient(binding, [&](const int row, const int column) {
                    converted[row] = component_values[column];
                    ++visits[static_cast<std::size_t>(row)];
                });
                CHECK(equation_fixture::relative_error(converted, expected) == 0.0);
                CHECK(std::ranges::all_of(visits, [](const int count) { return count == 1; }));

                dual = equation_fixture::deterministic(binding.size(), 0.4);
                const mfem::Vector scattered = binding.scatter(dual);
                expected_transpose.SetSize(binding.full_size());
                for (int component = 0; component < space.GetVDim(); ++component) {
                    for (int node = 0; node < space.GetNDofs(); ++node) {
                        expected_transpose[node + component * space.GetNDofs()] = scattered[space.DofToVDof(node, component)];
                    }
                }
                return serif::model::backend::mfem::displacement_projection(binding);
            }();

            // The immutable projection must not depend on the departed mesh or binding.
            REQUIRE(projection.has_assembly());
            const auto sparse = projection.assemble();
            REQUIRE(sparse->Finalized());
            mfem::Vector actual(expected.Size()), transpose(expected_transpose.Size());
            projection.action()->Mult(component_values, actual);
            projection.action()->MultTranspose(dual, transpose);
            CHECK(equation_fixture::relative_error(actual, expected) == 0.0);
            CHECK(equation_fixture::relative_error(transpose, expected_transpose) == 0.0);
            CHECK(std::abs(actual * dual - component_values * transpose) < 1.0e-12);
            CHECK(sparse->NumNonZeroElems() == expected.Size());
        }
    }
}