// Ported from MeanField tests/mapping/prepared_cache.cpp.

#include <stdexcept>

#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>

#include "serif/backend/mfem/mapping/volume_cache.hpp"

#include "serif/tests/test_tags.hpp"

namespace {
    using serif::mapping::VolumeMappingContext;

    void fill_vector(mfem::Vector& vector, const int dimension, const double offset) {
        vector.SetSize(dimension);
        for (int i = 0; i < dimension; ++i) {
            vector(i) = offset + i;
        }
    }

    void fill_matrix(mfem::DenseMatrix& matrix, const int dimension, const double offset) {
        matrix.SetSize(dimension);
        for (int j = 0; j < dimension; ++j) {
            for (int i = 0; i < dimension; ++i) {
                matrix(i, j) = offset + 10 * j + i;
            }
        }
    }

    VolumeMappingContext make_context(const int dimension, const double offset, const bool compactified) {
        VolumeMappingContext context;
        fill_vector(context.mapping.reference_position, dimension, offset + 1);
        fill_vector(context.mapping.displaced_position, dimension, offset + 2);
        fill_vector(context.mapping.physical_position, dimension, offset + 3);
        fill_matrix(context.mapping.displacement_jacobian, dimension, offset + 4);
        fill_matrix(context.mapping.mapping_jacobian, dimension, offset + 5);
        fill_matrix(context.mapping.inverse_mapping_jacobian, dimension, offset + 6);
        fill_matrix(context.quadrature.J_element_to_physical_inverse, dimension, offset + 7);
        context.mapping.mapping_determinate = offset + 8;
        context.mapping.compactified = compactified;
        context.quadrature.mapping_determinate = offset + 9;
        context.quadrature.physical_integration_weight = offset + 10;
        return context;
    }

    void check_vector(const mfem::Vector& actual, const mfem::Vector& expected) {
        REQUIRE(actual.Size() == expected.Size());
        for (int i = 0; i < expected.Size(); ++i) {
            CHECK(actual(i) == expected(i));
        }
    }

    void check_matrix(const mfem::DenseMatrix& actual, const mfem::DenseMatrix& expected) {
        REQUIRE(actual.Height() == expected.Height());
        REQUIRE(actual.Width() == expected.Width());
        for (int j = 0; j < expected.Width(); ++j) {
            for (int i = 0; i < expected.Height(); ++i) {
                CHECK(actual(i, j) == expected(i, j));
            }
        }
    }

    void check_context(const VolumeMappingContext& actual, const VolumeMappingContext& expected) {
        check_vector(actual.mapping.reference_position, expected.mapping.reference_position);
        check_vector(actual.mapping.displaced_position, expected.mapping.displaced_position);
        check_vector(actual.mapping.physical_position, expected.mapping.physical_position);
        check_matrix(actual.mapping.displacement_jacobian, expected.mapping.displacement_jacobian);
        check_matrix(actual.mapping.mapping_jacobian, expected.mapping.mapping_jacobian);
        check_matrix(actual.mapping.inverse_mapping_jacobian, expected.mapping.inverse_mapping_jacobian);
        check_matrix(
            actual.quadrature.J_element_to_physical_inverse, expected.quadrature.J_element_to_physical_inverse
        );
        CHECK(actual.mapping.mapping_determinate == expected.mapping.mapping_determinate);
        CHECK(actual.mapping.compactified == expected.mapping.compactified);
        CHECK(actual.quadrature.mapping_determinate == expected.quadrature.mapping_determinate);
        CHECK(actual.quadrature.physical_integration_weight == expected.quadrature.physical_integration_weight);
    }
} // namespace

TEST_CASE(
    "Flat Volume Mapping Cache Preserves Every Context Field",
    tags::unit & tags::volume_cache
) {
    for (const int dimension : {1, 2, 3}) {
        CAPTURE(dimension);
        serif::mapping::VolumeMappingCache cache;
        cache.set_size(2, dimension);
        CHECK(cache.get_point_count() == 2);
        CHECK(cache.get_dimension() == dimension);
        auto first = make_context(dimension, 0.125, false);
        auto second = make_context(dimension, -30.25, true);
        cache.store(0, first);
        cache.store(1, second);

        VolumeMappingContext workspace;
        cache.load(1, workspace);
        check_context(workspace, second);
        const double* inverse_buffer = workspace.quadrature.J_element_to_physical_inverse.HostRead();
        const double* position_buffer = workspace.mapping.physical_position.HostRead();
        cache.load(0, workspace);
        check_context(workspace, first);
        CHECK(workspace.quadrature.J_element_to_physical_inverse.HostRead() == inverse_buffer);
        CHECK(workspace.mapping.physical_position.HostRead() == position_buffer);

        mfem::DenseMatrix inverse;
        cache.load_inverse_jacobian(1, inverse);
        check_matrix(inverse, second.quadrature.J_element_to_physical_inverse);
        const auto copy = cache;
        cache.store(1, first);
        copy.load(1, workspace);
        check_context(workspace, second);
        cache.load(1, workspace);
        check_context(workspace, first);

        cache.set_size(1, 3);
        auto resized = make_context(3, 13.5, true);
        cache.store(0, resized);
        cache.load(0, workspace);
        check_context(workspace, resized);
    }
}

TEST_CASE(
    "Flat Volume Mapping Cache Rejects Invalid Indices And Dimensions",
    tags::unit & tags::volume_cache
) {
    serif::mapping::VolumeMappingCache cache;
    VolumeMappingContext workspace;
    CHECK_THROWS_AS(cache.set_size(-1, 3), std::logic_error);
    CHECK_THROWS_AS(cache.set_size(1, 0), std::logic_error);
    CHECK_THROWS_AS(cache.set_size(1, 4), std::logic_error);
    cache.set_size(1, 3);
    CHECK_THROWS_AS(cache.load(-1, workspace), std::logic_error);
    CHECK_THROWS_AS(cache.load(1, workspace), std::logic_error);
    auto wrong_dimension = make_context(2, 0.0, false);
    CHECK_THROWS_AS(cache.store(0, wrong_dimension), std::logic_error);
    cache.set_size(0, 2);
    CHECK(cache.get_point_count() == 0);
    CHECK_THROWS_AS(cache.load(0, workspace), std::logic_error);
}

TEST_CASE(
    "Flat Volume Mapping Cache Rejects Malformed Packed Fields",
    tags::unit & tags::volume_cache
) {
    serif::mapping::VolumeMappingCache cache;
    cache.set_size(1, 3);
    auto context = make_context(3, 0.0, false);
    for (auto* vector :
        {&context.mapping.reference_position, &context.mapping.displaced_position,
            &context.mapping.physical_position}) {
        for (const int size : {2, 4}) {
            vector->SetSize(size);
            CHECK_THROWS_AS(cache.store(0, context), std::logic_error);
        }
        vector->SetSize(3);
        *vector = 0.0;
    }
    for (auto* matrix :
        {&context.mapping.displacement_jacobian, &context.mapping.mapping_jacobian,
            &context.mapping.inverse_mapping_jacobian, &context.quadrature.J_element_to_physical_inverse}) {
        matrix->SetSize(4, 4);
        CHECK_THROWS_AS(cache.store(0, context), std::logic_error);
        matrix->SetSize(3, 2);
        CHECK_THROWS_AS(cache.store(0, context), std::logic_error);
        matrix->SetSize(3, 3);
        *matrix = 0.0;
    }
    CHECK_NOTHROW(cache.store(0, context));
}