#include <algorithm>
#include <array>
#include <catch2/catch_test_macros.hpp>
#include <cstddef>
#include <mfem.hpp>
#include <mpi.h>
#include <optional>
#include <stdexcept>
#include <vector>

import mean_field;
import test_helpers;

namespace field_dof_map_test_utils {
    namespace field  = mean_field::field;

    namespace domain = mean_field::utils::domain;

    using Schema     = domain::CoreEnvelopeVacuumDomainSchema;

    [[nodiscard]]
    mfem::Array<int> make_array(const std::initializer_list<int> values) {
        mfem::Array<int> result(static_cast<int>(values.size()));

        int index = 0;

        for (const int value : values) {
            result[index++] = value;
        }

        return result;
    }

    [[nodiscard]]
    mfem::Mesh make_split_mesh(
        const int stellarAttribute = 2,
        const int vacuumAttribute  = 3
    ) {
        int communicatorSize = 1;

        MPI_Comm_size(MPI_COMM_WORLD, &communicatorSize);

        /*
         * Ensure there are enough cells that every reasonable MPI test
         * configuration has useful work available.
         */
        const int xElementCount     = std::max(4, 2 * communicatorSize);

        constexpr int yElementCount = 2;

        mfem::Mesh mesh             = mfem::Mesh::MakeCartesian2D(
            xElementCount, yElementCount, mfem::Element::QUADRILATERAL, true, static_cast<double>(xElementCount),
            static_cast<double>(yElementCount)
        );

        for (int elementId = 0; elementId < mesh.GetNE(); ++elementId) {
            const int xIndex = elementId % xElementCount;

            mesh.GetElement(elementId)->SetAttribute(xIndex < xElementCount / 2 ? stellarAttribute : vacuumAttribute);
        }

        mesh.SetAttributes();

        return mesh;
    }

    [[nodiscard]]
    long long global_sum(const int localValue) {
        const long long local = static_cast<long long>(localValue);

        long long global      = 0;

        MPI_Allreduce(&local, &global, 1, MPI_LONG_LONG, MPI_SUM, MPI_COMM_WORLD);

        return global;
    }

    template <typename FieldT>
    concept CanMakeFieldDofMap =
        requires(const mfem::ParFiniteElementSpace &space) { field::make_field_dof_map<FieldT, Schema>(space); };

    using AlternateSchema = domain::DomainSchema<
        domain::MaterialList<
            domain::Material<domain::Core, 11>,
            domain::Material<domain::Envelope, 17>,
            domain::Material<domain::Vacuum, 29>>,
        domain::BoundaryList<>,
        domain::RelationList<>>;
} // namespace field_dof_map_test_utils

TEST_CASE(
    "Field DOF Map Preserves Canonical Bidirectional Indexing",
    tags::unit &tags::field
) {
    namespace field               = mean_field::field;

    const mfem::Array<int> active = field_dof_map_test_utils::make_array({0, 2, 5, 7});

    const field::FieldDofMap map(9, active);

    CHECK(map.full_size() == 9);

    CHECK(map.reduced_size() == 4);

    CHECK(map.inactive_size() == 5);

    CHECK_FALSE(map.is_identity());

    CHECK(map.true_dof(0) == 0);

    CHECK(map.true_dof(1) == 2);

    CHECK(map.true_dof(2) == 5);

    CHECK(map.true_dof(3) == 7);

    REQUIRE(map.reduced_dof(0).has_value());

    REQUIRE(map.reduced_dof(2).has_value());

    REQUIRE(map.reduced_dof(5).has_value());

    REQUIRE(map.reduced_dof(7).has_value());

    CHECK(*map.reduced_dof(0) == 0);

    CHECK(*map.reduced_dof(2) == 1);

    CHECK(*map.reduced_dof(5) == 2);

    CHECK(*map.reduced_dof(7) == 3);

    CHECK_FALSE(map.reduced_dof(1).has_value());

    CHECK_FALSE(map.reduced_dof(3).has_value());

    CHECK(map.contains_true_dof(0));

    CHECK(map.contains_true_dof(2));

    CHECK_FALSE(map.contains_true_dof(1));

    const mfem::Array<int> &forward = map.reduced_to_true();

    const mfem::Array<int> &inverse = map.true_to_reduced();

    REQUIRE(forward.Size() == 4);

    REQUIRE(inverse.Size() == 9);

    CHECK(forward[0] == 0);
    CHECK(forward[1] == 2);
    CHECK(forward[2] == 5);
    CHECK(forward[3] == 7);

    CHECK(inverse[0] == 0);
    CHECK(inverse[1] == -1);
    CHECK(inverse[2] == 1);
    CHECK(inverse[3] == -1);
    CHECK(inverse[4] == -1);
    CHECK(inverse[5] == 2);
    CHECK(inverse[6] == -1);
    CHECK(inverse[7] == 3);
    CHECK(inverse[8] == -1);
}

TEST_CASE(
    "Field DOF Map Rejects Invalid Canonical Mappings",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const mfem::Array<int> empty;

    CHECK_THROWS_AS((field::FieldDofMap(-1, empty)), std::invalid_argument);

    CHECK_THROWS_AS((field::FieldDofMap(4, field_dof_map_test_utils::make_array({-1, 2}))), std::invalid_argument);

    CHECK_THROWS_AS((field::FieldDofMap(4, field_dof_map_test_utils::make_array({1, 4}))), std::invalid_argument);

    /*
     * Duplicate true DOF.
     */
    CHECK_THROWS_AS((field::FieldDofMap(5, field_dof_map_test_utils::make_array({1, 1, 3}))), std::invalid_argument);

    /*
     * Non-canonical unsorted ordering.
     */
    CHECK_THROWS_AS((field::FieldDofMap(5, field_dof_map_test_utils::make_array({1, 3, 2}))), std::invalid_argument);
}

TEST_CASE(
    "Field DOF Map Rejects Out Of Range Index Queries",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(5, field_dof_map_test_utils::make_array({1, 3}));

    CHECK_THROWS_AS(map.true_dof(-1), std::out_of_range);

    CHECK_THROWS_AS(map.true_dof(2), std::out_of_range);

    CHECK_THROWS_AS(map.reduced_dof(-1), std::out_of_range);

    CHECK_THROWS_AS(map.reduced_dof(5), std::out_of_range);

    CHECK_THROWS_AS(map.contains_true_dof(-1), std::out_of_range);

    CHECK_THROWS_AS(map.contains_true_dof(5), std::out_of_range);
}

TEST_CASE(
    "Field DOF Map Gather Selects Exactly The Active True DOFs",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(6, field_dof_map_test_utils::make_array({1, 3, 5}));

    mfem::Vector full(6);

    for (int trueDof = 0; trueDof < full.Size(); ++trueDof) {
        full(trueDof) = 10.0 + static_cast<double>(trueDof);
    }

    const mfem::Vector reduced = map.gather(full);

    REQUIRE(reduced.Size() == 3);

    CHECK(reduced(0) == 11.0);

    CHECK(reduced(1) == 13.0);

    CHECK(reduced(2) == 15.0);

    mfem::Vector output(3);

    map.gather(full, output);

    CHECK(output(0) == 11.0);

    CHECK(output(1) == 13.0);

    CHECK(output(2) == 15.0);
}

TEST_CASE(
    "Field DOF Map Scatter Produces The Canonical Supported Projection",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(6, field_dof_map_test_utils::make_array({1, 3, 5}));

    mfem::Vector reduced(3);

    reduced(0)              = 2.0;
    reduced(1)              = 4.0;
    reduced(2)              = 6.0;

    const mfem::Vector full = map.scatter(reduced);

    REQUIRE(full.Size() == 6);

    CHECK(full(0) == 0.0);
    CHECK(full(1) == 2.0);
    CHECK(full(2) == 0.0);
    CHECK(full(3) == 4.0);
    CHECK(full(4) == 0.0);
    CHECK(full(5) == 6.0);

    const mfem::Vector roundTrip = map.gather(full);

    REQUIRE(roundTrip.Size() == reduced.Size());

    for (int index = 0; index < reduced.Size(); ++index) {
        CHECK(roundTrip(index) == reduced(index));
    }
}

TEST_CASE(
    "Field DOF Map Gather Scatter Projects A Full Vector Onto Field Support",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(7, field_dof_map_test_utils::make_array({0, 2, 3, 6}));

    mfem::Vector original(7);

    for (int index = 0; index < original.Size(); ++index) {
        original(index) = 0.25 + static_cast<double>(index);
    }

    const mfem::Vector reduced   = map.gather(original);

    const mfem::Vector projected = map.scatter(reduced);

    for (int trueDof = 0; trueDof < original.Size(); ++trueDof) {
        CAPTURE(trueDof);

        if (map.contains_true_dof(trueDof)) {
            CHECK(projected(trueDof) == original(trueDof));
        } else {
            CHECK(projected(trueDof) == 0.0);
        }
    }
}

TEST_CASE(
    "Field DOF Map Scatter Into Preserves Unsupported True DOFs",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(6, field_dof_map_test_utils::make_array({1, 4}));

    mfem::Vector reduced(2);

    reduced(0) = 7.0;
    reduced(1) = 9.0;

    mfem::Vector full(6);

    full = -3.0;

    map.scatter_into(reduced, full);

    CHECK(full(0) == -3.0);
    CHECK(full(1) == 7.0);
    CHECK(full(2) == -3.0);
    CHECK(full(3) == -3.0);
    CHECK(full(4) == 9.0);
    CHECK(full(5) == -3.0);
}

TEST_CASE(
    "Field DOF Map Scatter Add Accumulates Only Onto Active True DOFs",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(5, field_dof_map_test_utils::make_array({0, 2, 4}));

    mfem::Vector reduced(3);

    reduced(0) = 1.0;
    reduced(1) = 2.0;
    reduced(2) = 3.0;

    mfem::Vector full(5);

    full = 10.0;

    map.scatter_add(reduced, full, 2.0);

    CHECK(full(0) == 12.0);
    CHECK(full(1) == 10.0);
    CHECK(full(2) == 14.0);
    CHECK(full(3) == 10.0);
    CHECK(full(4) == 16.0);
}

TEST_CASE(
    "Field DOF Map Operations Support MFEM Vector Views Without Resizing",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(5, field_dof_map_test_utils::make_array({1, 3}));

    mfem::Vector storage(9);

    storage = -8.0;

    /*
     * View [2, 7) of the parent vector.
     */
    mfem::Vector fullView(storage.GetData() + 2, 5);

    mfem::Vector reduced(2);

    reduced(0) = 4.0;
    reduced(1) = 6.0;

    map.scatter_into(reduced, fullView);

    /*
     * Storage outside the view must remain untouched.
     */
    CHECK(storage(0) == -8.0);
    CHECK(storage(1) == -8.0);
    CHECK(storage(7) == -8.0);
    CHECK(storage(8) == -8.0);

    /*
     * Within the view, only active true DOFs change.
     */
    CHECK(storage(2) == -8.0);
    CHECK(storage(3) == 4.0);
    CHECK(storage(4) == -8.0);
    CHECK(storage(5) == 6.0);
    CHECK(storage(6) == -8.0);
}

TEST_CASE(
    "Field DOF Map Operations Reject Incompatible Vector Sizes",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(5, field_dof_map_test_utils::make_array({1, 3}));

    mfem::Vector correctFull(5);
    mfem::Vector wrongFull(4);

    mfem::Vector correctReduced(2);
    mfem::Vector wrongReduced(3);

    CHECK_THROWS_AS(map.gather(wrongFull), std::invalid_argument);

    CHECK_THROWS_AS(map.gather(correctFull, wrongReduced), std::invalid_argument);

    CHECK_THROWS_AS(map.scatter(wrongReduced), std::invalid_argument);

    CHECK_THROWS_AS(map.scatter(correctReduced, wrongFull), std::invalid_argument);

    CHECK_THROWS_AS(map.scatter_into(wrongReduced, correctFull), std::invalid_argument);

    CHECK_THROWS_AS(map.scatter_add(correctReduced, wrongFull), std::invalid_argument);
}

TEST_CASE(
    "Field DOF Map Identity Mapping Is An Exact Vector Identity",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    const field::FieldDofMap map(4, field_dof_map_test_utils::make_array({0, 1, 2, 3}));

    REQUIRE(map.is_identity());

    REQUIRE(map.inactive_size() == 0);

    mfem::Vector full(4);

    full(0)                     = 0.1;
    full(1)                     = -0.2;
    full(2)                     = 3.7;
    full(3)                     = 8.1;

    const mfem::Vector reduced  = map.gather(full);

    const mfem::Vector restored = map.scatter(reduced);

    for (int index = 0; index < full.Size(); ++index) {
        CHECK(reduced(index) == full(index));

        CHECK(restored(index) == full(index));
    }
}

TEST_CASE(
    "Field DOF Map Validates Field DOF Support Consistency",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    field::FieldDofSupport support;

    support.activeTrueDofMarker.SetSize(5);

    support.activeTrueDofMarker    = 0;

    support.activeTrueDofMarker[1] = 1;

    support.activeTrueDofMarker[3] = 1;

    support.activeTrueDofs         = field_dof_map_test_utils::make_array({1, 3});

    const field::FieldDofMap validMap(support);

    CHECK(validMap.full_size() == 5);

    CHECK(validMap.reduced_size() == 2);

    /*
     * Make the marker disagree with the list.
     */
    support.activeTrueDofMarker[3] = 0;

    CHECK_THROWS_AS((field::FieldDofMap(support)), std::invalid_argument);
}

TEST_CASE(
    "Field DOF Map Factory Is Available Only For Spatial Registered Fields",
    tags::unit &tags::field
) {
    namespace field = mean_field::field;

    STATIC_REQUIRE(field_dof_map_test_utils::CanMakeFieldDofMap<field::Density>);

    STATIC_REQUIRE(field_dof_map_test_utils::CanMakeFieldDofMap<field::Enthalpy>);

    STATIC_REQUIRE(field_dof_map_test_utils::CanMakeFieldDofMap<field::Gravity>);

    STATIC_REQUIRE(field_dof_map_test_utils::CanMakeFieldDofMap<field::Displacement>);

    STATIC_REQUIRE_FALSE(field_dof_map_test_utils::CanMakeFieldDofMap<field::BarotropicConstant>);

    CHECK(true);
}

TEST_CASE(
    "Field DOF Map Factory Exactly Preserves Density Support",
    tags::integration &tags::field
) {
    namespace field       = mean_field::field;

    mfem::Mesh serialMesh = field_dof_map_test_utils::make_split_mesh();

    mfem::ParMesh mesh(MPI_COMM_WORLD, serialMesh);

    auto fec                = field::Field<field::Density>::make_fec<field::Density::Scalar>(2);

    auto finiteElementSpace = field::Field<field::Density>::make_fespace<field::Density::Scalar>(mesh, *fec);

    REQUIRE(finiteElementSpace != nullptr);

    const field::FieldDofSupport support =
        field::resolve_field_dof_support<field::Density, field_dof_map_test_utils::Schema>(*finiteElementSpace);

    const field::FieldDofMap map =
        field::make_field_dof_map<field::Density, field_dof_map_test_utils::Schema>(*finiteElementSpace);

    REQUIRE(map.full_size() == finiteElementSpace->GetTrueVSize());

    REQUIRE(map.reduced_size() == support.activeTrueDofs.Size());

    REQUIRE(map.full_size() == support.activeTrueDofMarker.Size());

    for (int reducedDof = 0; reducedDof < map.reduced_size(); ++reducedDof) {
        CAPTURE(reducedDof);

        CHECK(map.true_dof(reducedDof) == support.activeTrueDofs[reducedDof]);
    }

    for (int trueDof = 0; trueDof < map.full_size(); ++trueDof) {
        CAPTURE(trueDof);

        CHECK(map.contains_true_dof(trueDof) == (support.activeTrueDofMarker[trueDof] != 0));
    }

    const long long globalFullSize    = field_dof_map_test_utils::global_sum(map.full_size());

    const long long globalReducedSize = field_dof_map_test_utils::global_sum(map.reduced_size());

    /*
     * L2 density has independent vacuum element DOFs, so removing vacuum
     * support must genuinely reduce the global nonlinear block.
     */
    CHECK(globalReducedSize > 0);

    CHECK(globalReducedSize < globalFullSize);
}

TEST_CASE(
    "Field DOF Map Factory Exactly Preserves H1 Enthalpy Support",
    tags::integration &tags::field
) {
    namespace field       = mean_field::field;

    mfem::Mesh serialMesh = field_dof_map_test_utils::make_split_mesh();

    mfem::ParMesh mesh(MPI_COMM_WORLD, serialMesh);

    auto fec                = field::Field<field::Enthalpy>::make_fec<field::Enthalpy::Scalar>(2);

    auto finiteElementSpace = field::Field<field::Enthalpy>::make_fespace<field::Enthalpy::Scalar>(mesh, *fec);

    REQUIRE(finiteElementSpace != nullptr);

    const field::FieldDofSupport support =
        field::resolve_field_dof_support<field::Enthalpy, field_dof_map_test_utils::Schema>(*finiteElementSpace);

    const field::FieldDofMap map =
        field::make_field_dof_map<field::Enthalpy, field_dof_map_test_utils::Schema>(*finiteElementSpace);

    REQUIRE(map.reduced_size() == support.activeTrueDofs.Size());

    for (int reducedDof = 0; reducedDof < map.reduced_size(); ++reducedDof) {
        CHECK(map.true_dof(reducedDof) == support.activeTrueDofs[reducedDof]);
    }

    /*
     * The separate field_mfem support tests already establish that shared
     * Stellar/Vacuum H1 trace DOFs are active. This test establishes that
     * FieldDofMap preserves that active set exactly, rather than applying
     * a second reduction or reinterpretation.
     */
    for (int trueDof = 0; trueDof < map.full_size(); ++trueDof) {
        CHECK(map.contains_true_dof(trueDof) == (support.activeTrueDofMarker[trueDof] != 0));
    }

    const long long globalFullSize    = field_dof_map_test_utils::global_sum(map.full_size());

    const long long globalReducedSize = field_dof_map_test_utils::global_sum(map.reduced_size());

    CHECK(globalReducedSize > 0);

    CHECK(globalReducedSize < globalFullSize);
}

TEST_CASE(
    "Field DOF Map Factory Produces Identity Maps For All Supported Fields",
    tags::integration &tags::field
) {
    namespace field       = mean_field::field;

    mfem::Mesh serialMesh = field_dof_map_test_utils::make_split_mesh();

    mfem::ParMesh mesh(MPI_COMM_WORLD, serialMesh);

    auto fec                = field::Field<field::Displacement>::make_fec<field::Displacement::Vector>(2);

    auto finiteElementSpace = field::Field<field::Displacement>::make_fespace<field::Displacement::Vector>(mesh, *fec);

    REQUIRE(finiteElementSpace != nullptr);

    const field::FieldDofMap map =
        field::make_field_dof_map<field::Displacement, field_dof_map_test_utils::Schema>(*finiteElementSpace);

    CHECK(map.is_identity());

    CHECK(map.full_size() == finiteElementSpace->GetTrueVSize());

    CHECK(map.reduced_size() == finiteElementSpace->GetTrueVSize());

    CHECK(map.inactive_size() == 0);

    for (int trueDof = 0; trueDof < map.full_size(); ++trueDof) {
        CHECK(map.true_dof(trueDof) == trueDof);

        CHECK(map.contains_true_dof(trueDof));
    }
}

TEST_CASE(
    "Field DOF Map Factory Uses Schema Material Bindings Rather Than Numeric Conventions",
    tags::integration &tags::field
) {
    namespace field       = mean_field::field;

    mfem::Mesh serialMesh = field_dof_map_test_utils::make_split_mesh(17, 29);

    mfem::ParMesh mesh(MPI_COMM_WORLD, serialMesh);

    auto fec                = field::Field<field::Density>::make_fec<field::Density::Scalar>(2);

    auto finiteElementSpace = field::Field<field::Density>::make_fespace<field::Density::Scalar>(mesh, *fec);

    REQUIRE(finiteElementSpace != nullptr);

    const field::FieldDofMap map =
        field::make_field_dof_map<field::Density, field_dof_map_test_utils::AlternateSchema>(*finiteElementSpace);

    const field::FieldDofSupport support =
        field::resolve_field_dof_support<field::Density, field_dof_map_test_utils::AlternateSchema>(
            *finiteElementSpace
        );

    CHECK(map.full_size() == support.activeTrueDofMarker.Size());

    CHECK(map.reduced_size() == support.activeTrueDofs.Size());

    for (int trueDof = 0; trueDof < map.full_size(); ++trueDof) {
        CHECK(map.contains_true_dof(trueDof) == (support.activeTrueDofMarker[trueDof] != 0));
    }
}

TEST_CASE(
    "Field DOF Map Reduced Vectors Round Trip Through Real Field Support",
    tags::integration &tags::field
) {
    namespace field       = mean_field::field;

    mfem::Mesh serialMesh = field_dof_map_test_utils::make_split_mesh();

    mfem::ParMesh mesh(MPI_COMM_WORLD, serialMesh);

    auto fec                = field::Field<field::Enthalpy>::make_fec<field::Enthalpy::Scalar>(2);

    auto finiteElementSpace = field::Field<field::Enthalpy>::make_fespace<field::Enthalpy::Scalar>(mesh, *fec);

    REQUIRE(finiteElementSpace != nullptr);

    const field::FieldDofMap map =
        field::make_field_dof_map<field::Enthalpy, field_dof_map_test_utils::Schema>(*finiteElementSpace);

    mfem::Vector reduced(map.reduced_size());

    for (int reducedDof = 0; reducedDof < reduced.Size(); ++reducedDof) {
        reduced(reducedDof) = 0.125 + 0.031 * static_cast<double>(reducedDof + 1);
    }

    const mfem::Vector full      = map.scatter(reduced);

    const mfem::Vector recovered = map.gather(full);

    REQUIRE(recovered.Size() == reduced.Size());

    for (int reducedDof = 0; reducedDof < reduced.Size(); ++reducedDof) {
        CAPTURE(reducedDof);

        CHECK(recovered(reducedDof) == reduced(reducedDof));
    }

    for (int trueDof = 0; trueDof < full.Size(); ++trueDof) {
        if (!map.contains_true_dof(trueDof)) {
            CHECK(full(trueDof) == 0.0);
        }
    }
}