#include <algorithm>
#include <numeric>
#include <optional>
#include <stdexcept>
#include <vector>

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>
#include <mpi.h>

import mean_field;
import test_helpers;

namespace surface_scalar_dof_test_utils {
    namespace domain          = mean_field::utils::domain;
    namespace field           = mean_field::field;

    using DefaultSchema       = domain::CoreEnvelopeVacuumDomainSchema;

    using FirstBoundarySchema = domain::DomainSchema<
        domain::MaterialList<>,
        domain::BoundaryList<domain::BoundaryAttribute<domain::StellarSurface, 1>>,
        domain::RelationList<>>;

    using SecondBoundarySchema = domain::DomainSchema<
        domain::MaterialList<>,
        domain::BoundaryList<domain::BoundaryAttribute<domain::StellarSurface, 2>>,
        domain::RelationList<>>;

    using MissingBoundarySchema =
        domain::DomainSchema<domain::MaterialList<>, domain::BoundaryList<>, domain::RelationList<>>;

    [[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::Array<int> expected_boundary_true_dofs(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const int boundaryAttribute
    ) {
        const mfem::Mesh *mesh = finiteElementSpace.GetMesh();
        REQUIRE(mesh != nullptr);
        REQUIRE(boundaryAttribute > 0);
        REQUIRE(boundaryAttribute <= mesh->bdr_attributes.Max());

        mfem::Array<int> boundaryMarker(mesh->bdr_attributes.Max());
        boundaryMarker                        = 0;
        boundaryMarker[boundaryAttribute - 1] = 1;

        mfem::Array<int> expected;
        finiteElementSpace.GetEssentialTrueDofs(boundaryMarker, expected);
        return expected;
    }

    void check_equal(
        const mfem::Array<int> &actual,
        const mfem::Array<int> &expected
    ) {
        REQUIRE(actual.Size() == expected.Size());
        for (int index = 0; index < actual.Size(); ++index) {
            CAPTURE(index);
            CHECK(actual[index] == expected[index]);
        }
    }

    [[nodiscard]] mfem::Mesh make_boundary_mesh() {
        return mfem::Mesh::MakeCartesian2D(6, 4, mfem::Element::QUADRILATERAL, true, 3.0, 2.0);
    }

    template <typename SchemaT>
    concept CanMakeStellarSurfaceMap = requires(const mfem::ParFiniteElementSpace &finiteElementSpace) {
        field::make_stellar_surface_scalar_dof_map<SchemaT>(finiteElementSpace);
    };
} // namespace surface_scalar_dof_test_utils

TEST_CASE(
    "Scalar Boundary DOF Map Preserves Canonical Surface Coordinate Indexing",
    tags::surface_deformation_dof_unit
) {
    namespace field = mean_field::field;

    const field::ScalarBoundaryDofMap map(8, surface_scalar_dof_test_utils::make_array({1, 3, 6}), 5, 12);

    CHECK(map.volume_true_dof_size() == 8);
    CHECK(map.local_size() == 3);
    CHECK(map.global_offset() == 5);
    CHECK(map.global_size() == 12);
    CHECK_FALSE(map.empty());

    CHECK(map.volume_true_dof(0) == 1);
    CHECK(map.volume_true_dof(1) == 3);
    CHECK(map.volume_true_dof(2) == 6);
    CHECK(map.global_boundary_dof(0) == 5);
    CHECK(map.global_boundary_dof(1) == 6);
    CHECK(map.global_boundary_dof(2) == 7);

    REQUIRE(map.local_boundary_dof(1).has_value());
    REQUIRE(map.local_boundary_dof(3).has_value());
    REQUIRE(map.local_boundary_dof(6).has_value());
    CHECK(*map.local_boundary_dof(1) == 0);
    CHECK(*map.local_boundary_dof(3) == 1);
    CHECK(*map.local_boundary_dof(6) == 2);
    CHECK_FALSE(map.local_boundary_dof(0).has_value());

    mfem::Vector volumeValues(8);
    for (int trueDof = 0; trueDof < volumeValues.Size(); ++trueDof) {
        volumeValues(trueDof) = 10.0 + trueDof;
    }

    const mfem::Vector boundaryValues = map.gather(volumeValues);
    REQUIRE(boundaryValues.Size() == map.local_size());
    CHECK(boundaryValues(0) == 11.0);
    CHECK(boundaryValues(1) == 13.0);
    CHECK(boundaryValues(2) == 16.0);

    const mfem::Vector scattered = map.scatter(boundaryValues);
    REQUIRE(scattered.Size() == map.volume_true_dof_size());
    for (int trueDof = 0; trueDof < scattered.Size(); ++trueDof) {
        const std::optional<int> localBoundaryDof = map.local_boundary_dof(trueDof);
        if (localBoundaryDof.has_value()) {
            CHECK(scattered(trueDof) == boundaryValues(*localBoundaryDof));
        } else {
            CHECK(scattered(trueDof) == 0.0);
        }
    }

    CHECK_THROWS_AS(
        (field::ScalarBoundaryDofMap(8, surface_scalar_dof_test_utils::make_array({3, 1}), 0, 2)), std::invalid_argument
    );
    CHECK_THROWS_AS(
        (field::ScalarBoundaryDofMap(8, surface_scalar_dof_test_utils::make_array({1, 3, 6}), -1, 3)),
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        (field::ScalarBoundaryDofMap(8, surface_scalar_dof_test_utils::make_array({1, 3, 6}), 4, 6)),
        std::invalid_argument
    );
    CHECK_THROWS_AS(map.global_boundary_dof(-1), std::out_of_range);
    CHECK_THROWS_AS(map.global_boundary_dof(map.local_size()), std::out_of_range);
}

TEST_CASE(
    "Stellar Surface Coordinates Use The Scalar Displacement Basis And Only Surface DOFs",
    tags::surface_deformation_dof_topology
) {
    namespace domain             = mean_field::utils::domain;
    namespace field              = mean_field::field;

    mean_field::utils::Args args = test_utils::setup_args();
    mean_field::fem::FEM fem     = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(fem.okay());

    REQUIRE(fem.surfaceDeformationFes != nullptr);
    REQUIRE(fem.displacementFes != nullptr);
    CHECK(fem.surfaceDeformationFes->GetVDim() == 1);
    CHECK(fem.surfaceDeformationFes->FEColl() == fem.displacementFes->FEColl());
    CHECK(fem.surfaceDeformationFes.get() != fem.enthalpyFes.get());

    const field::ScalarBoundaryDofMap surfaceCoordinates =
        field::make_stellar_surface_scalar_dof_map<surface_scalar_dof_test_utils::DefaultSchema>(
            *fem.surfaceDeformationFes
        );

    const mfem::Array<int> expected = surface_scalar_dof_test_utils::expected_boundary_true_dofs(
        *fem.surfaceDeformationFes,
        surface_scalar_dof_test_utils::DefaultSchema::template boundary_attribute<domain::StellarSurface>()
    );
    surface_scalar_dof_test_utils::check_equal(surfaceCoordinates.boundary_true_dofs(), expected);

    CHECK(surfaceCoordinates.global_size() > 0);
    CHECK(surfaceCoordinates.global_size() < fem.surfaceDeformationFes->GlobalTrueVSize());

    const field::FieldDofMap displacementMap =
        field::make_field_dof_map<field::Displacement, surface_scalar_dof_test_utils::DefaultSchema>(
            *fem.displacementFes
        );
    const field::FieldBoundaryDofMap vectorSurface = field::make_field_boundary_dof_map<
        field::Displacement, domain::StellarSurface, surface_scalar_dof_test_utils::DefaultSchema>(
        *fem.displacementFes, displacementMap
    );

    long long localVectorSurfaceSize  = vectorSurface.size();
    long long globalVectorSurfaceSize = 0;
    MPI_Allreduce(
        &localVectorSurfaceSize, &globalVectorSurfaceSize, 1, MPI_LONG_LONG, MPI_SUM,
        fem.surfaceDeformationFes->GetComm()
    );

    CHECK(
        globalVectorSurfaceSize == static_cast<long long>(fem.mesh->SpaceDimension()) * surfaceCoordinates.global_size()
    );
}

TEST_CASE(
    "Scalar Boundary DOF Resolution Uses Semantic Schema Boundary Attributes",
    tags::surface_deformation_dof_schema
) {
    namespace domain = mean_field::utils::domain;
    namespace field  = mean_field::field;

    STATIC_CHECK(
        surface_scalar_dof_test_utils::CanMakeStellarSurfaceMap<surface_scalar_dof_test_utils::FirstBoundarySchema>
    );
    STATIC_CHECK_FALSE(
        surface_scalar_dof_test_utils::CanMakeStellarSurfaceMap<surface_scalar_dof_test_utils::MissingBoundarySchema>
    );

    mfem::Mesh serialMesh = surface_scalar_dof_test_utils::make_boundary_mesh();
    mfem::ParMesh mesh(MPI_COMM_WORLD, serialMesh);
    mfem::H1_FECollection finiteElementCollection(2, mesh.Dimension());
    mfem::ParFiniteElementSpace finiteElementSpace(&mesh, &finiteElementCollection);

    const field::ScalarBoundaryDofMap firstBoundary =
        field::make_stellar_surface_scalar_dof_map<surface_scalar_dof_test_utils::FirstBoundarySchema>(
            finiteElementSpace
        );
    const field::ScalarBoundaryDofMap secondBoundary =
        field::make_stellar_surface_scalar_dof_map<surface_scalar_dof_test_utils::SecondBoundarySchema>(
            finiteElementSpace
        );

    const mfem::Array<int> expectedFirst = surface_scalar_dof_test_utils::expected_boundary_true_dofs(
        finiteElementSpace,
        surface_scalar_dof_test_utils::FirstBoundarySchema::template boundary_attribute<domain::StellarSurface>()
    );
    const mfem::Array<int> expectedSecond = surface_scalar_dof_test_utils::expected_boundary_true_dofs(
        finiteElementSpace,
        surface_scalar_dof_test_utils::SecondBoundarySchema::template boundary_attribute<domain::StellarSurface>()
    );

    surface_scalar_dof_test_utils::check_equal(firstBoundary.boundary_true_dofs(), expectedFirst);
    surface_scalar_dof_test_utils::check_equal(secondBoundary.boundary_true_dofs(), expectedSecond);

    bool localMapsDiffer = firstBoundary.local_size() != secondBoundary.local_size();
    if (!localMapsDiffer) {
        for (int localDof = 0; localDof < firstBoundary.local_size(); ++localDof) {
            if (firstBoundary.volume_true_dof(localDof) != secondBoundary.volume_true_dof(localDof)) {
                localMapsDiffer = true;
                break;
            }
        }
    }

    int localDifference  = localMapsDiffer ? 1 : 0;
    int globalDifference = 0;
    MPI_Allreduce(&localDifference, &globalDifference, 1, MPI_INT, MPI_MAX, finiteElementSpace.GetComm());
    CHECK(globalDifference == 1);
}

TEST_CASE(
    "Scalar Boundary Coordinates Have Deterministic Contiguous Parallel Ordering",
    tags::surface_deformation_dof_parallel
) {
    namespace field       = mean_field::field;

    mfem::Mesh serialMesh = surface_scalar_dof_test_utils::make_boundary_mesh();
    mfem::ParMesh mesh(MPI_COMM_WORLD, serialMesh);
    mfem::H1_FECollection finiteElementCollection(2, mesh.Dimension());
    mfem::ParFiniteElementSpace finiteElementSpace(&mesh, &finiteElementCollection);

    const field::ScalarBoundaryDofMap first =
        field::make_stellar_surface_scalar_dof_map<surface_scalar_dof_test_utils::FirstBoundarySchema>(
            finiteElementSpace
        );
    const field::ScalarBoundaryDofMap second =
        field::make_stellar_surface_scalar_dof_map<surface_scalar_dof_test_utils::FirstBoundarySchema>(
            finiteElementSpace
        );

    surface_scalar_dof_test_utils::check_equal(first.boundary_true_dofs(), second.boundary_true_dofs());
    CHECK(first.global_offset() == second.global_offset());
    CHECK(first.global_size() == second.global_size());

    for (int localDof = 0; localDof < first.local_size(); ++localDof) {
        CAPTURE(localDof);
        CHECK(first.global_boundary_dof(localDof) == first.global_offset() + localDof);
        if (localDof > 0) {
            CHECK(first.volume_true_dof(localDof - 1) < first.volume_true_dof(localDof));
        }
    }

    int communicatorSize = 1;
    int communicatorRank = 0;
    MPI_Comm_size(finiteElementSpace.GetComm(), &communicatorSize);
    MPI_Comm_rank(finiteElementSpace.GetComm(), &communicatorRank);

    const long long localSize = first.local_size();
    std::vector<long long> localSizes(static_cast<std::size_t>(communicatorSize));
    MPI_Allgather(&localSize, 1, MPI_LONG_LONG, localSizes.data(), 1, MPI_LONG_LONG, finiteElementSpace.GetComm());

    const long long expectedOffset = std::accumulate(localSizes.begin(), localSizes.begin() + communicatorRank, 0LL);
    const long long expectedGlobalSize = std::accumulate(localSizes.begin(), localSizes.end(), 0LL);

    CHECK(first.global_offset() == expectedOffset);
    CHECK(first.global_size() == expectedGlobalSize);
}