#include <algorithm>
#include <chrono>
#include <concepts>
#include <cmath>
#include <filesystem>
#include <memory>
#include <vector>

#include <catch2/catch_test_macros.hpp>
#include "mfem/fem/bilinearform.hpp"
#include "mfem/fem/bilininteg.hpp"
#include "serif/backend/mfem/discretization/element_coordinates.hpp"
#include "serif/backend/mfem/discretization/mass_diagonal.hpp"
#include "serif/backend/mfem/discretization/schema.hpp"
#include "serif/tests/nonconforming.hpp"
#include "serif/discretization/domain/schema/schemas.hpp"

namespace {
    namespace backend = serif::discretization::backend::mfem;
    namespace blocks = serif::discretization::blocks;
    using Schema = serif::discretization::domain::schema::CoreEnvelopeVacuumDomainSchema;

    /** @brief Keep each saved STROID fixture private to this invocation and remove it after any assertion. */
    struct TemporaryMesh {
        std::filesystem::path path = std::filesystem::temp_directory_path() /
            ("serif-nc-" + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()) + ".smesh");
        /** @brief Release only this test's generated file. */
        ~TemporaryMesh() {
            std::error_code error;
            std::filesystem::remove(path, error);
        }
    };

    /** @brief Use nonconstant signed data so primal restriction cannot masquerade as dual accumulation. */
    [[nodiscard]] ::mfem::Vector sequence(const int size) {
        ::mfem::Vector values(size);
        for (int index = 0; index < size; ++index) {
            values[index] = std::sin(0.17 * (index + 1)) + 0.3 * std::cos(0.23 * index);
        }
        return values;
    }

    /** @brief Compare complete vectors with a norm-scaled roundoff tolerance. */
    void check_values(const ::mfem::Vector& actual, const ::mfem::Vector& expected) {
        REQUIRE(actual.Size() == expected.Size());
        ::mfem::Vector difference(actual);
        difference -= expected;
        CHECK(difference.Norml2() <= 2.0e-12 * std::max(1.0, expected.Norml2()));
    }

    /** @brief Independently close element support through native MFEM prolongation rows. */
    [[nodiscard]] std::vector<int> supported_true_dofs(const backend::DofMap& binding) {
        const ::mfem::SparseMatrix& prolongation = *binding.space().GetConformingProlongation();
        std::vector<bool> supported(prolongation.Width(), false);
        ::mfem::Array<int> dofs;
        for (int element = 0; element < binding.mesh().GetNE(); ++element) {
            if (!binding.supports(element)) {
                continue;
            }
            binding.space().GetElementVDofs(element, dofs);
            for (const int encoded : dofs) {
                const int local = ::mfem::FiniteElementSpace::DecodeDof(encoded);
                for (int entry = prolongation.GetI()[local]; entry < prolongation.GetI()[local + 1]; ++entry) {
                    if (prolongation.GetData()[entry] != 0.0) {
                        supported[prolongation.GetJ()[entry]] = true;
                    }
                }
            }
        }
        std::vector<int> result;
        for (int index = 0; index < prolongation.Width(); ++index) {
            if (supported[index]) {
                result.push_back(index);
            }
        }
        return result;
    }

    /** @brief Compare primal, dual, and oriented element coordinates with independent native matrices. */
    void check_transfers(const backend::DofMap& binding) {
        REQUIRE(binding.constrained());
        const auto supported = supported_true_dofs(binding);
        REQUIRE(static_cast<int>(supported.size()) == binding.size());
        const ::mfem::Vector reduced = sequence(binding.size());
        ::mfem::Vector independent(binding.true_size());
        independent = 0.0;
        for (int index = 0; index < binding.size(); ++index) {
            independent[supported[index]] = reduced[index];
        }
        ::mfem::Vector expected(binding.full_size());
        binding.space().GetConformingProlongation()->Mult(independent, expected);
        check_values(binding.scatter(reduced), expected);
        check_values(binding.scatter_true(reduced), independent);
        check_values(binding.gather(expected), reduced);
        const ::mfem::Vector dual = sequence(binding.full_size());
        ::mfem::Vector true_dual(binding.true_size()), expected_dual(binding.size());
        binding.space().GetConformingProlongation()->MultTranspose(dual, true_dual);
        for (int index = 0; index < binding.size(); ++index) {
            expected_dual[index] = true_dual[supported[index]];
        }
        check_values(binding.accumulate_dual(dual), expected_dual);
        const double full_work = expected * dual;
        CHECK(std::abs(reduced * expected_dual - full_work) <= 2.0e-11 * std::max(1.0, std::abs(full_work)));
        ::mfem::Vector arbitrary_true(binding.true_size()), expected_primal(binding.size());
        binding.space().GetConformingRestriction()->Mult(dual, arbitrary_true);
        for (int index = 0; index < binding.size(); ++index) {
            expected_primal[index] = arbitrary_true[supported[index]];
        }
        check_values(binding.gather(dual), expected_primal);
        const auto elements = binding.validated_elements();
        for (int element = 0; element < binding.mesh().GetNE(); ++element) {
            if (!binding.supports(element)) {
                continue;
            }
            ::mfem::Array<int> dofs;
            const auto* transformation = binding.space().GetElementVDofs(element, dofs);
            ::mfem::Vector native, actual;
            expected.GetSubVector(dofs, native);
            if (transformation != nullptr) {
                transformation->InvTransformPrimal(native);
            }
            elements.element_values(element, reduced, actual);
            check_values(actual, native);
        }
    }

    /** @brief Retain every local mass coupling before reducing; local diagonal restriction is not an oracle. */
    template <typename Term>
    void check_mass(const backend::TermBinding<Term, Schema>& binding) {
        ::mfem::Array<int> marker(binding.mesh().attributes.Max());
        marker = 0;
        for (int element = 0; element < binding.mesh().GetNE(); ++element) {
            if (binding.supports(element)) {
                marker[binding.mesh().GetAttribute(element) - 1] = 1;
            }
        }
        ::mfem::BilinearForm native(&binding.space());
        if constexpr (std::same_as<Term, blocks::displacement::vector>) {
            native.AddDomainIntegrator(new ::mfem::VectorMassIntegrator(), marker);
        } else {
            native.AddDomainIntegrator(new ::mfem::MassIntegrator(), marker);
        }
        native.Assemble();
        native.Finalize();
        const ::mfem::SparseMatrix& prolongation = *binding.space().GetConformingProlongation();
        const std::unique_ptr<::mfem::SparseMatrix> constrained(::mfem::RAP(prolongation, native.SpMat(), prolongation));
        ::mfem::Vector diagonal;
        constrained->GetDiag(diagonal);
        const auto supported = supported_true_dofs(binding);
        const auto actual = backend::reference_mass_diagonal(binding);
        REQUIRE(actual);
        REQUIRE(actual->size() == supported.size());
        for (int index = 0; index < binding.size(); ++index) {
            const double expected = diagonal[supported[index]];
            CHECK(std::abs((*actual)[index] - expected) <= 2.0e-11 * expected);
        }
    }
}

TEST_CASE("STROID nonconforming meshes preserve native coefficient and dual transfers", "[integration][nonconforming]") {
    auto generated = serif::tests::nonconforming::mesh(0, 3);
    REQUIRE(generated.mesh->Nonconforming());
    CHECK(serif::backend::mfem::discretization::validate_schema<Schema>(*generated.mesh).valid());
    CHECK(serif::backend::mfem::discretization::validate_schema<Schema>(*generated.reference_mesh).valid());
    ::mfem::H1_FECollection scalar_collection(blocks::density::mass::FunctionalSpace::order, 3);
    ::mfem::FiniteElementSpace scalar(generated.mesh.get(), &scalar_collection);
    REQUIRE(scalar.GetTrueVSize() < scalar.GetVSize());
    backend::TermBinding<blocks::density::mass, Schema> density(scalar);
    ::mfem::H1_FECollection potential_collection(blocks::gravity::potential::FunctionalSpace::order, 3);
    ::mfem::FiniteElementSpace potential_space(generated.mesh.get(), &potential_collection);
    REQUIRE(potential_space.GetTrueVSize() < potential_space.GetVSize());
    backend::TermBinding<blocks::gravity::potential, Schema> potential(potential_space);
    CHECK(density.size() < potential.size());
    REQUIRE(potential.size() == potential_space.GetTrueVSize());
    const ::mfem::SparseMatrix& density_prolongation = *scalar.GetConformingProlongation();
    int material_hanging_coefficients = 0;
    ::mfem::Array<int> material_dofs;
    for (int element = 0; element < scalar.GetNE(); ++element) {
        if (!density.supports(element)) {
            continue;
        }
        scalar.GetElementDofs(element, material_dofs);
        for (const int local : material_dofs) {
            const int begin = density_prolongation.GetI()[local];
            if (density_prolongation.GetI()[local + 1] != begin + 1 || density_prolongation.GetData()[begin] != 1.0) {
                ++material_hanging_coefficients;
            }
        }
    }
    REQUIRE(material_hanging_coefficients > 0);
    check_transfers(density);
    check_transfers(potential);
    check_mass(density);
    check_mass(potential);
    for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
        ::mfem::FiniteElementSpace vector(generated.mesh.get(), &scalar_collection, 3, ordering);
        backend::TermBinding<blocks::displacement::vector, Schema> displacement(vector);
        check_transfers(displacement);
        check_mass(displacement);
    }
}

TEST_CASE("STROID saved meshes retain genuine hanging faces when loaded directly", "[integration][nonconforming]") {
    auto generated = serif::tests::nonconforming::mesh();
    const TemporaryMesh file;
    stroid::IO::SaveStroidMesh(generated, file.path.string());
    auto restored = stroid::IO::LoadStroidMesh(file.path.string());
    INFO((restored ? "restored" : restored.error()));
    REQUIRE(restored);
    REQUIRE(restored->mesh->Nonconforming());
    REQUIRE(restored->reference_mesh->Nonconforming());
    CHECK(restored->mesh->GetNE() == generated.mesh->GetNE());
    ::mfem::H1_FECollection collection(blocks::gravity::potential::FunctionalSpace::order, 3);
    ::mfem::FiniteElementSpace space(restored->mesh.get(), &collection);
    REQUIRE(space.GetTrueVSize() < space.GetVSize());
    backend::TermBinding<blocks::gravity::potential, Schema> potential(space);
    check_transfers(potential);
    CHECK(serif::backend::mfem::discretization::validate_schema<Schema>(*restored->mesh).valid());
}