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

import mean_field;

namespace {
    using mean_field::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_inv, dimension, offset + 7);
        context.mapping.mapping_determinant = offset + 8;
        context.mapping.compactified        = compactified;
        context.quadrature.detJ             = offset + 9;
        context.quadrature.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_inv, expected.quadrature.J_inv);
        CHECK(actual.mapping.mapping_determinant == expected.mapping.mapping_determinant);
        CHECK(actual.mapping.compactified == expected.mapping.compactified);
        CHECK(actual.quadrature.detJ == expected.quadrature.detJ);
        CHECK(actual.quadrature.weight == expected.quadrature.weight);
    }
} // namespace

TEST_CASE(
    "Flat Volume Mapping Cache Preserves Every Context Field",
    "[mapping][prepared-cache]"
) {
    for (const int dimension : {1, 2, 3}) {
        CAPTURE(dimension);
        mean_field::mapping::VolumeMappingCache cache;
        cache.SetSize(2, dimension);
        CHECK(cache.GetPointCount() == 2);
        CHECK(cache.GetDimension() == dimension);
        const auto first  = make_context(dimension, 0.125, false);
        const 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_inv.HostRead();
        const double *position_buffer = workspace.mapping.physical_position.HostRead();
        cache.Load(0, workspace);
        check_context(workspace, first);
        CHECK(workspace.quadrature.J_inv.HostRead() == inverse_buffer);
        CHECK(workspace.mapping.physical_position.HostRead() == position_buffer);

        mfem::DenseMatrix inverse;
        cache.LoadInverseJacobian(1, inverse);
        check_matrix(inverse, second.quadrature.J_inv);
        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.SetSize(1, 3);
        const 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",
    "[mapping][prepared-cache]"
) {
    mean_field::mapping::VolumeMappingCache cache;
    VolumeMappingContext workspace;
    CHECK_THROWS_AS(cache.SetSize(-1, 3), std::invalid_argument);
    CHECK_THROWS_AS(cache.SetSize(1, 0), std::invalid_argument);
    CHECK_THROWS_AS(cache.SetSize(1, 4), std::invalid_argument);
    cache.SetSize(1, 3);
    CHECK_THROWS_AS(cache.Load(-1, workspace), std::out_of_range);
    CHECK_THROWS_AS(cache.Load(1, workspace), std::out_of_range);
    CHECK_THROWS_AS(cache.Store(0, make_context(2, 0.0, false)), std::invalid_argument);
    cache.SetSize(0, 2);
    CHECK(cache.GetPointCount() == 0);
    CHECK_THROWS_AS(cache.Load(0, workspace), std::out_of_range);
}