#include <array>
#include <limits>
#include <memory>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include "mfem/linalg/densemat.hpp"
#include "serif/backend/mfem/solver/preconditioner/border.hpp"
#include "serif/backend/mfem/solver/preconditioner/inverse/dense.hpp"

namespace
{
    namespace backend = serif::solver::preconditioner::backend::mfem;
    using serif::solver::preconditioner::ApplicationContract;
    using serif::solver::preconditioner::strategies::coupling::CouplingOrder;

    /** @brief Construct a small row-major matrix independently of the implementation's block algebra. */
    [[nodiscard]] std::shared_ptr<mfem::DenseMatrix> matrix(int height, int width, std::initializer_list<double> entries)
    {
        REQUIRE(entries.size() == static_cast<std::size_t>(height * width));
        auto result = std::make_shared<mfem::DenseMatrix>(height, width);
        auto entry = entries.begin();
        for (int row = 0; row < height; ++row)
        {
            for (int column = 0; column < width; ++column)
            {
                (*result)(row, column) = *entry++;
            }
        }
        return result;
    }

    /** @brief Compare complete numerical actions without ignoring small absolute discrepancies. */
    [[nodiscard]] double difference(const mfem::Vector& left, const mfem::Vector& right)
    {
        mfem::Vector result(left);
        result -= right;
        return result.Norml2();
    }
} // namespace

TEST_CASE("Two-group coupling policies reproduce their complete nonsymmetric factorizations", "[preconditioner][coupling][unit]")
{
    const backend::Inverse first{.action = matrix(2, 2, {0.5, 0.0, 0.0, 0.25})};
    const backend::Inverse second{.action = matrix(2, 2, {0.2, 0.0, 0.0, 0.1})};
    const backend::TwoBlockCouplings couplings{.first_from_second = matrix(2, 2, {1, 2, -1, 3}), .second_from_first = matrix(2, 2, {3, -2, 4, 1})};
    const std::array orders{CouplingOrder::independent, CouplingOrder::second_then_first, CouplingOrder::lower_diagonal_upper};
    const std::array systems{
        matrix(4, 4, {2, 0, 0, 0, 0, 4, 0, 0, 0, 0, 5, 0, 0, 0, 0, 10}),
        matrix(4, 4, {2, 0, 1, 2, 0, 4, -1, 3, 0, 0, 5, 0, 0, 0, 0, 10}),
        matrix(4, 4, {2, 0, 1, 2, 0, 4, -1, 3, 3, -2, 7, 1.5, 4, 1, 1.75, 14.75})};
    mfem::Vector rhs(4), actual(4), expected(4);
    rhs[0] = 1.0;
    rhs[1] = -2.0;
    rhs[2] = 3.0;
    rhs[3] = 4.0;
    for (std::size_t policy = 0; policy < orders.size(); ++policy)
    {
        CAPTURE(policy);
        auto inverse = backend::try_coupled_inverse(first, second, couplings, orders[policy]);
        REQUIRE(inverse);
        CHECK(inverse->contract == ApplicationContract::fixed_linear);
        mfem::DenseMatrixInverse exact(*systems[policy]);
        exact.Mult(rhs, expected);
        double* const storage = actual.GetData();
        inverse->action->Mult(rhs, actual);
        CHECK(actual.GetData() == storage);
        CHECK(difference(actual, expected) < 2.0e-13);
        mfem::Vector aliased(rhs);
        inverse->action->Mult(aliased, aliased);
        CHECK(difference(aliased, expected) < 2.0e-13);
        inverse->action->Mult(rhs, actual);
        CHECK(difference(actual, expected) < 2.0e-13);
    }
}

TEST_CASE("A matrix-free Schur action shares the eliminated-field orientation", "[preconditioner][coupling][unit]")
{
    const backend::Inverse eliminated{.action = matrix(2, 2, {0.5, 0.0, 0.0, 0.25})};
    const backend::TwoBlockCouplings couplings{.first_from_second = matrix(2, 2, {1, 2, -1, 3}), .second_from_first = matrix(2, 2, {3, -2, 4, 1})};
    auto schur = backend::try_schur_operator(matrix(2, 2, {7, 1.5, 1.75, 14.75}), eliminated, couplings);
    REQUIRE(schur);
    mfem::Vector values(2);
    values[0] = 3.0;
    values[1] = -2.0;
    (*schur)->Mult(values, values);
    CHECK(values[0] == Catch::Approx(15.0).margin(2.0e-14));
    CHECK(values[1] == Catch::Approx(-20.0).margin(2.0e-14));
    backend::Inverse variable = eliminated;
    variable.contract = ApplicationContract::variable;
    CHECK_FALSE(backend::try_schur_operator(matrix(2, 2, {7, 1.5, 1.75, 14.75}), variable, couplings));
}

TEST_CASE("Coupled inverse validation rejects unsupported policies and preserves outputs", "[preconditioner][coupling][validation]")
{
    using serif::utils::errors::SERiFErrorCode;

    const backend::Inverse identity{.action = matrix(1, 1, {1.0})};
    const backend::TwoBlockCouplings full{.first_from_second = matrix(1, 1, {2.0}), .second_from_first = matrix(1, 1, {3.0})};
    CHECK(backend::try_coupled_inverse(identity, identity, {}, CouplingOrder::independent));
    CHECK(backend::try_coupled_inverse(identity, identity, {.first_from_second = full.first_from_second}, CouplingOrder::second_then_first));
    auto missing = backend::try_coupled_inverse({}, identity, full, CouplingOrder::independent);
    REQUIRE_FALSE(missing);
    CHECK(missing.error().code == SERiFErrorCode::PRECONDITIONER_INVALID_OPERATOR);
    auto incompatible = backend::try_coupled_inverse(identity, identity, {}, CouplingOrder::lower_diagonal_upper);
    REQUIRE_FALSE(incompatible);
    CHECK(incompatible.error().code == SERiFErrorCode::PRECONDITIONER_SIZE_MISMATCH);
    auto unknown = backend::try_coupled_inverse(identity, identity, full, static_cast<CouplingOrder>(99));
    REQUIRE_FALSE(unknown);
    CHECK(unknown.error().code == SERiFErrorCode::PRECONDITIONER_INVALID_OPTIONS);
    backend::Inverse variable = identity;
    variable.contract = ApplicationContract::variable;
    auto composed = backend::try_coupled_inverse(identity, variable, {}, CouplingOrder::independent);
    REQUIRE(composed);
    CHECK(composed->contract == ApplicationContract::variable);

    mfem::Vector rhs(2), output(2), short_rhs(1);
    rhs = 1.0;
    output = 7.0;
    CHECK_THROWS(composed->action->Mult(short_rhs, output));
    CHECK(output[0] == 7.0);
    CHECK(output[1] == 7.0);
    rhs[1] = std::numeric_limits<double>::quiet_NaN();
    CHECK_THROWS(composed->action->Mult(rhs, output));
    CHECK(output[0] == 7.0);
    CHECK(output[1] == 7.0);
}

TEST_CASE("Selected operators preserve distinct equation and unknown orders and their transpose", "[preconditioner][coupling][coordinates]")
{
    const auto source = matrix(3, 4, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
    const std::array rows{2, 0};
    const std::array columns{3, 1};
    auto selected = backend::try_restrict_operator(source, rows, columns);
    REQUIRE(selected);
    mfem::Vector values(2);
    values[0] = 2.0;
    values[1] = -1.0;
    (*selected)->Mult(values, values);
    CHECK(values[0] == 14.0);
    CHECK(values[1] == 6.0);
    values[0] = 2.0;
    values[1] = -1.0;
    (*selected)->MultTranspose(values, values);
    CHECK(values[0] == 20.0);
    CHECK(values[1] == 18.0);
    CHECK_FALSE(backend::try_restrict_operator(source, std::array{2, 2}, columns));
    CHECK_FALSE(backend::try_restrict_operator(source, std::array{3}, columns));
    CHECK_FALSE(backend::try_restrict_operator(source, rows, std::array{-1}));
}