#include <cmath>
#include <limits>
#include <stdexcept>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include "mfem/linalg/densemat.hpp"
#include "serif/backend/mfem/solver/linear/solve.hpp"
#include "serif/utils/error/errors.hpp"
#include "serif/backend/mfem/operators/actions.hpp"

namespace {
    namespace linear = serif::solver::linear;
    constexpr auto fixed = serif::solver::preconditioner::ApplicationContract::fixed_linear;

    /** @brief Keep an exact identity inverse while checking the original operator address is preserved. */
    class IdentityInverse final : public ::mfem::Solver {
    public:
        /** @brief Bind one test linearization. */
        explicit IdentityInverse(const ::mfem::Operator& operation) : ::mfem::Solver(operation.Width()), m_operation(operation) {}

        /** @brief Reject a counting wrapper or unrelated operator replacing the accepted linearization. */
        void SetOperator(const ::mfem::Operator& operation) override { REQUIRE(&operation == &m_operation); }

        /** @brief Copy the correction without changing its coordinates. */
        void Mult(const ::mfem::Vector& input, ::mfem::Vector& output) const override { output = input; }

    private:
        const ::mfem::Operator& m_operation;
    };

    /** @brief Use MeanField's nonsymmetric two-coordinate problem, with solution (0.1, 0.6). */
    [[nodiscard]] ::mfem::DenseMatrix example_matrix() {
        ::mfem::DenseMatrix matrix(2);
        matrix(0, 0) = 4.0;
        matrix(0, 1) = 1.0;
        matrix(1, 0) = 2.0;
        matrix(1, 1) = 3.0;
        return matrix;
    }
}

TEST_CASE("Checked FGMRES preserves restart counts and RHS-relative warm-start accuracy", "[solver][linear]") {
    const ::mfem::DenseMatrix matrix = example_matrix();
    IdentityInverse inverse(matrix);
    ::mfem::Vector rhs(2), correction(2);
    rhs[0] = 1.0;
    rhs[1] = 2.0;
    correction = 0.0;
    const linear::FGMRES options{.restart_length = 1};
    const linear::Control control{.relative_tolerance = 1.0e-12, .absolute_tolerance = 1.0e-14, .maximum_iterations = 200};
    const auto cold = linear::backend::mfem::solve(options, matrix, inverse, fixed, rhs, correction, control);
    REQUIRE(cold.converged());
    CHECK(correction[0] == Catch::Approx(0.1).margin(1.0e-11));
    CHECK(correction[1] == Catch::Approx(0.6).margin(1.0e-11));
    CHECK(cold.right_hand_side_norm == Catch::Approx(std::sqrt(5.0)));
    CHECK(cold.true_residual_norm <= control.convergence_threshold(cold.right_hand_side_norm));
    CHECK(cold.iterations == static_cast<int>(cold.preconditioner_applications));
    CHECK(cold.iterations > 1);
    CHECK(cold.iterations <= control.maximum_iterations);
    CHECK(cold.restarts == cold.iterations - 1);
    CHECK(cold.operator_applications > cold.preconditioner_applications);
    CHECK(cold.solve_seconds >= 0.0);
    CHECK(cold.operator_seconds >= 0.0);
    CHECK(cold.preconditioner_seconds >= 0.0);

    correction[0] += 1.0e-5;
    correction[1] -= 1.0e-5;
    const auto warm = linear::backend::mfem::solve(options, matrix, inverse, fixed, rhs, correction, control);
    REQUIRE(warm.converged());
    CHECK(warm.initial_residual_norm < cold.initial_residual_norm);
    CHECK(warm.true_residual_norm <= control.convergence_threshold(warm.right_hand_side_norm));
}

TEST_CASE("Checked FGMRES discards an adverse or overflowing finite guess", "[solver][linear]") {
    const ::mfem::DenseMatrix matrix = example_matrix();
    IdentityInverse inverse(matrix);
    ::mfem::Vector rhs(2), correction(2);
    rhs[0] = 1.0;
    rhs[1] = 2.0;
    correction = 0.0;
    const linear::FGMRES options{.restart_length = 1};
    const linear::Control control{.relative_tolerance = 1.0e-12, .absolute_tolerance = 1.0e-14, .maximum_iterations = 200};
    const auto cold = linear::backend::mfem::solve(options, matrix, inverse, fixed, rhs, correction, control);
    REQUIRE(cold.converged());
    SECTION("Finite adverse coefficients") {
        correction[0] = 100.0;
        correction[1] = -100.0;
    }
    SECTION("Finite coefficients with overflowing action") {
        correction = std::numeric_limits<double>::max() / 2.0;
    }
    const auto guarded = linear::backend::mfem::solve(options, matrix, inverse, fixed, rhs, correction, control);
    REQUIRE(guarded.converged());
    CHECK(guarded.initial_residual_norm == Catch::Approx(cold.right_hand_side_norm));
    CHECK(guarded.iterations == cold.iterations);
    CHECK(guarded.operator_applications == cold.operator_applications + 1);
}

TEST_CASE("Checked FGMRES verifies exact guesses and zero right-hand sides without Krylov iterations", "[solver][linear]") {
    const ::mfem::DenseMatrix matrix = example_matrix();
    IdentityInverse inverse(matrix);
    ::mfem::Vector rhs(2), correction(2);
    correction[0] = 1.0;
    correction[1] = 2.0;
    matrix.Mult(correction, rhs);
    SECTION("Exact nonzero initial guess") {
        const auto report = linear::backend::mfem::solve({}, matrix, inverse, fixed, rhs, correction);
        REQUIRE(report.converged());
        CHECK(report.iterations == 0);
        CHECK(report.restarts == 0);
        CHECK(report.operator_applications == 2);
        CHECK(report.true_residual_norm == 0.0);
    }
    SECTION("Zero equation discards a nonzero guess") {
        rhs = 0.0;
        const auto report = linear::backend::mfem::solve({}, matrix, inverse, fixed, rhs, correction);
        REQUIRE(report.converged());
        CHECK(report.iterations == 0);
        CHECK(report.relative_true_residual_norm == 0.0);
        CHECK(correction.Norml2() == 0.0);
    }
    SECTION("Variable preconditioner contract is supported") {
        const auto report = linear::backend::mfem::solve({}, matrix, inverse,
            serif::solver::preconditioner::ApplicationContract::variable, rhs, correction);
        CHECK(report.converged());
    }
}

TEST_CASE("Checked FGMRES reports exhausted budgets and nonfinite numerical results", "[solver][linear]") {
    ::mfem::DenseMatrix matrix = example_matrix();
    IdentityInverse inverse(matrix);
    ::mfem::Vector rhs(2), correction(2);
    rhs[0] = 1.0;
    rhs[1] = 2.0;
    correction = 0.0;
    SECTION("A single Krylov iteration is insufficient") {
        const auto report = linear::backend::mfem::solve({}, matrix, inverse, fixed, rhs, correction,
            {.relative_tolerance = 1.0e-14, .maximum_iterations = 1});
        CHECK(report.status == linear::Status::maximum_iterations);
        CHECK(report.iterations == 1);
        CHECK(report.true_residual_norm > report.control.convergence_threshold(report.right_hand_side_norm));
    }
    SECTION("Nonfinite RHS is numerical failure") {
        rhs[0] = std::numeric_limits<double>::infinity();
        const auto report = linear::backend::mfem::solve({}, matrix, inverse, fixed, rhs, correction);
        CHECK(report.status == linear::Status::non_finite);
        CHECK(report.operator_applications == 0);
    }
    SECTION("Nonfinite initial correction is numerical failure") {
        correction[0] = std::numeric_limits<double>::quiet_NaN();
        CHECK(linear::backend::mfem::solve({}, matrix, inverse, fixed, rhs, correction).status == linear::Status::non_finite);
    }
    SECTION("SERiF action overflow is a numerical report, without catching unrelated callback errors") {
        const auto action = serif::operators::backend::mfem::linearized_action(2, 2,
            [](const ::mfem::Vector&, ::mfem::Vector& output) { output = std::numeric_limits<double>::infinity(); },
            [](const ::mfem::Vector&, ::mfem::Vector& output) { output = std::numeric_limits<double>::infinity(); });
        IdentityInverse action_inverse(*action.action());
        CHECK(linear::backend::mfem::solve({}, *action.action(), action_inverse, fixed, rhs, correction).status == linear::Status::non_finite);
    }
    SECTION("Nonfinite operator action cannot converge") {
        matrix(0, 0) = std::numeric_limits<double>::quiet_NaN();
        CHECK(linear::backend::mfem::solve({}, matrix, inverse, fixed, rhs, correction).status == linear::Status::non_finite);
    }
}

TEST_CASE("Checked FGMRES independently rejects misleading backend convergence", "[solver][linear]") {
    /** @brief Mimic a stale callback whose final action disagrees with the Krylov recurrence. */
    class ChangingOperator final : public ::mfem::Operator {
    public:
        /** @brief Use one coordinate so the recurrence converges in one iteration. */
        ChangingOperator() : ::mfem::Operator(1) {}
        /** @brief Change only after MFEM's initial and Arnoldi operator applications. */
        void Mult(const ::mfem::Vector& input, ::mfem::Vector& output) const override {
            ++applications;
            output[0] = input[0] * (applications > 2 ? 2.0 : 1.0);
        }
        mutable int applications{0};
    } operation;
    IdentityInverse inverse(operation);
    ::mfem::Vector rhs(1), correction(1);
    rhs = 1.0;
    correction = 0.0;
    const auto report = linear::backend::mfem::solve({}, operation, inverse, fixed, rhs, correction);
    CHECK(report.status == linear::Status::backend_failure);
    CHECK(report.reported_residual_norm == 0.0);
    CHECK(report.true_residual_norm == 1.0);
}

TEST_CASE("Checked FGMRES validates configuration and preserves callback errors", "[solver][linear]") {
    CHECK_THROWS_AS(linear::FGMRES{.restart_length = 0}.validate(), std::invalid_argument);
    CHECK_THROWS_AS(linear::FGMRES{.print_level = 4}.validate(), std::invalid_argument);
    for (const double invalid : {-1.0, std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::infinity()}) {
        CHECK_THROWS_AS(linear::Control{.relative_tolerance = invalid}.validate(), std::invalid_argument);
        CHECK_THROWS_AS(linear::Control{.absolute_tolerance = invalid}.validate(), std::invalid_argument);
        CHECK_THROWS_AS(linear::Control{}.convergence_threshold(invalid), std::invalid_argument);
    }
    CHECK_THROWS_AS(linear::Control{.maximum_iterations = 0}.validate(), std::invalid_argument);
    CHECK_THROWS_AS(linear::Control{.relative_tolerance = std::numeric_limits<double>::max()}.convergence_threshold(2.0), std::invalid_argument);
    const linear::Control control{.relative_tolerance = 0.25, .absolute_tolerance = 0.5};
    CHECK(control.convergence_threshold(4.0) == 1.0);
    CHECK(control.convergence_threshold(0.0) == 0.5);
    /** @brief Model a genuine programming error rather than a numerical rejection. */
    class BrokenOperator final : public ::mfem::Operator {
    public:
        /** @brief Match one valid coefficient. */
        BrokenOperator() : ::mfem::Operator(1) {}
        /** @brief Exercise the error boundary without introducing a broad catch in the solver. */
        void Mult(const ::mfem::Vector&, ::mfem::Vector&) const override {
            serif::utils::errors::serif_error<std::logic_error, serif::utils::errors::SERiFErrorCode::SOLVER_ERROR>("Broken callback contract.");
        }
    } operation;
    IdentityInverse inverse(operation);
    ::mfem::Vector rhs(1), correction(1);
    rhs = 1.0;
    correction = 0.0;
    CHECK_THROWS_AS(linear::backend::mfem::solve({}, operation, inverse, fixed, rhs, correction), std::logic_error);
    correction.SetSize(2);
    CHECK_THROWS_AS(linear::backend::mfem::solve({}, operation, inverse, fixed, rhs, correction), std::invalid_argument);
}