#include <initializer_list>
#include <memory>
#include <array>
#include <cmath>
#include <limits>
#include <vector>
#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include "serif/backend/mfem/solver/preconditioner/inverse/amg.hpp"
#include "serif/backend/mfem/solver/preconditioner/inverse/dense.hpp"
#include "serif/backend/mfem/solver/preconditioner/inverse/diagonal.hpp"

#ifdef MFEM_USE_MPI
#include "mfem/general/communication.hpp"
#endif

namespace
{
    namespace preconditioner = serif::solver::preconditioner;
    namespace backend = preconditioner::backend::mfem;
    namespace options = preconditioner::strategies::inverse;
    using Code = serif::utils::errors::SERiFErrorCode;

    /** @brief Make small right-hand sides readable without coupling expectations to a numerical implementation. */
    mfem::Vector vector(std::initializer_list<double> values)
    {
        mfem::Vector result(static_cast<int>(values.size()));
        int index = 0;
        for (const double value : values)
        {
            result[index++] = value;
        }
        return result;
    }

    /** @brief Compare inverse actions at the original MeanField tolerances. */
    void check_vector(const mfem::Vector& actual, const mfem::Vector& expected, const double tolerance = 2.0e-13)
    {
        REQUIRE(actual.Size() == expected.Size());
        for (int index = 0; index < actual.Size(); ++index)
        {
            CHECK(actual[index] == Catch::Approx(expected[index]).epsilon(tolerance).margin(tolerance));
        }
    }

    /** @brief Reuse one positive definite mass fixture for polynomial linearity and frozen ownership. */
    std::shared_ptr<mfem::DenseMatrix> mass_matrix()
    {
        auto matrix = std::make_shared<mfem::DenseMatrix>(3);
        *matrix = 0.0;
        (*matrix)(0, 0) = 4.0;
        (*matrix)(1, 1) = 3.0;
        (*matrix)(2, 2) = 2.0;
        (*matrix)(0, 1) = (*matrix)(1, 0) = 1.0;
        (*matrix)(1, 2) = (*matrix)(2, 1) = 1.0;
        return matrix;
    }

    /** @brief Check the fixed-map contract independently of a particular inverse algorithm. */
    void check_fixed_linear(const backend::Inverse& inverse, const mfem::Vector& x, const mfem::Vector& y, const double tolerance)
    {
        REQUIRE(inverse.contract == preconditioner::ApplicationContract::fixed_linear);
        mfem::Vector px(x.Size()), py(x.Size()), actual(x.Size()), expected(x.Size()), combination(x);
        inverse.action->Mult(x, px);
        inverse.action->Mult(y, py);
        combination *= 1.7;
        combination.Add(-0.4, y);
        actual = std::numeric_limits<double>::quiet_NaN();
        const auto* storage = actual.GetData();
        inverse.action->Mult(combination, actual);
        CHECK(actual.GetData() == storage);
        expected = px;
        expected *= 1.7;
        expected.Add(-0.4, py);
        check_vector(actual, expected, tolerance);
        inverse.action->Mult(x, actual);
        check_vector(actual, px, 0.0);
        combination = 0.0;
        inverse.action->Mult(combination, actual);
        CHECK(actual.Norml2() == 0.0);
        actual = x;
        inverse.action->Mult(actual, actual);
        check_vector(actual, px, tolerance);
        mfem::Vector overlapping_storage(x.Size() + 1);
        for (const int shift : {0, 1})
        {
            mfem::Vector overlapping_input(overlapping_storage.GetData() + shift, x.Size());
            mfem::Vector overlapping_output(overlapping_storage.GetData() + 1 - shift, x.Size());
            overlapping_input = x;
            inverse.action->Mult(overlapping_input, overlapping_output);
            check_vector(overlapping_output, px, tolerance);
        }
    }
} // namespace

TEST_CASE("Signed diagonal inverses own their factors and reject unrepresentable reciprocals", "[preconditioner][inverse][unit]")
{
    std::vector<double> diagonal{2.0, -4.0, 0.5};
    const auto inverse = backend::try_diagonal_inverse(diagonal);
    REQUIRE(inverse);
    diagonal.assign(3, 77.0);
    mfem::Vector action(3);
    inverse->action->Mult(vector({4.0, 8.0, -1.0}), action);
    check_vector(action, vector({2.0, -2.0, -2.0}));
    inverse->action->MultTranspose(vector({4.0, 8.0, -1.0}), action);
    check_vector(action, vector({2.0, -2.0, -2.0}));
    check_fixed_linear(*inverse, vector({1.0, -2.0, 0.25}), vector({-0.5, 0.75, 3.0}), 2.0e-13);
    for (const double bad : {0.0, std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::denorm_min()})
    {
        CHECK_FALSE(backend::try_diagonal_inverse(std::array{bad}));
    }
    CHECK_FALSE(backend::try_diagonal_inverse({}));
}

TEST_CASE("Dense inverse retains signed pivoted solves transpose and independent snapshots", "[preconditioner][inverse][unit]")
{
    mfem::DenseMatrix matrix(2);
    matrix(0, 0) = 4.0;
    matrix(0, 1) = 1.0;
    matrix(1, 0) = 2.0;
    matrix(1, 1) = 3.0;
    const auto inverse = backend::try_dense_inverse(matrix);
    REQUIRE(inverse);
    matrix = 0.0;
    matrix(0, 0) = 2.0;
    matrix(1, 1) = 4.0;
    const auto replacement = backend::try_dense_inverse(matrix);
    REQUIRE(replacement);
    mfem::Vector action(2);
    inverse->action->Mult(vector({7.0, 1.0}), action);
    check_vector(action, vector({2.0, -1.0}));
    inverse->action->MultTranspose(vector({6.0, -1.0}), action);
    check_vector(action, vector({2.0, -1.0}));
    replacement->action->Mult(vector({7.0, 1.0}), action);
    check_vector(action, vector({3.5, 0.25}));
    check_fixed_linear(*inverse, vector({1.0, -2.0}), vector({-0.5, 0.75}), 2.0e-13);
}

TEST_CASE("Dense conditioning uses scale-independent infinity norms and recoverable failures", "[preconditioner][inverse][unit]")
{
    for (const double scale : {1.0e-250, 1.0, 1.0e250})
    {
        mfem::DenseMatrix matrix(2);
        matrix(0, 0) = 0.0;
        matrix(0, 1) = 2.0 * scale;
        matrix(1, 0) = 3.0 * scale;
        matrix(1, 1) = scale;
        const auto inverse = backend::try_dense_inverse(matrix, {.minimum_reciprocal_condition = 0.49});
        REQUIRE(inverse);
        CHECK(backend::try_dense_inverse(matrix, {.minimum_reciprocal_condition = 0.5}));
        const auto ill_conditioned = backend::try_dense_inverse(matrix, {.minimum_reciprocal_condition = 0.51});
        REQUIRE_FALSE(ill_conditioned);
        CHECK(ill_conditioned.error().code == Code::PRECONDITIONER_ILL_CONDITIONED_OPERATOR);
        mfem::Vector action(2);
        inverse->action->Mult(vector({-2.0 * scale, 5.0 * scale}), action);
        check_vector(action, vector({2.0, -1.0}));
    }
    mfem::DenseMatrix singular(2);
    singular = 1.0;
    const auto rejected = backend::try_dense_inverse(singular);
    REQUIRE_FALSE(rejected);
    CHECK(rejected.error().code == Code::PRECONDITIONER_SINGULAR_OPERATOR);
    singular(0, 0) = std::numeric_limits<double>::quiet_NaN();
    const auto nonfinite = backend::try_dense_inverse(singular);
    REQUIRE_FALSE(nonfinite);
    CHECK(nonfinite.error().code == Code::PRECONDITIONER_NONFINITE_VALUE);
    CHECK_FALSE(backend::try_dense_inverse(singular, {.minimum_reciprocal_condition = -1.0}));
    CHECK_FALSE(backend::try_dense_inverse(mfem::DenseMatrix(2, 3)));

    mfem::DenseMatrix disparate(2);
    disparate = 0.0;
    disparate(0, 0) = 1.0e308;
    disparate(1, 1) = 1.0e-308;
    const auto disparate_inverse = backend::try_dense_inverse(disparate);
    REQUIRE(disparate_inverse);
    mfem::Vector action(2);
    disparate_inverse->action->Mult(vector({1.0e308, -1.0e-308}), action);
    check_vector(action, vector({1.0, -1.0}));
    disparate_inverse->action->MultTranspose(vector({1.0e308, -1.0e-308}), action);
    check_vector(action, vector({1.0, -1.0}));
    CHECK_FALSE(backend::try_dense_inverse(disparate, {.minimum_reciprocal_condition = 1.0e-14}));
}

TEST_CASE("One-process BoomerAMG preserves fixed-cycle linearity", "[preconditioner][inverse][integration]")
{
    auto matrix = std::make_shared<mfem::SparseMatrix>(32);
    for (int row = 0; row < matrix->Height(); ++row)
    {
        matrix->Add(row, row, 2.0);
        if (row > 0)
        {
            matrix->Add(row, row - 1, -1.0);
            matrix->Add(row - 1, row, -1.0);
        }
    }
    matrix->Finalize();
    const auto zero = std::make_shared<mfem::SparseMatrix>(2);
    zero->Add(0, 0, 0.0);
    zero->Add(1, 1, 0.0);
    zero->Finalize();
    const auto unsupported = backend::try_amg_inverse(zero);
    REQUIRE_FALSE(unsupported);
    CHECK(unsupported.error().code == Code::PRECONDITIONER_INVALID_OPERATOR);
    CHECK_FALSE(backend::try_amg_inverse(matrix, options::FixedAMG{.cycles = 0}));
#ifdef MFEM_USE_MPI
    if (!mfem::Mpi::IsInitialized())
    {
        mfem::Mpi::Init();
    }
    const auto fixed = backend::try_amg_inverse(matrix, options::FixedAMG{.cycles = 2});
    REQUIRE(fixed);
    mfem::Vector exact(32), second(32), rhs(32), second_rhs(32), action(32), residual(32);
    for (int index = 0; index < 32; ++index)
    {
        exact[index] = 1.0;
        second[index] = 0.25 + index;
    }
    matrix->Mult(exact, rhs);
    matrix->Mult(second, second_rhs);
    check_fixed_linear(*fixed, rhs, second_rhs, 1.0e-11);
    fixed->action->Mult(rhs, action);
    matrix->Mult(action, residual);
    residual -= rhs;
    CHECK(residual.Norml2() < rhs.Norml2());
    matrix.reset();
    fixed->action->Mult(rhs, action);
    CHECK(std::isfinite(action.Norml2()));
#else
    const auto unavailable = backend::try_amg_inverse(matrix);
    REQUIRE_FALSE(unavailable);
    CHECK(unavailable.error().code == Code::PRECONDITIONER_INCOMPATIBLE_STRATEGY);
#endif
}