#include <array>
#include <expected>
#include <memory>
#include <type_traits>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

#include "serif/solver/preconditioner/preconditioner.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace
{
    namespace pc = serif::solver::preconditioner;
    namespace dependencies = serif::utils::dependencies;
    namespace physical = serif::operators::backend::mfem;
    using Code = serif::utils::errors::SERiFErrorCode;

    struct Density
    {
    };
    struct DensityEquation
    {
    };
    struct PhysicalScale
    {
    };
    struct PreparedScale
    {
    };
    using DensityKeys = dependencies::KeyList<Density>;
    using DensityRows = dependencies::KeyList<DensityEquation>;
    using DensityBlock = pc::keys::Block<DensityRows, DensityKeys>;
} // namespace

TEST_CASE("Explicit upstream computations retain move-only owners and share ordinary preconditioner invalidation", "[unit][preconditioner][graph]")
{
    using Layout = serif::discretization::blocks::SystemLayout<DensityKeys, DensityRows>;
    const Layout layout{.state_sizes = {2}, .residual_sizes = {2}};
    const auto choice = pc::make_preconditioner<Density, DensityEquation>(pc::strategies::inverse::Diagonal{});
    auto prepared_block = dependencies::make_computation<DensityBlock, dependencies::InputList<PreparedScale>>(
        [coefficients = std::make_unique<std::array<double, 2>>(std::array<double, 2>{2.0, 3.0})](double scale) -> std::expected<physical::LinearizedBlock, pc::Failure>
    {
        if (scale <= 0.0)
        {
            return std::unexpected(pc::Failure{Code::PRECONDITIONER_INVALID_CALIBRATION, "The upstream physical scale must be positive."});
        }
        auto matrix = std::make_shared<mfem::DenseMatrix>(2);
        *matrix = 0.0;
        (*matrix)(0, 0) = (*coefficients)[0] * scale;
        (*matrix)(1, 1) = (*coefficients)[1] * scale;
        return physical::linearized_block(std::move(matrix));
    });
    static_assert(!std::copy_constructible<decltype(prepared_block)>);
    auto prepared_scale = dependencies::make_computation<PreparedScale, dependencies::InputList<PhysicalScale>>([](double scale)
    {
        return 2.0 * scale;
    });
    // The explicitly supplied nodes are deliberately in the opposite order from their numerical dependency.
    auto graph = pc::make_preconditioner_graph(choice, layout, std::move(prepared_block), std::move(prepared_scale));
    double physical_scale = 1.0;
    const auto context = dependencies::make_dependency_context(dependencies::borrow<PhysicalScale>(physical_scale), dependencies::borrow<pc::keys::Layout>(layout));
    auto cache = dependencies::make_computation_cache<pc::Failure>(std::move(graph), context);
    REQUIRE(cache.try_prepare(context));
    const auto accepted = cache.share<pc::keys::Preconditioner>();
    mfem::Vector rhs(2), result(2);
    rhs[0] = 6.0;
    rhs[1] = 12.0;
    accepted->action->Mult(rhs, result);
    CHECK(result[0] == Catch::Approx(1.5));
    CHECK(result[1] == Catch::Approx(2.0));
    const auto reused = cache.try_prepare(context);
    REQUIRE(reused);
    CHECK(reused->computed_count() == 0);
    CHECK(cache.share<pc::keys::Preconditioner>() == accepted);

    physical_scale = 2.0;
    const auto changed = cache.try_prepare(context);
    REQUIRE(changed);
    CHECK(changed->status<PreparedScale>() == dependencies::PreparationStatus::computed);
    CHECK(changed->status<DensityBlock>() == dependencies::PreparationStatus::computed);
    CHECK(changed->status<pc::keys::Preconditioner>() == dependencies::PreparationStatus::computed);
    cache.get<pc::keys::Preconditioner>().action->Mult(rhs, result);
    CHECK(result[0] == Catch::Approx(0.75));
    CHECK(result[1] == Catch::Approx(1.0));

    physical_scale = 0.0;
    const auto rejected = cache.try_prepare(context);
    REQUIRE_FALSE(rejected);
    CHECK(rejected.error().error.code == Code::PRECONDITIONER_INVALID_CALIBRATION);
    CHECK(rejected.error().error.message == "The upstream physical scale must be positive.");
    CHECK_FALSE(cache.ready());
    accepted->action->Mult(rhs, result);
    CHECK(result[0] == Catch::Approx(1.5));
    CHECK(result[1] == Catch::Approx(2.0));
    physical_scale = 3.0;
    REQUIRE(cache.try_prepare(context));
    cache.get<pc::keys::Preconditioner>().action->Mult(rhs, result);
    CHECK(result[0] == Catch::Approx(0.5));
    CHECK(result[1] == Catch::Approx(2.0 / 3.0));
}