#include <algorithm>
#include <limits>
#include <vector>
#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

#include "serif/tests/equations/fixture.hpp"
#include "serif/backend/mfem/discretization/element_coordinates.hpp"
#include "serif/tests/test_tags.hpp"

namespace serif::tests::equations {
    TEST_CASE("Scalar Operand Evaluation Rejects Incompatible MFEM Elements", tags::unit & tags::equations) {
        Fixture f;
        ::mfem::ND_FECollection collection(1, 3);
        ::mfem::FiniteElementSpace space(&f.mesh, &collection);
        const auto& element = *space.GetFE(0);
        const auto& rule = ::mfem::IntRules.Get(element.GetGeomType(), 2);
        backend::BasisEvaluation basis;
        backend::TensorBasisCache cache;
        backend::BasisWorkspace workspace;
        auto& transformation = *f.mesh.GetElementTransformation(0);
        using Operand = discretization::forms::Operand<blocks::density::mass>;

        // Direct sampling must enforce the scalar representation even without a TermBinding.
        CHECK_THROWS_AS(
            backend::OperandEvaluator<Operand>::sample_basis(element, rule, basis, cache, backend::ContractionMode::automatic, transformation, workspace),
            std::invalid_argument
        );
    }

    TEST_CASE("Equation State Guards Remain Active For Invalid Sizes And Nonfinite Values", tags::unit & tags::equations) {
        Fixture f;
        auto rho = f.density(), h = f.enthalpy(), u = f.displacement(0.0);
        PreparedClosure prepared(
            Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
            f.mapper
        );
        ::mfem::Vector result;
        CHECK_THROWS_AS(prepared.residual(f.workspace, result), std::logic_error);
        CHECK_THROWS_AS(prepared.width(), std::logic_error);
        const auto wrong = constant(rho.Size() + 1, 0.0);
        CHECK_THROWS_AS(prepared.prepare(f.state(wrong, h), u, f.workspace), std::invalid_argument);
        CHECK_FALSE(prepared.is_prepared());
        rho[0] = std::numeric_limits<double>::quiet_NaN();
        const auto invalid = prepared.try_prepare(f.state(rho, h), u, f.workspace);
        REQUIRE_FALSE(invalid.has_value());
        CHECK(invalid.error().reason == backend::PreparationRejectionReason::NON_FINITE);
        rho = f.density();
        prepared.prepare(f.state(rho, h), u, f.workspace);
        CHECK_THROWS_AS(prepared.apply_jacobian(f.state(wrong, h), u, f.workspace, result), std::invalid_argument);
        CHECK_THROWS_AS(prepared.apply_jacobian(wrong, f.workspace, result), std::invalid_argument);
        h[0] = std::numeric_limits<double>::infinity();
        CHECK_THROWS_AS(prepared.apply_jacobian(f.state(rho, h), u, f.workspace, result), std::domain_error);
    }

    TEST_CASE("Equation Bindings Reject Wrong Spaces And Different Meshes", tags::unit & tags::equations) {
        Fixture f, other;
        ::mfem::H1_FECollection wrong_collection(2, 3);
        ::mfem::FiniteElementSpace wrong_order(&f.mesh, &wrong_collection);
        CHECK_THROWS_AS(
            (discretization::backend::mfem::TermBinding<::serif::tests::equations::Enthalpy, Schema>{wrong_order}), std::invalid_argument
        );
        CHECK_THROWS_AS(
            (discretization::backend::mfem::TermBinding<blocks::density::mass, Schema>{wrong_order}), std::invalid_argument
        );
        const auto rho = f.density(), h = other.enthalpy(), u = f.displacement(0.0);
        PreparedClosure prepared(
            Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
            f.mapper
        );
        const auto wrong_state =
            backend::make_term_coefficients(backend::bind(f.density_binding, rho), backend::bind(other.enthalpy_binding, h));
        CHECK_THROWS_AS(prepared.prepare(wrong_state, u, f.workspace), std::invalid_argument);
        const auto local_h = f.enthalpy();
        prepared.prepare(f.state(rho, local_h), u, f.workspace);
        ::mfem::Vector result;
        CHECK_THROWS_AS(prepared.apply_jacobian(wrong_state, u, f.workspace, result), std::invalid_argument);
    }

    TEST_CASE("Prepared Equations Reject Changed Discretization And Domain Support", tags::unit & tags::equations) {
        Fixture f;
        const auto rho = f.density(), h = f.enthalpy(), u = f.displacement(0.0);
        PreparedClosure prepared(
            Closure{eos::models::Polytrope{3.0, 1.5}}, f.density_binding, f.displacement_binding, f.compactification,
            f.mapper
        );
        prepared.prepare(f.state(rho, h), u, f.workspace);
        ::mfem::Vector result;
        SECTION("Changed domain attributes") {
            f.mesh.GetElement(0)->SetAttribute(3);
            CHECK_THROWS_AS(prepared.residual(f.workspace, result), std::logic_error);
        }
        SECTION("Mesh refinement") {
            f.mesh.UniformRefinement();
            CHECK_THROWS_AS(prepared.residual(f.workspace, result), std::logic_error);
        }
        SECTION("Updated node geometry") {
            f.mesh.NodesUpdated();
            CHECK_THROWS_AS(prepared.residual(f.workspace, result), std::logic_error);
            prepared.prepare(f.state(rho, h), u, f.workspace);
            CHECK_NOTHROW(prepared.residual(f.workspace, result));
        }
        SECTION("Rebound compactification") {
            Fixture other;
            f.compactification.SetSpace(&other.enthalpy_space);
            CHECK_THROWS_AS(prepared.residual(f.workspace, result), std::logic_error);
        }
    }

    TEST_CASE("Equation Field Restriction Excludes Only Vacuum-Only Coefficients", tags::unit & tags::equations) {
        Fixture f;
        auto full = constant(f.density_space.GetVSize(), 0.0);
        ::mfem::Array<int> dofs;
        std::vector<bool> stellar_dofs(full.Size(), false);
        for (int element = 0; element < f.mesh.GetNE(); ++element) {
            if (f.mesh.GetAttribute(element) == 3) {
                continue;
            }
            f.density_space.GetElementDofs(element, dofs);
            for (const int dof : dofs) {
                stellar_dofs[dof] = true;
            }
        }
        for (int dof = 0; dof < full.Size(); ++dof) {
            full[dof] = stellar_dofs[dof] ? 0.5 : std::numeric_limits<double>::quiet_NaN();
        }
        const auto selected = std::count(stellar_dofs.begin(), stellar_dofs.end(), true);
        REQUIRE(selected > 0);
        REQUIRE(selected < full.Size());
        CHECK(f.density_binding.size() == selected);
        const auto reduced = f.density_binding.gather(full);
        CHECK_NOTHROW(f.density_binding.validate_values(reduced));
        const auto expanded = f.density_binding.scatter(reduced);
        CHECK(expanded.CheckFinite() == 0);
        CHECK(reduced.Min() == 0.5);
    }
    TEST_CASE("Equation Bindings Reject Unsupported Meshes And Element Indices", tags::unit & tags::equations) {
        Fixture f;
        CHECK_THROWS_AS(f.density_binding.supports(-1), std::out_of_range);
        CHECK_THROWS_AS(f.density_binding.supports(f.mesh.GetNE()), std::out_of_range);
        ::mfem::Vector empty, output;
        CHECK_THROWS_AS(f.density_binding.element_values(0, empty, output), std::invalid_argument);
        auto flat_mesh = ::mfem::Mesh::MakeCartesian2D(2, 2, ::mfem::Element::QUADRILATERAL);
        ::mfem::L2_FECollection flat_collection(2, 2);
        ::mfem::FiniteElementSpace flat_space(&flat_mesh, &flat_collection);
        CHECK_THROWS_AS((discretization::backend::mfem::TermBinding<blocks::density::mass, Schema>{flat_space}), std::invalid_argument);
        auto nonconforming_mesh = make_mesh(::mfem::Element::HEXAHEDRON);
        nonconforming_mesh.EnsureNCMesh();
        ::mfem::FiniteElementSpace nonconforming_space(&nonconforming_mesh, &f.density_collection);
        CHECK_NOTHROW((discretization::backend::mfem::TermBinding<blocks::density::mass, Schema>{nonconforming_space}));
    }
    namespace {
        /** @brief Compare primal/dual coordinate access with independent MFEM extraction and assembly. */
        void check_element_coordinates(const discretization::backend::mfem::DofMap& binding) {
            const auto reduced = deterministic(binding.size(), 0.31);
            const auto full = binding.scatter(reduced);
            const auto elements = binding.validated_elements();
            for (int element = 0; element < binding.mesh().GetNE(); ++element) {
                ::mfem::Array<int> dofs;
                ::mfem::DofTransformation transformation;
                binding.space().GetElementVDofs(element, dofs, transformation);
                ::mfem::Vector expected, actual;
                full.GetSubVector(dofs, expected);
                if (transformation.GetDofTransformation() != nullptr) {
                    transformation.InvTransformPrimal(expected);
                }
                elements.element_values(element, reduced, actual);
                CHECK(relative_error(actual, expected) < 1.0e-13);

                const auto coordinates = elements.coordinates(element);
                REQUIRE(coordinates.size() == dofs.Size());
                for (int local = 0; local < dofs.Size(); ++local) {
                    const int full_index = dofs[local] < 0 ? -1 - dofs[local] : dofs[local];
                    const int index = binding.reduced_dof(full_index).value_or(-1);
                    CHECK(coordinates.reduced_dof(local) == index);
                    if (index >= 0) {
                        CHECK(coordinates.local_dof(index) == local);
                        ::mfem::Vector unit = constant(binding.size(), 0.0), direction;
                        unit[index] = 1.0;
                        coordinates.element_values(unit, expected);
                        coordinates.unit_direction(local, direction);
                        CHECK(relative_error(direction, expected) < 1.0e-13);
                    }
                }

                const auto dual = deterministic(dofs.Size(), 0.78);
                ::mfem::Vector expected_dual(dual), actual_dual(dual);
                auto expected_full = constant(binding.full_size(), 0.0);
                auto assembled = constant(binding.full_size(), 0.0);
                if (transformation.GetDofTransformation() != nullptr) {
                    transformation.TransformDual(expected_dual);
                }
                expected_full.AddElementVector(dofs, expected_dual);
                elements.add_element(element, actual_dual, assembled);
                CHECK(relative_error(assembled, expected_full) < 1.0e-13);
                CHECK(full * assembled == Catch::Approx(actual * dual).margin(1.0e-12));

                actual_dual = dual;
                coordinates.orient_dual(actual_dual);
                for (int local = 0; local < dofs.Size(); ++local) {
                    const double expected_sign = dofs[local] < 0 ? -1.0 : 1.0;
                    CHECK(actual_dual[local] == Catch::Approx(expected_sign * expected_dual[local]).margin(1.0e-14));
                }
            }
        }
    }

    TEST_CASE("Validated Element Traversals Preserve Native MFEM Coordinates", tags::unit & tags::equations) {
        for (const auto geometry : {::mfem::Element::HEXAHEDRON, ::mfem::Element::TETRAHEDRON}) {
            for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
                Fixture fixture(geometry, ordering);
                check_element_coordinates(fixture.density_binding);
                check_element_coordinates(fixture.enthalpy_binding);
                check_element_coordinates(fixture.displacement_binding);
            }
        }
    }

    TEST_CASE("Validated Element Traversals Preserve Global Scalar Support And Guards", tags::unit & tags::equations) {
        Fixture fixture;
        const auto& binding = fixture.constant_binding;
        const auto elements = binding.validated_elements();
        const auto reduced = constant(1, 2.0);
        for (int element = 0; element < fixture.mesh.GetNE(); ++element) {
            ::mfem::Vector local;
            elements.element_values(element, reduced, local);
            CHECK(local[0] == (binding.supports(element) ? 2.0 : 0.0));
            auto assembled = constant(1, 0.0);
            auto dual = constant(1, 3.0);
            elements.add_element(element, dual, assembled);
            CHECK(assembled[0] == (binding.supports(element) ? 3.0 : 0.0));
        }
        ::mfem::Vector output, empty;
        CHECK_THROWS_AS(elements.element_values(-1, reduced, output), std::out_of_range);
        CHECK_THROWS_AS(elements.element_values(0, empty, output), std::invalid_argument);
        CHECK_THROWS_AS(elements.add_element(0, empty, output), std::invalid_argument);
        fixture.mesh.GetElement(0)->SetAttribute(3);
        CHECK_THROWS_AS(binding.validated_elements(), std::logic_error);
        CHECK_THROWS_AS(binding.element_values(1, reduced, output), std::logic_error);
        CHECK_THROWS_AS(elements.element_values(0, reduced, output), std::logic_error);
    }

    TEST_CASE("Validated Element Traversals Reject Changed Mesh Revisions", tags::unit & tags::equations) {
        Fixture fixture;
        const auto elements = fixture.enthalpy_binding.validated_elements();
        const auto values = fixture.enthalpy();
        fixture.mesh.UniformRefinement();
        ::mfem::Vector output;
        CHECK_THROWS_AS(elements.element_values(0, values, output), std::logic_error);
        CHECK_THROWS_AS(fixture.enthalpy_binding.validated_elements(), std::logic_error);
    }
} // namespace serif::tests::equations