#include <array>
#include <limits>
#include <memory>
#include <optional>
#include <vector>
#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include <catch2/matchers/catch_matchers_string.hpp>
#include "mfem/linalg/densemat.hpp"
#include "serif/backend/mfem/model/system.hpp"

namespace {
    namespace backend = serif::model::backend::mfem;
    using Catch::Matchers::WithinAbs;

    /** @brief Make small independent matrices whose entries reveal coordinate permutations. */
    [[nodiscard]] backend::LinearizedBlock matrix(int rows, int columns, std::initializer_list<double> entries) {
        auto values = std::make_shared<::mfem::DenseMatrix>(rows, columns);
        auto entry = entries.begin();
        for (int row = 0; row < rows; ++row) {
            for (int column = 0; column < columns; ++column) {
                (*values)(row, column) = *entry++;
            }
        }
        return serif::operators::backend::mfem::linearized_block(std::shared_ptr<const ::mfem::DenseMatrix>(values));
    }

    /** @brief Compare all coordinate entries, independent of the implementation's block selection. */
    void check_vector(const ::mfem::Vector& actual, std::initializer_list<double> expected) {
        REQUIRE(actual.Size() == static_cast<int>(expected.size()));
        int index = 0;
        for (const double value : expected) {
            CHECK_THAT(actual[index++], WithinAbs(value, 1.0e-13));
        }
    }
} // namespace

TEST_CASE("Model runtime adds generated physics before replacing complete boundary rows", "[model][runtime][unit]") {
    const auto accepted = backend::make_physical_system(4, 5,
        {{0, {1.0, 2.0, 3.0}}, {1, {10.0, 20.0, 30.0}}},
        {{0, 2, matrix(3, 3, {1, 2, 3, 4, 5, 6, 7, 8, 9})},
            {1, 0, matrix(3, 2, {10, 11, 12, 13, 14, 15})},
            {0, 1, matrix(4, 1, {20, 21, 22, 23})}},
        {{{3, 1}, {300, 100}, {{2, matrix(2, 3, {2, 3, 4, 5, 6, 7})}}},
            {{0}, {400}, {{0, matrix(1, 1, {8})}}}});
    CHECK(accepted->residual() == std::vector<double>{400, 100, 23, 300});
    const auto& jacobian = accepted->jacobian();
    REQUIRE(jacobian.has_assembly());
    CHECK_FALSE(jacobian.has_diagonal());
    const auto assembled = jacobian.assemble();
    const std::array<std::array<double, 5>, 4> expected{{{8, 0, 0, 0, 0}, {0, 0, 5, 6, 7}, {12, 35, 7, 8, 9}, {0, 0, 2, 3, 4}}};
    for (int row = 0; row < 4; ++row) {
        for (int column = 0; column < 5; ++column) {
            CHECK_THAT((*assembled)(row, column), WithinAbs(expected[row][column], 1.0e-14));
        }
    }
    ::mfem::Vector direction(5), dual(4), action, transpose;
    for (int index = 0; index < 5; ++index) {
        direction[index] = index + 1;
    }
    for (int index = 0; index < 4; ++index) {
        dual[index] = 0.25 + index;
    }
    jacobian.action()->Mult(direction, action);
    check_vector(action, {8, 74, 180, 38});
    jacobian.action()->MultTranspose(dual, transpose);
    CHECK_THAT(dual * action, WithinAbs(direction * transpose, 1.0e-12));
    const auto selected = accepted->block(1, 2, 2, 2);
    CHECK(selected.has_diagonal());
    CHECK(selected.diagonal() == std::vector<double>{5, 8});
    ::mfem::Vector local(2), selected_action;
    local[0] = 2;
    local[1] = 3;
    selected.action()->Mult(local, selected_action);
    check_vector(selected_action, {28, 38});
}

TEST_CASE("Model selected blocks retain owners and skip unrelated equation actions", "[model][runtime][unit]") {
    auto calls = std::make_shared<int>(0);
    auto owner = std::make_shared<int>(37);
    std::weak_ptr<const int> lifetime = owner;
    const auto unrelated = serif::operators::backend::mfem::linearized_action(1, 1, [calls](const auto& input, auto& output) {
        ++*calls;
        output[0] = input[0];
    }, [calls](const auto& input, auto& output) {
        ++*calls;
        output[0] = input[0];
    });
    auto accepted = backend::make_physical_system(3, 3, {}, {{0, 0, matrix(2, 2, {2, 1, -1, 3})}, {2, 2, unrelated}}, {}, owner);
    std::optional<backend::LinearizedBlock> selected = accepted->block(0, 2, 0, 2);
    owner.reset();
    accepted.reset();
    CHECK_FALSE(lifetime.expired());
    ::mfem::Vector input(2), output;
    input = 1.0;
    selected->action()->Mult(input, output);
    check_vector(output, {3, 2});
    CHECK(*calls == 0);
    CHECK(selected->has_diagonal());
    CHECK(selected->has_assembly());
    CHECK(selected->diagonal() == std::vector<double>{2, 3});
    selected.reset();
    CHECK(lifetime.expired());
}

TEST_CASE("Grouped model blocks preserve ordered ranges and boundary rows with one action per child", "[model][runtime][unit]") {
    const auto original = matrix(4, 5, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20});
    auto calls = std::make_shared<std::array<int, 2>>();
    const auto action = serif::operators::backend::mfem::linearized_action(4, 5, [calls, original](const auto& input, auto& output) {
        ++(*calls)[0];
        original.action()->Mult(input, output);
    }, [calls, original](const auto& input, auto& output) {
        ++(*calls)[1];
        original.action()->MultTranspose(input, output);
    });
    const backend::LinearizedBlock counted(action.action(), {}, [original]() {
        return original.assemble();
    });
    auto unrelated_calls = std::make_shared<int>(0);
    const auto unrelated_action = [unrelated_calls](const auto& input, auto& output) {
        ++*unrelated_calls;
        output = input;
    };
    auto owner = std::make_shared<int>(42);
    std::weak_ptr<const int> lifetime = owner;
    auto accepted = backend::make_physical_system(5, 6, {}, {{0, 0, counted}, {4, 5, serif::operators::backend::mfem::linearized_action(1, 1, unrelated_action, unrelated_action)}},
        {{{1, 3}, {0.0, 0.0}, {{0, matrix(2, 5, {101, 102, 103, 104, 105, 106, 107, 108, 109, 110})}}}}, owner);
    const std::array rows{backend::CoordinateRange{2, 2}, backend::CoordinateRange{0, 2}};
    const std::array columns{backend::CoordinateRange{3, 2}, backend::CoordinateRange{0, 2}};
    std::optional<backend::LinearizedBlock> selected = accepted->block(rows, columns);
    accepted.reset();
    owner.reset();
    CHECK_FALSE(lifetime.expired());

    const std::array<std::array<double, 4>, 4> expected{{{14, 15, 11, 12}, {109, 110, 106, 107}, {4, 5, 1, 2}, {104, 105, 101, 102}}};
    const auto assembled = selected->assemble();
    for (int row = 0; row < 4; ++row) {
        for (int column = 0; column < 4; ++column) {
            CHECK((*assembled)(row, column) == expected[row][column]);
        }
    }
    CHECK(selected->diagonal() == std::vector<double>{14, 110, 1, 102});
    CHECK(*calls == std::array<int, 2>{0, 0});
    ::mfem::Vector input(4), dual(4), result, transpose;
    for (int index = 0; index < 4; ++index) {
        input[index] = index + 1;
        dual[index] = index - 0.5;
    }
    selected->action()->Mult(input, result);
    selected->action()->MultTranspose(dual, transpose);
    check_vector(result, {125, 1075, 25, 1025});
    CHECK_THAT(result * dual, WithinAbs(input * transpose, 1.0e-12));
    CHECK(*calls == std::array<int, 2>{1, 1});
    CHECK(*unrelated_calls == 0);
    selected.reset();
    CHECK(lifetime.expired());
}

TEST_CASE("Grouped model selections retain native diagonals and exact transposes for repeated ranges", "[model][runtime][unit]") {
    const auto original = matrix(3, 3, {1, 2, 3, 4, 5, 6, 7, 8, 9});
    const backend::LinearizedBlock diagonal_only(original.action(), [original]() {
        return original.diagonal();
    });
    const auto accepted = backend::make_physical_system(3, 3, {}, {{0, 0, diagonal_only}});
    const std::array ranges{backend::CoordinateRange{2, 1}, backend::CoordinateRange{0, 2}, backend::CoordinateRange{2, 1}};
    const auto selected = accepted->block(ranges, ranges);
    CHECK_FALSE(selected.has_assembly());
    REQUIRE(selected.has_diagonal());
    CHECK(selected.diagonal() == std::vector<double>{9, 1, 5, 9});
    ::mfem::Vector input(4), dual(4), output, transpose;
    for (int index = 0; index < 4; ++index) {
        input[index] = index + 1;
        dual[index] = 0.5 * index - 0.3;
    }
    selected.action()->Mult(input, output);
    selected.action()->MultTranspose(dual, transpose);
    check_vector(output, {83, 23, 53, 83});
    CHECK_THAT(output * dual, WithinAbs(input * transpose, 1.0e-12));

    const auto empty = accepted->block(std::span<const backend::CoordinateRange>{}, ranges);
    CHECK(empty.action()->Height() == 0);
    CHECK(empty.action()->Width() == 4);
    empty.action()->Mult(input, output);
    CHECK(output.Size() == 0);
}

TEST_CASE("Model runtime actions permit both shifted buffer overlaps", "[model][runtime][unit]") {
    const auto accepted = backend::make_physical_system(2, 2, {}, {{0, 0, matrix(2, 2, {2, 1, -1, 3})}});
    for (const bool transpose : {false, true}) {
        for (const bool forward_overlap : {false, true}) {
            std::array<double, 3> storage{};
            ::mfem::Vector input(storage.data() + (forward_overlap ? 0 : 1), 2);
            ::mfem::Vector output(storage.data() + (forward_overlap ? 1 : 0), 2);
            input[0] = 4;
            input[1] = 5;
            if (transpose) {
                accepted->jacobian().action()->MultTranspose(input, output);
                check_vector(output, {3, 19});
            }
            else {
                accepted->jacobian().action()->Mult(input, output);
                check_vector(output, {13, 11});
            }
        }
    }
}

TEST_CASE("Model runtime rejects malformed contribution and replacement coordinates", "[model][runtime][unit]") {
    const auto unit = matrix(1, 1, {1});
    CHECK_THROWS(backend::make_physical_system(-1, 1, {}, {}));
    CHECK_THROWS(backend::make_physical_system(1, 1, {{1, {1}}}, {}));
    CHECK_THROWS(backend::make_physical_system(1, 1, {}, {{0, 1, unit}}));
    CHECK_THROWS(backend::make_physical_system(1, 1, {}, {}, {{{0, 0}, {1, 2}, {}}}));
    CHECK_THROWS(backend::make_physical_system(1, 1, {}, {}, {{{0}, {}, {}}}));
    CHECK_THROWS(backend::make_physical_system(1, 1, {}, {}, {{{1}, {0}, {}}}));
    CHECK_THROWS(backend::make_physical_system(2, 1, {}, {}, {{{0, 1}, {0, 0}, {{0, unit}}}}));
    CHECK_THROWS(backend::make_physical_system(1, 1, {{0, {std::numeric_limits<double>::quiet_NaN()}}}, {}));
    CHECK_THROWS(backend::make_physical_system(1, 1, {}, {}, {{{0}, {std::numeric_limits<double>::infinity()}, {}}}));
    const auto accepted = backend::make_physical_system(1, 2, {}, {});
    CHECK_THROWS(accepted->block(0, 2, 0, 1));
    const std::array row{backend::CoordinateRange{0, 1}};
    for (const backend::CoordinateRange invalid : {backend::CoordinateRange{-1, 1}, {0, -1}, {2, 1}, {std::numeric_limits<int>::max(), 1}}) {
        CHECK_THROWS(accepted->block(row, std::array{invalid}));
    }
    const auto wide = backend::make_physical_system(0, std::numeric_limits<int>::max(), {}, {});
    CHECK_THROWS_WITH(wide->block(std::span<const backend::CoordinateRange>{}, std::array{backend::CoordinateRange{0, std::numeric_limits<int>::max()}, backend::CoordinateRange{0, 1}}),
        Catch::Matchers::ContainsSubstring("INVALID_DOF_COUNT"));
    ::mfem::Vector wrong(3), result(1);
    wrong = 0.0;
    result = 19;
    CHECK_THROWS(accepted->jacobian().action()->Mult(wrong, result));
    CHECK(result[0] == 19);
}

TEST_CASE("Model action composition preserves forward transpose and native assembly", "[model][runtime][unit]") {
    const auto left = matrix(2, 3, {1, 2, 3, 4, 5, 6});
    const auto right = matrix(3, 2, {2, -1, 3, 2, 1, 4});
    const auto product = serif::operators::backend::mfem::compose_blocks(left, right);
    const auto result = serif::operators::backend::mfem::sum_blocks(2, 2, {product, serif::operators::backend::mfem::transpose_block(product)});
    REQUIRE(result.has_assembly());
    const auto assembled = result.assemble();
    CHECK((*assembled)(0, 0) == 22);
    CHECK((*assembled)(0, 1) == 44);
    CHECK((*assembled)(1, 0) == 44);
    CHECK((*assembled)(1, 1) == 60);
    CHECK_THROWS(serif::operators::backend::mfem::compose_blocks(left, left));
    CHECK_THROWS(serif::operators::backend::mfem::sum_blocks(1, 1, {left}));
}

TEST_CASE("Model native diagonals do not probe actions or require sparse assembly", "[model][runtime][unit]") {
    auto actions = std::make_shared<int>(0);
    auto diagonals = std::make_shared<int>(0);
    const auto action = serif::operators::backend::mfem::linearized_action(2, 2, [actions](const auto& input, auto& output) {
        ++*actions;
        output = input;
    }, [actions](const auto& input, auto& output) {
        ++*actions;
        output = input;
    });
    const backend::LinearizedBlock block(action.action(), [diagonals]() {
        ++*diagonals;
        return std::vector<double>{1, 1};
    });
    const auto accepted = backend::make_physical_system(2, 2, {}, {{0, 0, block}, {0, 0, block}});
    REQUIRE(accepted->jacobian().has_diagonal());
    CHECK_FALSE(accepted->jacobian().has_assembly());
    CHECK(accepted->jacobian().diagonal() == std::vector<double>{2, 2});
    CHECK(*actions == 0);
    CHECK(*diagonals == 2);
}

TEST_CASE("Model composition and selection reject malformed physical block callbacks", "[model][runtime][unit]") {
    const auto original = matrix(2, 2, {1, 0, 0, 1});
    const auto check_rejection = [&original](backend::LinearizedBlock invalid) {
        const auto accepted = backend::make_physical_system(2, 2, {}, {{0, 0, invalid}});
        CHECK_THROWS_WITH(accepted->jacobian().assemble(), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
        CHECK_THROWS_WITH(serif::operators::backend::mfem::transpose_block(invalid).assemble(), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
        CHECK_THROWS_WITH(serif::operators::backend::mfem::compose_blocks(original, invalid).assemble(), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
        CHECK_THROWS_WITH(serif::operators::backend::mfem::compose_blocks(invalid, original).assemble(), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
        CHECK_THROWS_WITH(serif::operators::backend::mfem::sum_blocks(2, 2, {original, invalid}).assemble(), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
    };
    check_rejection(backend::LinearizedBlock(original.action(), {}, []() {
        return std::shared_ptr<const ::mfem::SparseMatrix>{};
    }));
    check_rejection(backend::LinearizedBlock(original.action(), {}, []() {
        return std::make_shared<const ::mfem::SparseMatrix>(2);
    }));
    check_rejection(backend::LinearizedBlock(original.action(), {}, []() {
        return matrix(1, 2, {1, 0}).assemble();
    }));
    const backend::LinearizedBlock invalid_diagonal(original.action(), []() {
        return std::vector<double>{1};
    });
    const auto accepted = backend::make_physical_system(2, 2, {}, {{0, 0, invalid_diagonal}});
    CHECK_THROWS_WITH(accepted->jacobian().diagonal(), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
    CHECK_THROWS_WITH(serif::operators::backend::mfem::transpose_block(invalid_diagonal).diagonal(), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
    CHECK_THROWS_WITH(serif::operators::backend::mfem::sum_blocks(2, 2, {original, invalid_diagonal}).diagonal(), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
}

TEST_CASE("Model invalid callback outputs never overwrite caller storage", "[model][runtime][unit]") {
    const auto transpose = [](const auto& input, auto& output) {
        output = input;
    };
    const auto wrong_size = serif::operators::backend::mfem::linearized_action(1, 1, [](const auto&, auto& output) {
        output.SetSize(0);
    }, transpose);
    const auto nonfinite = serif::operators::backend::mfem::linearized_action(1, 1, [](const auto&, auto& output) {
        output[0] = std::numeric_limits<double>::infinity();
    }, transpose);
    ::mfem::Vector input(1), output(1);
    input = 1.0;
    output = 42.0;
    CHECK_THROWS(wrong_size.action()->Mult(input, output));
    CHECK(output[0] == 42.0);
    CHECK_THROWS(nonfinite.action()->Mult(input, output));
    CHECK(output[0] == 42.0);
    input = std::numeric_limits<double>::quiet_NaN();
    CHECK_THROWS(nonfinite.action()->MultTranspose(input, output));
    CHECK(output[0] == 42.0);
}

TEST_CASE("Complete model actions traverse each equation once while selected blocks keep their native representation", "[model][runtime][unit]") {
    auto calls = std::make_shared<std::array<int, 4>>();
    const auto counted = [calls](backend::LinearizedBlock block, const int index) {
        const auto action = serif::operators::backend::mfem::linearized_action(block.action()->Height(), block.action()->Width(),
            [calls, block, index](const auto& input, auto& output) {
            ++(*calls)[index];
            block.action()->Mult(input, output);
        },
            [calls, block, index](const auto& input, auto& output) {
            ++(*calls)[index];
            block.action()->MultTranspose(input, output);
        });
        return backend::LinearizedBlock(action.action(), {}, [block]() {
            return block.assemble();
        });
    };
    const auto left = counted(matrix(3, 2, {1, 2, 3, 4, 5, 6}), 0);
    const auto right = counted(matrix(3, 2, {7, 8, 9, 10, 11, 12}), 1);
    const auto complete_matrix = matrix(3, 4, {1, 2, 7, 8, 3, 4, 9, 10, 5, 6, 11, 12});
    const auto complete = serif::operators::backend::mfem::linearized_action(3, 4, [calls, complete_matrix](const auto& input, auto& output) {
        ++(*calls)[2];
        complete_matrix.action()->Mult(input, output);
    }, [calls, complete_matrix](const auto& input, auto& output) {
        ++(*calls)[3];
        complete_matrix.action()->MultTranspose(input, output);
    });
    const auto accepted = backend::make_physical_system(3, 4, {}, {{0, 0, left}, {0, 2, right}},
        {{{1}, {0.0}, {{0, matrix(1, 4, {2, 3, 4, 5})}}}}, {},
        std::vector<backend::JacobianContribution>{{0, 0, complete}});
    mfem::Vector direction(4), dual(3), action, transpose;
    for (int column = 0; column < 4; ++column) {
        direction[column] = column + 1;
    }
    dual[0] = 0.2;
    dual[1] = -0.7;
    dual[2] = 0.4;
    accepted->jacobian().action()->Mult(direction, action);
    accepted->jacobian().action()->MultTranspose(dual, transpose);
    check_vector(action, {58, 40, 98});
    CHECK_THAT(action * dual, WithinAbs(direction * transpose, 1.0e-12));
    CHECK(*calls == std::array<int, 4>{0, 0, 1, 1});
    REQUIRE(accepted->jacobian().has_assembly());
    const auto assembled = accepted->jacobian().assemble();
    mfem::Vector assembled_action(3), assembled_transpose(4);
    assembled->Mult(direction, assembled_action);
    assembled->MultTranspose(dual, assembled_transpose);
    assembled_action -= action;
    assembled_transpose -= transpose;
    CHECK(assembled_action.Norml2() == 0.0);
    CHECK(assembled_transpose.Norml2() < 1.0e-13);
    CHECK(*calls == std::array<int, 4>{0, 0, 1, 1});

    const auto selected = accepted->block(0, 3, 0, 2);
    mfem::Vector local(2), selected_action;
    local[0] = 1.0;
    local[1] = 2.0;
    selected.action()->Mult(local, selected_action);
    check_vector(selected_action, {5, 8, 17});
    CHECK(*calls == std::array<int, 4>{1, 0, 1, 1});
    CHECK(accepted->block(0, 2, 0, 2).diagonal() == std::vector<double>{1, 3});
    CHECK(*calls == std::array<int, 4>{1, 0, 1, 1});
    CHECK_THROWS(backend::make_physical_system(2, 4, {}, {}, {}, {}, std::vector<backend::JacobianContribution>{{0, 0, complete}}));
    CHECK_THROWS(backend::make_physical_system(3, 3, {}, {}, {}, {}, std::vector<backend::JacobianContribution>{{0, 0, complete}}));
}