#include <algorithm>
#include <bit>
#include <cmath>
#include <cstdint>
#include <memory>
#include <numeric>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include "mfem/fem/gridfunc.hpp"
#include "serif/backend/mfem/geometry/deformation/computations.hpp"
#include "serif/backend/mfem/model/deformation.hpp"
#include "serif/equations/registry.hpp"
#include "serif/tests/nonconforming.hpp"
#include "serif/tests/solver/stellar_fixture.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace {
    namespace deformation = serif::geometry::deformation;
    namespace backend = deformation::backend::mfem;
    namespace fixture = serif::tests::solver;

    /** @brief Compare complete deformed coordinates on both sides of every hanging face. */
    void check_hanging_geometry(mfem::FiniteElementSpace& space, const std::vector<double>& displacement) {
        mfem::GridFunction field(&space);
        REQUIRE(displacement.size() == static_cast<std::size_t>(field.Size()));
        std::ranges::copy(displacement, field.GetData());
        mfem::Mesh& mesh = *space.GetMesh();
        mfem::Vector first(3), second(3), first_motion(3), second_motion(3);
        int hanging_faces = 0;
        for (int face = 0; face < mesh.GetNumFaces(); ++face) {
            const mfem::Mesh::FaceInformation information = mesh.GetFaceInformation(face);
            if (!information.IsNonconformingFine() || !information.IsLocal()) {
                continue;
            }
            ++hanging_faces;
            mfem::FaceElementTransformations* transformation = mesh.GetFaceElementTransformations(face);
            REQUIRE(transformation->Elem1 != nullptr);
            REQUIRE(transformation->Elem2 != nullptr);
            const mfem::IntegrationRule& rule = mfem::IntRules.Get(mfem::Geometry::SQUARE, 6);
            for (int sample = 0; sample < rule.GetNPoints(); ++sample) {
                transformation->SetAllIntPoints(&rule.IntPoint(sample));
                const mfem::IntegrationPoint& first_point = transformation->Elem1->GetIntPoint();
                const mfem::IntegrationPoint& second_point = transformation->Elem2->GetIntPoint();
                transformation->Elem1->Transform(first_point, first);
                transformation->Elem2->Transform(second_point, second);
                field.GetVectorValue(transformation->Elem1No, first_point, first_motion);
                field.GetVectorValue(transformation->Elem2No, second_point, second_motion);
                first += first_motion;
                second += second_motion;
                first -= second;
                CHECK(first.Norml2() < 2.0e-11);
            }
        }
        REQUIRE(hanging_faces > 0);
    }

    /** @brief Check independent displacement extraction and its assembled transpose for either MFEM ordering. */
    void check_displacement_projection(mfem::FiniteElementSpace& space, const std::vector<double>& component_values) {
        serif::discretization::backend::mfem::TermBinding<fixture::Displacement, fixture::Schema> binding{space};
        const auto projection = serif::model::backend::mfem::displacement_projection(binding);
        const mfem::SparseMatrix* restriction = space.GetConformingRestriction();
        REQUIRE(restriction != nullptr);
        mfem::Vector local(space.GetVSize()), component(space.GetVSize()), expected(space.GetTrueVSize());
        for (int node = 0; node < space.GetNDofs(); ++node) {
            for (int axis = 0; axis < 3; ++axis) {
                const int column = node + axis * space.GetNDofs();
                component[column] = component_values[column];
                local[space.DofToVDof(node, axis)] = component[column];
            }
        }
        restriction->Mult(local, expected);
        REQUIRE(binding.size() == expected.Size());
        mfem::Vector actual(binding.size()), dual(binding.size()), transpose(space.GetVSize());
        projection.action()->Mult(component, actual);
        actual -= expected;
        CHECK(actual.Norml2() < 2.0e-13);
        for (int index = 0; index < dual.Size(); ++index) {
            dual[index] = std::sin(0.13 * (index + 1));
        }
        projection.action()->MultTranspose(dual, transpose);
        const double reduced_work = expected * dual;
        CHECK(std::abs(component * transpose - reduced_work) < 2.0e-11 * std::max(1.0, std::abs(reduced_work)));
        const auto sparse = projection.assemble();
        REQUIRE(sparse);
        mfem::Vector sparse_action(binding.size()), sparse_transpose(space.GetVSize());
        sparse->Mult(component, sparse_action);
        sparse->MultTranspose(dual, sparse_transpose);
        sparse_action -= expected;
        sparse_transpose -= transpose;
        CHECK(sparse_action.Norml2() < 2.0e-13);
        CHECK(sparse_transpose.Norml2() < 2.0e-13);
    }

    /** @brief Preserve convenience calls while leaving reusable scratch available for repeated actions. */
    void check_extension_call_forms(const backend::ConstrainedRadialExtension& extension, const std::vector<double>& surface) {
        auto workspace = extension.make_workspace();
        std::vector<double> expected(extension.output_size()), actual(extension.output_size());
        REQUIRE(extension.apply(surface, expected, workspace));
        REQUIRE(extension.apply(surface, actual));
        CHECK(actual == expected);
        std::vector<double> force(extension.output_size());
        for (std::size_t index = 0; index < force.size(); ++index) {
            force[index] = std::sin(0.19 * (index + 1));
        }
        expected.resize(extension.input_size());
        actual.resize(extension.input_size());
        REQUIRE(extension.apply_transpose(force, expected, workspace));
        REQUIRE(extension.apply_transpose(force, actual));
        CHECK(actual == expected);
    }
}

TEST_CASE("STROID nonconforming geometry retains continuity and complete shape derivatives", "[integration][deformation][nonconforming]") {
    const auto resources = std::make_shared<fixture::StellarResources>(serif::tests::nonconforming::mesh(2, 0, 2));
    const auto finite_elements = fixture::make_discretization(resources, 2);
    std::vector<double> shape(9), direction(9);
    for (std::size_t index = 0; index < shape.size(); ++index) {
        shape[index] = 0.002 * std::sin(0.37 * (index + 1));
        direction[index] = std::cos(0.23 * (index + 1));
    }
    const auto inputs = backend::make_deformation_inputs(finite_elements.deformation.model, finite_elements.deformation.reference, shape);
    auto cache = serif::utils::dependencies::make_computation_cache<deformation::DeformationFailure>(
        backend::make_deformation_graph<deformation::SphericalHarmonicSurface>(), inputs);
    REQUIRE(cache.try_prepare(inputs));
    const auto map = cache.get<deformation::keys::DomainMap>();
    const auto extension = cache.get<deformation::keys::VolumeExtension>();
    const std::vector<double>& volume = cache.get<deformation::keys::VolumeDisplacement>();
    mfem::FiniteElementSpace& space = resources->reference_displacement;
    REQUIRE(space.GetTrueVSize() < space.GetVSize());
    REQUIRE(map->output_size() == static_cast<std::size_t>(space.GetVSize()));

    const auto raw_extension = deformation::make_radial_extension(cache.get<deformation::keys::RadialInterpolation>());
    REQUIRE(raw_extension);
    const auto raw = deformation::evaluate_displacement(*raw_extension, cache.get<deformation::keys::SurfaceDisplacement>());
    REQUIRE(raw);
    mfem::Vector local(space.GetVSize()), independent(space.GetTrueVSize()), expected(space.GetVSize());
    std::ranges::copy(*raw, local.GetData());
    space.GetConformingRestriction()->Mult(local, independent);
    space.GetConformingProlongation()->Mult(independent, expected);
    mfem::Vector defect(local);
    defect -= expected;
    REQUIRE(defect.Norml2() > 1.0e-10);
    for (int index = 0; index < expected.Size(); ++index) {
        CHECK(std::abs(volume[index] - expected[index]) < 2.0e-13);
    }
    check_hanging_geometry(space, volume);
    check_displacement_projection(resources->reference_displacement, volume);
    check_displacement_projection(resources->displacement_space, volume);

    const auto& reference = *cache.get<deformation::keys::ReferenceSurface>();
    const std::vector<double>& surface = cache.get<deformation::keys::SurfaceDisplacement>();
    check_extension_call_forms(*extension, surface);
    for (std::size_t node = 0; node < reference.scalar_dofs.size(); ++node) {
        for (int component = 0; component < 3; ++component) {
            CHECK(std::abs(volume[reference.scalar_dofs[node] + component * space.GetNDofs()] - surface[3 * node + component]) < 2.0e-13);
        }
    }
    mfem::Array<int> boundary(space.GetMesh()->bdr_attributes.Max()), boundary_dofs;
    boundary = 0;
    boundary[fixture::Schema::boundary_id<serif::discretization::domain::InfinitySurfaceBoundary>() - 1] = 1;
    space.GetEssentialVDofs(boundary, boundary_dofs);
    for (int dof = 0; dof < boundary_dofs.Size(); ++dof) {
        if (boundary_dofs[dof] != 0) {
            CHECK(std::abs(volume[dof]) < 2.0e-13);
        }
    }

    auto workspace = deformation::make_deformation_workspace(*map);
    std::vector<double> tangent(volume.size()), force(volume.size()), pullback(shape.size());
    REQUIRE(deformation::apply_jacobian(*map, shape, direction, tangent, workspace));
    for (std::size_t index = 0; index < force.size(); ++index) {
        force[index] = std::sin(0.17 * (index + 1));
    }
    REQUIRE(deformation::apply_jacobian_transpose(*map, shape, force, pullback, workspace));
    const double volume_work = std::inner_product(tangent.begin(), tangent.end(), force.begin(), 0.0);
    const double shape_work = std::inner_product(direction.begin(), direction.end(), pullback.begin(), 0.0);
    CHECK(std::abs(volume_work - shape_work) < 2.0e-11 * std::max({1.0, std::abs(volume_work), std::abs(shape_work)}));
    std::vector<double> plus(shape), minus(shape);
    constexpr double step = 1.0e-6;
    for (std::size_t index = 0; index < shape.size(); ++index) {
        plus[index] += step * direction[index];
        minus[index] -= step * direction[index];
    }
    const auto positive = deformation::evaluate_displacement(*map, plus);
    const auto negative = deformation::evaluate_displacement(*map, minus);
    REQUIRE(positive);
    REQUIRE(negative);
    double error_squared = 0.0, norm_squared = 0.0;
    for (std::size_t index = 0; index < tangent.size(); ++index) {
        const double error = ((*positive)[index] - (*negative)[index]) / (2.0 * step) - tangent[index];
        error_squared += error * error;
        norm_squared += tangent[index] * tangent[index];
    }
    REQUIRE(norm_squared > 0.0);
    CHECK(std::sqrt(error_squared / norm_squared) < 2.0e-9);
    backend::ConstrainedRadialExtension::Workspace wrong_workspace;
    std::vector<double> output(volume.size());
    CHECK_FALSE(extension->apply(surface, output, wrong_workspace));
}

TEST_CASE("Conforming deformation graph preserves the original radial action bit for bit", "[unit][deformation][nonconforming][parity]") {
    const auto resources = fixture::make_resources(1.0, 0, 2);
    const auto finite_elements = fixture::make_discretization(resources, 2);
    const std::vector<double> shape{0.001, -0.002, 0.003, -0.001, 0.002, 0.001, -0.003, 0.002, 0.001};
    const auto inputs = backend::make_deformation_inputs(finite_elements.deformation.model, finite_elements.deformation.reference, shape);
    auto cache = serif::utils::dependencies::make_computation_cache<deformation::DeformationFailure>(
        backend::make_deformation_graph<deformation::SphericalHarmonicSurface>(), inputs);
    REQUIRE(cache.try_prepare(inputs));
    REQUIRE(resources->reference_displacement.GetConformingProlongation() == nullptr);
    check_extension_call_forms(*cache.get<deformation::keys::VolumeExtension>(), cache.get<deformation::keys::SurfaceDisplacement>());
    const auto radial = deformation::make_radial_extension(cache.get<deformation::keys::RadialInterpolation>());
    REQUIRE(radial);
    const auto expected = deformation::evaluate_displacement(*radial, cache.get<deformation::keys::SurfaceDisplacement>());
    REQUIRE(expected);
    const std::vector<double>& actual = cache.get<deformation::keys::VolumeDisplacement>();
    REQUIRE(actual.size() == expected->size());
    for (std::size_t index = 0; index < actual.size(); ++index) {
        CHECK(std::bit_cast<std::uint64_t>(actual[index]) == std::bit_cast<std::uint64_t>((*expected)[index]));
    }
}