#include "serif/utils/error/errors.hpp"
#include <initializer_list>
#include <limits>
#include <stdexcept>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_string.hpp>
#include <catch2/matchers/catch_matchers_exception.hpp>
#include <mfem/linalg/densemat.hpp>
#include <mfem/linalg/solvers.hpp>

#include "serif/backend/mfem/solver/normalization/operators.hpp"

namespace
{
    namespace normalization = serif::solver::normalization;
    namespace backend = normalization::backend::mfem;

    /** @brief Form a vector from physically readable test coordinates. */
    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 coordinate actions at the original MeanField relative tolerances. */
    void check_vector(const mfem::Vector &actual, const mfem::Vector &expected, double epsilon = 2.0e-13)
    {
        REQUIRE(actual.Size() == expected.Size());
        for (int index = 0; index < actual.Size(); ++index)
        {
            CHECK(actual(index) == Catch::Approx(expected(index)).epsilon(epsilon).margin(1.0e-300));
        }
    }

    /** @brief Exact physical inverse which records what dimensional Jacobian MFEM binds. */
    class DenseInverseSolver final : public mfem::Solver
    {
    public:
        /** @brief Copy a small reference inverse for independent solver assertions. */
        explicit DenseInverseSolver(const mfem::DenseMatrix &inverse) : mfem::Solver(inverse.Height(), inverse.Width()), m_inverse(inverse)
        {
        }

        /** @brief Bind the dimensional Jacobian, permitting a rectangular inverse map. */
        void SetOperator(const mfem::Operator &operation) override
        {
            if (operation.Height() != Width() || operation.Width() != Height())
            {
                serif::utils::errors::serif_error<std::invalid_argument, serif::utils::errors::SERiFErrorCode::INVALID_ARGUMENT>("Dense test inverse received an incompatible Jacobian.");
            }
            bound_operator = &operation;
            ++bindings;
        }

        /** @brief Apply the exact inverse after checking the iterative initial guess. */
        void Mult(const mfem::Vector &input, mfem::Vector &output) const override
        {
            if (iterative_mode)
            {
                for (int index = 0; index < output.Size(); ++index)
                {
                    CHECK(output(index) == 0.0);
                }
            }
            m_inverse.Mult(input, output);
        }

        /** @brief Apply the transpose of the exact inverse. */
        void MultTranspose(const mfem::Vector &input, mfem::Vector &output) const override
        {
            m_inverse.MultTranspose(input, output);
        }

        const mfem::Operator *bound_operator{nullptr};
        int bindings{0};

    private:
        mfem::DenseMatrix m_inverse;
    };

    /** @brief Recover dimensional matrices from a well-conditioned normalized test problem. */
    struct PhysicalSystem final
    {
        normalization::DiagonalNormalization map = *normalization::make_diagonal_normalization({1.0e-9, 2.0e7}, {5.0e8, 3.0e-6});
        mfem::DenseMatrix jacobian{2};
        mfem::DenseMatrix inverse{2};

        /** @brief Use independently known normalized matrix and inverse entries. */
        PhysicalSystem()
        {
            constexpr double normalized_matrix[2][2]{{4.0, 1.0}, {2.0, 3.0}};
            constexpr double normalized_inverse[2][2]{{0.3, -0.1}, {-0.2, 0.4}};
            for (int row = 0; row < 2; ++row)
            {
                for (int column = 0; column < 2; ++column)
                {
                    jacobian(row, column) = normalized_matrix[row][column] * map.state_factors()[column] / map.residual_factors()[row];
                    inverse(row, column) = normalized_inverse[row][column] * map.residual_factors()[column] / map.state_factors()[row];
                }
            }
        }
    };
}

// Ported from MeanField tests/normalization/physical_riesz.cpp, with transpose coverage added.
TEST_CASE("Scaled Jacobian and inverse implement the exact coordinate change", "[normalization][unit][linear-algebra]")
{
    const PhysicalSystem physical;
    const backend::ScaledJacobianOperator jacobian(physical.jacobian, physical.map);
    const backend::ScaledInverseOperator inverse(physical.inverse, physical.map);
    const mfem::Vector direction = vector({0.75, -1.25});
    mfem::Vector action;
    jacobian.Mult(direction, action);
    check_vector(action, vector({1.75, -2.25}), 4.0e-15);
    mfem::Vector recovered;
    inverse.Mult(action, recovered);
    check_vector(recovered, direction);

    jacobian.MultTranspose(direction, action);
    check_vector(action, vector({0.5, -3.0}), 4.0e-15);
    inverse.MultTranspose(action, recovered);
    check_vector(recovered, direction);
}

TEST_CASE("Scaled physical preconditioning preserves dimensional binding through FGMRES", "[normalization][unit][solver]")
{
    const PhysicalSystem physical;
    const backend::ScaledJacobianOperator jacobian(physical.jacobian, physical.map);
    DenseInverseSolver inverse(physical.inverse);
    backend::ScaledPreconditioner preconditioner(inverse, jacobian);
    CHECK(inverse.bound_operator == &physical.jacobian);
    const mfem::Vector right_hand_side = vector({1.5, -0.75});
    const mfem::Vector expected = vector({0.525, -0.6});
    mfem::Vector correction(2);
    preconditioner.Mult(right_hand_side, correction);
    check_vector(correction, expected, 3.0e-13);

    mfem::FGMRESSolver krylov;
    krylov.SetRelTol(1.0e-13);
    krylov.SetAbsTol(1.0e-15);
    krylov.SetMaxIter(4);
    krylov.SetKDim(2);
    krylov.SetPrintLevel(0);
    krylov.SetPreconditioner(preconditioner);
    krylov.SetOperator(jacobian);
    mfem::Vector solution(2);
    solution = 0.0;
    krylov.Mult(right_hand_side, solution);
    CHECK(krylov.GetConverged());
    CHECK(krylov.GetNumIterations() <= 1);
    check_vector(solution, expected, 3.0e-13);
    CHECK(inverse.bound_operator == &physical.jacobian);
    CHECK(inverse.bindings >= 2);
    CHECK(preconditioner.statistics().operator_bindings >= 2);
    CHECK(preconditioner.statistics().applications >= 2);

    preconditioner.MultTranspose(right_hand_side, correction);
    check_vector(correction, vector({0.6, -0.45}));
}

TEST_CASE("Scaled preconditioners reject unrelated preparations and malformed vectors", "[normalization][unit][validation]")
{
    const PhysicalSystem physical;
    const backend::ScaledJacobianOperator jacobian(physical.jacobian, physical.map);
    DenseInverseSolver inverse(physical.inverse);
    backend::ScaledPreconditioner preconditioner(inverse, jacobian);
    const backend::ScaledJacobianOperator unrelated(physical.jacobian, physical.map);
    const mfem::IdentityOperator wrong_size(3);
    CHECK_THROWS_MATCHES(preconditioner.SetOperator(unrelated), std::invalid_argument, Catch::Matchers::MessageMatches(Catch::Matchers::ContainsSubstring("NORMALIZATION_ERROR")));
    CHECK_THROWS_MATCHES(preconditioner.SetOperator(wrong_size), std::invalid_argument, Catch::Matchers::MessageMatches(Catch::Matchers::ContainsSubstring("NORMALIZATION_ERROR")));
    CHECK(inverse.bindings == 1);

    const mfem::Vector right_hand_side = vector({1.0, 2.0});
    mfem::Vector wrong_output(1);
    CHECK_THROWS_MATCHES(preconditioner.Mult(right_hand_side, wrong_output), std::invalid_argument, Catch::Matchers::MessageMatches(Catch::Matchers::ContainsSubstring("NORMALIZATION_VECTOR_SIZE_MISMATCH")));
    CHECK_THROWS_MATCHES(preconditioner.MultTranspose(right_hand_side, wrong_output), std::invalid_argument, Catch::Matchers::MessageMatches(Catch::Matchers::ContainsSubstring("NORMALIZATION_VECTOR_SIZE_MISMATCH")));
    mfem::Vector output(2);
    output = 7.0;
    CHECK_THROWS_MATCHES(preconditioner.Mult(wrong_output, output), std::invalid_argument, Catch::Matchers::MessageMatches(Catch::Matchers::ContainsSubstring("NORMALIZATION_VECTOR_SIZE_MISMATCH")));
    check_vector(output, vector({7.0, 7.0}));
    CHECK(preconditioner.statistics().applications == 0);
    CHECK_THROWS_MATCHES((backend::ScaledJacobianOperator{wrong_size, physical.map}), std::invalid_argument, Catch::Matchers::MessageMatches(Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR")));
    CHECK_THROWS_MATCHES((backend::ScaledInverseOperator{wrong_size, physical.map}), std::invalid_argument, Catch::Matchers::MessageMatches(Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR")));
}

TEST_CASE("Physical Riesz scaling collapses a forty-eight-decade diagonal imbalance", "[normalization][unit][numerics]")
{
    const auto map = normalization::make_diagonal_normalization({1.0e-12, 1.0, 1.0e12}, {1.0e12, 1.0, 1.0e-12});
    REQUIRE(map);
    mfem::DenseMatrix physical(3);
    physical = 0.0;
    for (int index = 0; index < 3; ++index)
    {
        physical(index, index) = map->state_factors()[index] / map->residual_factors()[index];
    }
    CHECK(physical(2, 2) / physical(0, 0) == Catch::Approx(1.0e48));
    const backend::ScaledJacobianOperator jacobian(physical, *map);
    const mfem::Vector direction = vector({-2.0, 3.5, 0.125});
    mfem::Vector action;
    jacobian.Mult(direction, action);
    check_vector(action, direction, 4.0e-15);
}

TEST_CASE("Rectangular normalized operators and their transposes preserve independent coordinate sizes", "[normalization][unit][linear-algebra]")
{
    const auto map = normalization::make_diagonal_normalization({2.0, 3.0}, {5.0, 7.0, 11.0});
    REQUIRE(map);
    mfem::DenseMatrix physical(3, 2);
    mfem::DenseMatrix physical_inverse(2, 3);
    for (int row = 0; row < 3; ++row)
    {
        for (int column = 0; column < 2; ++column)
        {
            const double entry = 1.0 + 2.0 * row + column;
            physical(row, column) = entry * map->state_factors()[column] / map->residual_factors()[row];
            physical_inverse(column, row) = entry * map->residual_factors()[row] / map->state_factors()[column];
        }
    }
    const backend::ScaledJacobianOperator jacobian(physical, *map);
    const backend::ScaledInverseOperator inverse(physical_inverse, *map);
    mfem::Vector action;
    jacobian.Mult(vector({2.0, -0.5}), action);
    check_vector(action, vector({1.0, 4.0, 7.0}));
    jacobian.MultTranspose(vector({1.0, 2.0, 3.0}), action);
    check_vector(action, vector({22.0, 28.0}));
    inverse.Mult(vector({1.0, 2.0, 3.0}), action);
    check_vector(action, vector({22.0, 28.0}));
    inverse.MultTranspose(vector({2.0, -0.5}), action);
    check_vector(action, vector({1.0, 4.0, 7.0}));

    DenseInverseSolver physical_solver(physical_inverse);
    backend::ScaledPreconditioner preconditioner(physical_solver, jacobian);
    CHECK(physical_solver.bound_operator == &physical);
    action.SetSize(2);
    preconditioner.Mult(vector({1.0, 2.0, 3.0}), action);
    check_vector(action, vector({22.0, 28.0}));
    action.SetSize(3);
    preconditioner.MultTranspose(vector({2.0, -0.5}), action);
    check_vector(action, vector({1.0, 4.0, 7.0}));
}

TEST_CASE("Normalized adapters own frozen factors and support in-place applications", "[normalization][unit][lifetime]")
{
    PhysicalSystem physical;
    const backend::ScaledJacobianOperator jacobian(physical.jacobian, physical.map);
    physical.map = *normalization::make_diagonal_normalization({1.0, 1.0}, {1.0, 1.0});
    CHECK(jacobian.normalization().state_factors()[0] == 1.0e-9);
    CHECK(&jacobian.physical_operator() == &physical.jacobian);
    mfem::Vector values = vector({0.75, -1.25});
    jacobian.Mult(values, values);
    check_vector(values, vector({1.75, -2.25}));
    values = vector({0.75, -1.25});
    jacobian.MultTranspose(values, values);
    check_vector(values, vector({0.5, -3.0}));
}

TEST_CASE("Physical iterative inverses receive fresh zero correction guesses", "[normalization][unit][solver]")
{
    const PhysicalSystem physical;
    const backend::ScaledJacobianOperator jacobian(physical.jacobian, physical.map);
    DenseInverseSolver inverse(physical.inverse);
    inverse.iterative_mode = true;
    backend::ScaledPreconditioner preconditioner(inverse, jacobian);
    mfem::Vector correction(2);
    correction = 41.0;
    preconditioner.Mult(vector({1.5, -0.75}), correction);
    check_vector(correction, vector({0.525, -0.6}));
    preconditioner.Mult(vector({1.5, -0.75}), correction);
    check_vector(correction, vector({0.525, -0.6}));
}

TEST_CASE("Scaled preconditioners bind an initially unconfigured physical solver before checking its dimensions", "[normalization][unit][solver]")
{
    mfem::DenseMatrix physical(2);
    physical(0, 0) = 4.0;
    physical(0, 1) = 1.0;
    physical(1, 0) = 1.0;
    physical(1, 1) = 3.0;
    const auto map = normalization::make_diagonal_normalization({2.0, 5.0}, {7.0, 11.0});
    REQUIRE(map);
    const backend::ScaledJacobianOperator jacobian(physical, *map);
    mfem::CGSolver inverse;
    inverse.SetRelTol(1.0e-14);
    inverse.SetAbsTol(1.0e-15);
    inverse.SetMaxIter(4);
    inverse.SetPrintLevel(0);
    CHECK(inverse.Width() == 0);
    CHECK(inverse.Height() == 0);

    backend::ScaledPreconditioner preconditioner(inverse, jacobian);
    CHECK(inverse.Width() == 2);
    CHECK(inverse.Height() == 2);
    CHECK(preconditioner.statistics().operator_bindings == 1);
    mfem::Vector correction(2);
    preconditioner.Mult(vector({2.45, -3.575}), correction);
    CHECK(inverse.GetConverged());
    check_vector(correction, vector({0.25, -0.75}));

    preconditioner.SetOperator(jacobian);
    CHECK(preconditioner.statistics().operator_bindings == 2);
}

TEST_CASE("Normalization reports MFEM index overflow before narrowing", "[normalization][unit][validation]")
{
    const std::size_t largest_index = static_cast<std::size_t>(std::numeric_limits<int>::max());
    CHECK(backend::detail::mfem_size(0) == 0);
    CHECK(backend::detail::mfem_size(largest_index) == std::numeric_limits<int>::max());
    CHECK_THROWS_MATCHES(backend::detail::mfem_size(largest_index + 1), std::invalid_argument, Catch::Matchers::MessageMatches(Catch::Matchers::ContainsSubstring("INVALID_DOF_COUNT")));
}