#include <array>
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include <cmath>
#include <limits>
#include <mfem.hpp>
#include <string>
#include <vector>

import mean_field;
import test_helpers;
using namespace mean_field;
using Catch::Matchers::WithinAbs;

namespace {
    constexpr int dimension = 3;

    mfem::DenseMatrix make_matrix(
        const std::array<
            double,
            9> &values
    ) {
        mfem::DenseMatrix matrix(dimension);

        for (int row = 0; row < dimension; ++row) {
            for (int column = 0; column < dimension; ++column)
                matrix(row, column) = values[row * dimension + column];
        }

        return matrix;
    }

    mfem::DenseMatrix make_identity_matrix() {
        mfem::DenseMatrix identity(dimension);
        identity = 0.0;
        for (int i = 0; i < dimension; ++i)
            identity(i, i) = 1.0;
        return identity;
    }

    mapping::MappingPointContext
    make_context(const mfem::DenseMatrix &jacobian) {
        mapping::MappingPointContext context;
        context.mapping_jacobian    = jacobian;
        context.mapping_determinant = jacobian.Det();
        context.inverse_mapping_jacobian.SetSize(dimension);
        mfem::CalcInverse(jacobian, context.inverse_mapping_jacobian);
        context.physical_position.SetSize(dimension);
        context.physical_position = 0.0;
        return context;
    }

    double determinant_variation(
        const mfem::DenseMatrix &jacobian,
        const mfem::DenseMatrix &jacobian_variation
    ) {
        mfem::DenseMatrix inverse_jacobian(dimension);
        mfem::DenseMatrix product(dimension);
        mfem::CalcInverse(jacobian, inverse_jacobian);
        mfem::Mult(inverse_jacobian, jacobian_variation, product);

        double trace = 0.0;
        for (int i = 0; i < dimension; ++i)
            trace += product(i, i);
        return jacobian.Det() * trace;
    }

    mapping::MappingPointVariation make_variation(
        const mfem::DenseMatrix &jacobian,
        const mfem::DenseMatrix &jacobian_variation
    ) {
        mapping::MappingPointVariation variation;
        variation.mapping_jacobian_variation = jacobian_variation;
        variation.mapping_determinant_variation =
            determinant_variation(jacobian, jacobian_variation);
        variation.physical_position_variation.SetSize(dimension);
        variation.physical_position_variation = 0.0;
        return variation;
    }

    double matrix_norm(const mfem::DenseMatrix &matrix) {
        double norm_squared = 0.0;

        for (int row = 0; row < matrix.Height(); ++row) {
            for (int column = 0; column < matrix.Width(); ++column)
                norm_squared += matrix(row, column) * matrix(row, column);
        }

        return std::sqrt(norm_squared);
    }

    double relative_matrix_error(
        const mfem::DenseMatrix &computed,
        const mfem::DenseMatrix &reference
    ) {
        REQUIRE(computed.Height() == reference.Height());
        REQUIRE(computed.Width() == reference.Width());

        mfem::DenseMatrix difference(computed);
        difference -= reference;

        return matrix_norm(difference) /
               std::max(
                   matrix_norm(reference),
                   std::numeric_limits<double>::epsilon()
               );
    }

    double matrix_asymmetry(const mfem::DenseMatrix &matrix) {
        double asymmetry_squared = 0.0;

        for (int row = 0; row < matrix.Height(); ++row) {
            for (int column = 0; column < matrix.Width(); ++column) {
                const double difference =
                    matrix(row, column) - matrix(column, row);
                asymmetry_squared += difference * difference;
            }
        }

        return std::sqrt(asymmetry_squared);
    }

    mfem::DenseMatrix centered_mass_tensor_difference(
        const mfem::DenseMatrix &jacobian,
        const mfem::DenseMatrix &jacobian_variation,
        const double step
    ) {
        mfem::DenseMatrix plus_jacobian(jacobian);
        mfem::DenseMatrix minus_jacobian(jacobian);

        for (int row = 0; row < dimension; ++row) {
            for (int column = 0; column < dimension; ++column) {
                plus_jacobian(row, column) +=
                    step * jacobian_variation(row, column);
                minus_jacobian(row, column) -=
                    step * jacobian_variation(row, column);
            }
        }

        REQUIRE(plus_jacobian.Det() > 0.0);
        REQUIRE(minus_jacobian.Det() > 0.0);

        const mapping::MappingPointContext plus_context =
            make_context(plus_jacobian);
        const mapping::MappingPointContext minus_context =
            make_context(minus_jacobian);

        mfem::DenseMatrix plus_tensor;
        mfem::DenseMatrix minus_tensor;
        mapping::ComputeHDivMassTensor(plus_context, plus_tensor);
        mapping::ComputeHDivMassTensor(minus_context, minus_tensor);

        plus_tensor -= minus_tensor;
        plus_tensor *= 1 / (2.0 * step);
        return plus_tensor;
    }

    void check_zero_matrix(
        const mfem::DenseMatrix &matrix,
        const double tolerance
    ) {
        for (int row = 0; row < matrix.Height(); ++row) {
            for (int column = 0; column < matrix.Width(); ++column)
                CHECK_THAT(matrix(row, column), WithinAbs(0.0, tolerance));
        }
    }

    struct TensorVariationCase {
        std::string name;
        mfem::DenseMatrix jacobian;
        mfem::DenseMatrix jacobian_variation;
    };
} // namespace

TEST_CASE(
    "Hdiv Mass Tensor Variation Matches Centered Differences",
    tags::unit &tags::transformations
) {
    std::vector<TensorVariationCase> cases;

    cases.push_back(
        {"identity with general variation", make_identity_matrix(),
         make_matrix({0.12, -0.07, 0.03, 0.05, -0.09, 0.04, -0.02, 0.08, 0.06})}
    );

    cases.push_back(
        {"anisotropic stretch",
         make_matrix({1.20, 0.00, 0.00, 0.00, 0.85, 0.00, 0.00, 0.00, 1.10}),
         make_matrix({0.08, 0.01, -0.03, 0.02, -0.05, 0.04, 0.01, -0.02, 0.07})}
    );

    cases.push_back(
        {"sheared mapping",
         make_matrix({1.10, 0.20, -0.05, 0.04, 0.90, 0.12, -0.03, 0.08, 1.15}),
         make_matrix(
             {0.06, -0.04, 0.02, 0.03, 0.05, -0.07, -0.01, 0.04, -0.02}
         )}
    );

    cases.push_back(
        {"strong general mapping",
         make_matrix({1.35, 0.31, -0.18, -0.12, 0.78, 0.22, 0.09, -0.16, 1.27}),
         make_matrix({-0.11, 0.08, 0.05, 0.07, 0.09, -0.04, -0.06, 0.03, 0.12})}
    );

    for (const TensorVariationCase &test_case : cases) {
        DYNAMIC_SECTION(test_case.name) {
            REQUIRE(test_case.jacobian.Det() > 0.0);

            const mapping::MappingPointContext context =
                make_context(test_case.jacobian);
            const mapping::MappingPointVariation variation = make_variation(
                test_case.jacobian, test_case.jacobian_variation
            );

            mfem::DenseMatrix analytic_variation;
            mapping::ComputeHDivMassTensorVariation(
                context, variation, analytic_variation
            );

            const mfem::DenseMatrix finite_difference =
                centered_mass_tensor_difference(
                    test_case.jacobian, test_case.jacobian_variation, 1.0e-6
                );
            const double relative_error =
                relative_matrix_error(analytic_variation, finite_difference);
            const double asymmetry = matrix_asymmetry(analytic_variation);

            INFO("Mapping determinant = " << context.mapping_determinant);
            INFO(
                "Determinant variation = "
                << variation.mapping_determinant_variation
            );
            INFO(
                "Analytic variation norm = " << matrix_norm(analytic_variation)
            );
            INFO(
                "Finite-difference variation norm = "
                << matrix_norm(finite_difference)
            );
            INFO("Relative tensor-variation error = " << relative_error);
            INFO("Tensor-variation asymmetry = " << asymmetry);

            CHECK(relative_error < 2.0e-9);
            CHECK(asymmetry < 2.0e-14);
        }
    }
}

TEST_CASE(
    "Hdiv Mass Tensor Variation Has Second Order Centered Difference "
    "Convergence",
    tags::unit &tags::transformations &tags::convergence
) {
    const mfem::DenseMatrix jacobian =
        make_matrix({1.18, 0.17, -0.09, -0.04, 0.92, 0.14, 0.07, -0.11, 1.23});

    const mfem::DenseMatrix jacobian_variation =
        make_matrix({0.09, -0.06, 0.04, 0.03, 0.07, -0.05, -0.02, 0.08, -0.03});

    const mapping::MappingPointContext context = make_context(jacobian);
    const mapping::MappingPointVariation variation =
        make_variation(jacobian, jacobian_variation);

    mfem::DenseMatrix analytic_variation;
    mapping::ComputeHDivMassTensorVariation(
        context, variation, analytic_variation
    );

    const std::array<double, 3> steps{4.0e-2, 2.0e-2, 1.0e-2};
    std::array<double, 3> errors{};

    for (int i = 0; i < static_cast<int>(steps.size()); ++i) {
        const mfem::DenseMatrix finite_difference =
            centered_mass_tensor_difference(
                jacobian, jacobian_variation, steps[i]
            );
        errors[i] =
            relative_matrix_error(finite_difference, analytic_variation);
        INFO("Step = " << steps[i] << ", relative error = " << errors[i]);
    }

    const double first_reduction  = errors[1] / errors[0];
    const double second_reduction = errors[2] / errors[1];

    INFO("First error-reduction ratio = " << first_reduction);
    INFO("Second error-reduction ratio = " << second_reduction);

    CHECK(first_reduction < 0.30);
    CHECK(second_reduction < 0.30);
    CHECK(errors[2] < 1.0e-5);
}

TEST_CASE(
    "Hdiv Mass Tensor Variation Vanishes For Translation",
    tags::unit &tags::transformations
) {
    const mfem::DenseMatrix jacobian =
        make_matrix({1.12, 0.08, -0.03, 0.02, 0.94, 0.07, -0.01, 0.05, 1.09});

    mfem::DenseMatrix zero_jacobian_variation(dimension);
    zero_jacobian_variation              = 0.0;

    mapping::MappingPointContext context = make_context(jacobian);
    mapping::MappingPointVariation variation =
        make_variation(jacobian, zero_jacobian_variation);
    variation.physical_position_variation.SetSize(dimension);
    variation.physical_position_variation(0) = 0.7;
    variation.physical_position_variation(1) = -0.4;
    variation.physical_position_variation(2) = 0.9;

    mfem::DenseMatrix tensor_variation;
    mapping::ComputeHDivMassTensorVariation(
        context, variation, tensor_variation
    );

    CHECK_THAT(variation.mapping_determinant_variation, WithinAbs(0.0, 0.0));
    check_zero_matrix(tensor_variation, 1.0e-14);
}

TEST_CASE(
    "Hdiv Mass Tensor Variation Vanishes For Infinitesimal Rotation At "
    "Identity",
    tags::unit &tags::transformations
) {
    const mfem::DenseMatrix identity = make_identity_matrix();

    const mfem::DenseMatrix rotation_variation =
        make_matrix({0.0, -0.30, 0.20, 0.30, 0.0, -0.15, -0.20, 0.15, 0.0});

    const mapping::MappingPointContext context = make_context(identity);
    const mapping::MappingPointVariation variation =
        make_variation(identity, rotation_variation);

    mfem::DenseMatrix tensor_variation;
    mapping::ComputeHDivMassTensorVariation(
        context, variation, tensor_variation
    );

    CHECK_THAT(
        variation.mapping_determinant_variation, WithinAbs(0.0, 1.0e-15)
    );
    check_zero_matrix(tensor_variation, 1.0e-14);
}

TEST_CASE(
    "Hdiv Mass Tensor Variation Matches Isotropic Scaling At Identity",
    tags::unit &tags::transformations
) {
    constexpr double scaling_variation = 0.17;

    const mfem::DenseMatrix identity   = make_identity_matrix();
    mfem::DenseMatrix jacobian_variation(dimension);
    jacobian_variation = 0.0;
    for (int i = 0; i < dimension; ++i)
        jacobian_variation(i, i) = scaling_variation;

    const mapping::MappingPointContext context = make_context(identity);
    const mapping::MappingPointVariation variation =
        make_variation(identity, jacobian_variation);

    mfem::DenseMatrix tensor_variation;
    mapping::ComputeHDivMassTensorVariation(
        context, variation, tensor_variation
    );

    CHECK_THAT(
        variation.mapping_determinant_variation,
        WithinAbs(3.0 * scaling_variation, 1.0e-14)
    );

    for (int row = 0; row < dimension; ++row) {
        for (int column = 0; column < dimension; ++column) {
            const double expected = row == column ? -scaling_variation : 0.0;
            CHECK_THAT(
                tensor_variation(row, column), WithinAbs(expected, 1.0e-14)
            );
        }
    }
}

TEST_CASE(
    "Hdiv Mass Tensor Variation Symmetrizes Simple Shear At Identity",
    tags::unit &tags::transformations
) {
    constexpr double shear_variation = 0.23;

    const mfem::DenseMatrix identity = make_identity_matrix();
    mfem::DenseMatrix jacobian_variation(dimension);
    jacobian_variation                         = 0.0;
    jacobian_variation(0, 1)                   = shear_variation;

    const mapping::MappingPointContext context = make_context(identity);
    const mapping::MappingPointVariation variation =
        make_variation(identity, jacobian_variation);

    mfem::DenseMatrix tensor_variation;
    mapping::ComputeHDivMassTensorVariation(
        context, variation, tensor_variation
    );

    CHECK_THAT(
        variation.mapping_determinant_variation, WithinAbs(0.0, 1.0e-15)
    );
    CHECK_THAT(tensor_variation(0, 1), WithinAbs(shear_variation, 1.0e-14));
    CHECK_THAT(tensor_variation(1, 0), WithinAbs(shear_variation, 1.0e-14));

    for (int row = 0; row < dimension; ++row) {
        for (int column = 0; column < dimension; ++column) {
            if ((row == 0 && column == 1) || (row == 1 && column == 0))
                continue;
            CHECK_THAT(tensor_variation(row, column), WithinAbs(0.0, 1.0e-14));
        }
    }
}

TEST_CASE(
    "Mapping Determinant Variation Matches Jacobi Formula",
    tags::unit &tags::transformations
) {
    const mfem::DenseMatrix jacobian =
        make_matrix({1.24, 0.19, -0.07, -0.06, 0.88, 0.16, 0.04, -0.12, 1.19});

    const mfem::DenseMatrix jacobian_variation =
        make_matrix({0.08, -0.03, 0.05, 0.02, 0.06, -0.04, -0.01, 0.07, -0.02});

    const mapping::MappingPointContext context = make_context(jacobian);
    const mapping::MappingPointVariation variation =
        make_variation(jacobian, jacobian_variation);
    constexpr double difference_step = 1.0e-3;

    mfem::DenseMatrix plus_one(context.mapping_jacobian);
    mfem::DenseMatrix plus_two(context.mapping_jacobian);
    mfem::DenseMatrix minus_one(context.mapping_jacobian);
    mfem::DenseMatrix minus_two(context.mapping_jacobian);

    plus_one.Add(difference_step, variation.mapping_jacobian_variation);
    plus_two.Add(2.0 * difference_step, variation.mapping_jacobian_variation);
    minus_one.Add(-difference_step, variation.mapping_jacobian_variation);
    minus_two.Add(-2.0 * difference_step, variation.mapping_jacobian_variation);

    const double finite_difference = (minus_two.Det() - 8.0 * minus_one.Det() +
                                      8.0 * plus_one.Det() - plus_two.Det()) /
                                     (12.0 * difference_step);
    const double analytic          = variation.mapping_determinant_variation;
    const double relative_error    = std::abs(finite_difference - analytic) /
                                     std::max(std::abs(analytic), 1.0e-14);

    INFO("Analytic determinant variation = " << analytic);
    INFO("Finite-difference determinant variation = " << finite_difference);
    INFO("Relative determinant-variation error = " << relative_error);

    CHECK(relative_error < 2.0e-11);
}