#include <algorithm>
#include <cmath>
#include <memory>
#include <numeric>
#include <span>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include "serif/backend/mfem/geometry/deformation/reference_geometry.hpp"
#include "serif/geometry/deformation/composition.hpp"
#include "serif/geometry/deformation/linearization.hpp"
#include "serif/geometry/deformation/surface_embedding.hpp"
#include "serif/tests/deformation/mesh.hpp"

namespace
{
    namespace deformation = serif::geometry::deformation;
    namespace backend = deformation::backend::mfem;
    using DomainMap = deformation::ComposedDeformation<deformation::SurfaceEmbedding, deformation::RadialExtension>;

    /** @brief Pair complete surface and volume actions on the same serial STROID geometry. */
    template <typename Model>
    [[nodiscard]] DomainMap make_domain_map(serif::tests::deformation::PairedMesh& mesh, const Model& model)
    {
        const auto reference = backend::sample_reference_surface(mesh.scalar_space, mesh.surface_attribute);
        REQUIRE(reference);
        auto surface = deformation::build_surface_embedding(model, reference->positions);
        auto interpolation = mesh.interpolation();
        REQUIRE(surface);
        REQUIRE(interpolation);
        auto extension = deformation::make_radial_extension(std::make_shared<const deformation::RadialInterpolation>(std::move(*interpolation)));
        REQUIRE(extension);
        auto result = deformation::compose_deformation(std::make_shared<const deformation::SurfaceEmbedding>(std::move(*surface)), std::make_shared<const deformation::RadialExtension>(std::move(*extension)));
        REQUIRE(result);
        return std::move(*result);
    }

    /** @brief Compare the assembled physical map and complete transpose with the original MeanField tolerances. */
    void check_complete_actions(const DomainMap& map, const std::vector<double>& state, const std::vector<double>& direction, const double step, const double difference_tolerance, const double work_tolerance)
    {
        std::vector<double> plus(state), minus(state);
        for (std::size_t parameter = 0; parameter < state.size(); ++parameter)
        {
            plus[parameter] += step * direction[parameter];
            minus[parameter] -= step * direction[parameter];
        }
        const auto plus_value = deformation::evaluate_displacement(map, plus);
        const auto minus_value = deformation::evaluate_displacement(map, minus);
        REQUIRE(plus_value);
        REQUIRE(minus_value);
        auto workspace = deformation::make_deformation_workspace(map);
        std::vector<double> action(map.output_size()), dual(map.output_size()), force(map.input_size()), curvature(map.input_size(), 7.0);
        REQUIRE(deformation::apply_jacobian(map, state, direction, action, workspace));
        long double error_squared = 0.0L;
        long double reference_squared = 0.0L;
        for (std::size_t row = 0; row < action.size(); ++row)
        {
            const double difference = ((*plus_value)[row] - (*minus_value)[row]) / (2.0 * step);
            error_squared += static_cast<long double>(action[row] - difference) * (action[row] - difference);
            reference_squared += static_cast<long double>(difference) * difference;
            const double index = static_cast<double>(row + 1);
            dual[row] = std::sin(0.07 * index) + 0.4 * std::cos(0.11 * index);
        }
        REQUIRE(reference_squared > 0.0L);
        CHECK(std::sqrt(error_squared / reference_squared) < difference_tolerance);
        REQUIRE(deformation::apply_jacobian_transpose(map, state, dual, force, workspace));
        const double volume_work = std::inner_product(action.begin(), action.end(), dual.begin(), 0.0);
        const double parameter_work = std::inner_product(direction.begin(), direction.end(), force.begin(), 0.0);
        CHECK(std::abs(volume_work - parameter_work) <= work_tolerance * std::max({1.0, std::abs(volume_work), std::abs(parameter_work)}));
        REQUIRE(deformation::apply_pullback_derivative(map, state, direction, dual, curvature, workspace));
        CHECK(std::ranges::all_of(curvature, [](const double value)
        {
            return value == 0.0;
        }));

        std::vector<double> plus_force(map.input_size()), minus_force(map.input_size());
        REQUIRE(deformation::apply_jacobian_transpose(map, plus, dual, plus_force, workspace));
        REQUIRE(deformation::apply_jacobian_transpose(map, minus, dual, minus_force, workspace));
        CHECK(plus_force == minus_force);
    }
} // namespace

TEST_CASE("MeanField nodal domain directions preserve complete finite differences and virtual work", "[unit][deformation][meanfield]")
{
    // Preserve the original cubic displacement space, including interpolation between refined surface rays.
    for (const int refinement : {0, 1})
    {
        CAPTURE(refinement);
        serif::tests::deformation::PairedMesh mesh(refinement, 0.0, 3, 4, 5.0);
        const DomainMap map = make_domain_map(mesh, deformation::NodalRadialSurface{});
        std::vector<double> state(map.input_size()), direction(map.input_size());
        for (std::size_t parameter = 0; parameter < state.size(); ++parameter)
        {
            const double index = static_cast<double>(parameter + 1);
            state[parameter] = 0.008 * std::sin(0.13 * index);
            direction[parameter] = std::cos(0.17 * index) - 0.25 * std::sin(0.29 * index);
        }
        check_complete_actions(map, state, direction, 1.0e-6, 3.0e-10, 8.0e-13);
    }
}

TEST_CASE("MeanField harmonic domain directions preserve complete finite differences and dual work", "[unit][deformation][meanfield]")
{
    serif::tests::deformation::PairedMesh mesh(0, 0.0, 3, 4, 5.0);
    const DomainMap map = make_domain_map(mesh, deformation::SphericalHarmonicSurface{.maximum_degree = 2});
    REQUIRE(map.input_size() == 9);
    std::vector<double> state(map.input_size()), direction(map.input_size());
    for (std::size_t parameter = 0; parameter < state.size(); ++parameter)
    {
        const double value = std::sin(0.37 * static_cast<double>(parameter + 1));
        state[parameter] = 0.001 * value;
        direction[parameter] = 0.004 * value;
    }
    check_complete_actions(map, state, direction, 1.0e-5, 2.0e-9, 1.0e-11);
}