#include <limits>
#include <memory>
#include <optional>
#include <utility>
#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

#include "mfem/linalg/densemat.hpp"
#include "serif/backend/mfem/model/scalar_row.hpp"
#include "serif/backend/mfem/model/system.hpp"

TEST_CASE("A scalar residual row retains its exact functional without retaining or repeating the source action", "[model][scalar_row][unit]") {
    namespace backend = serif::model::backend::mfem;
    namespace execution = serif::operators::backend::mfem;
    int forward_calls = 0, transpose_calls = 0;
    std::optional<execution::LinearizedBlock> cached;
    std::weak_ptr<int> lifetime;
    {
        auto owner = std::make_shared<int>(1);
        lifetime = owner;
        const auto source = execution::linearized_action(1, 3, [&, owner](const auto& input, auto& output) {
            ++forward_calls;
            output[0] = 2.0 * input[0] - 3.0 * input[1] + 5.0 * input[2];
        }, [&, owner](const auto& input, auto& output) {
            ++transpose_calls;
            output[0] = 2.0 * input[0];
            output[1] = -3.0 * input[0];
            output[2] = 5.0 * input[0];
        });
        auto result = backend::cache_scalar_row(source);
        REQUIRE(result);
        cached.emplace(std::move(*result));
    }
    CHECK(lifetime.expired());
    CHECK(forward_calls == 0);
    CHECK(transpose_calls == 1);
    const auto system = backend::make_physical_system(1, 3, {}, {{0, 0, *cached}});
    mfem::Vector direction(3), action, transpose, dual(1);
    direction[0] = 7.0;
    direction[1] = 11.0;
    direction[2] = 13.0;
    dual[0] = 0.25;
    system->jacobian().action()->Mult(direction, action);
    CHECK(action[0] == Catch::Approx(46.0));
    system->jacobian().action()->MultTranspose(dual, transpose);
    CHECK(transpose[0] == Catch::Approx(0.5));
    CHECK(transpose[1] == Catch::Approx(-0.75));
    CHECK(transpose[2] == Catch::Approx(1.25));
    const auto selected = system->block(0, 1, 1, 2);
    mfem::Vector selected_direction(direction.GetData() + 1, 2);
    selected.action()->Mult(selected_direction, action);
    CHECK(action[0] == Catch::Approx(32.0));
    const auto matrix = selected.assemble();
    CHECK((*matrix)(0, 0) == Catch::Approx(-3.0));
    CHECK((*matrix)(0, 1) == Catch::Approx(5.0));
    CHECK(forward_calls == 0);
    CHECK(transpose_calls == 1);
}

TEST_CASE("Scalar row caching rejects a spatial residual without invoking its transpose", "[model][scalar_row][unit]") {
    namespace execution = serif::operators::backend::mfem;
    int calls = 0;
    const auto source = execution::linearized_action(2, 3, [](const auto&, auto&) {}, [&](const auto&, auto&) {
        ++calls;
    });
    CHECK_THROWS(serif::model::backend::mfem::cache_scalar_row(source));
    CHECK(calls == 0);
}

TEST_CASE("Scalar row caching preserves numerical failure classification for nonfinite derivatives", "[model][scalar_row][unit]") {
    namespace execution = serif::operators::backend::mfem;
    for (const double value : {std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::infinity()}) {
        auto matrix = std::make_shared<mfem::DenseMatrix>(1, 1);
        (*matrix)(0, 0) = value;
        // A raw matrix deliberately bypasses the checked callback adapter to exercise the cache guard.
        const auto source = execution::linearized_block(std::shared_ptr<const mfem::DenseMatrix>(matrix));
        const auto raw_result = serif::model::backend::mfem::cache_scalar_row(source);
        REQUIRE_FALSE(raw_result);
        CHECK(raw_result.error().reason == execution::PreparationRejectionReason::NON_FINITE);
        const auto checked = execution::linearized_action(1, 1, [source](const auto& input, auto& output) {
            source.action()->Mult(input, output);
        }, [source](const auto& input, auto& output) {
            source.action()->MultTranspose(input, output);
        });
        const auto checked_result = serif::model::backend::mfem::cache_scalar_row(checked);
        REQUIRE_FALSE(checked_result);
        CHECK(checked_result.error().reason == execution::PreparationRejectionReason::NON_FINITE);
    }
}