#include <algorithm>
#include <array>
#include <expected>
#include <initializer_list>
#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;
    namespace coupling = pc::strategies::coupling;
    using Inverse = pc::backend::mfem::Inverse;
    using Code = serif::utils::errors::SERiFErrorCode;

    struct Density
    {
    };
    struct Temperature
    {
    };
    struct HeatBalance
    {
    };
    using DensityKeys = dependencies::KeyList<Density>;
    using HeatKeys = dependencies::KeyList<HeatBalance>;
    using TemperatureKeys = dependencies::KeyList<Temperature>;
    using DensityBlock = pc::keys::Block<DensityKeys, DensityKeys>;
    using HeatBlock = pc::keys::Block<HeatKeys, TemperatureKeys>;

    /** @brief Create independently owned diagonal or one-entry physical matrices. */
    [[nodiscard]] std::shared_ptr<const mfem::DenseMatrix> diagonal(std::initializer_list<double> values)
    {
        auto matrix = std::make_shared<mfem::DenseMatrix>(values.size());
        *matrix = 0.0;
        int index = 0;
        for (double value : values)
        {
            (*matrix)(index, index) = value;
            ++index;
        }
        return matrix;
    }

    /** @brief Apply a compiled inverse without leaving initialization to the implementation. */
    [[nodiscard]] mfem::Vector apply(const Inverse& inverse, std::initializer_list<double> coefficients)
    {
        mfem::Vector input(coefficients.size());
        std::copy(coefficients.begin(), coefficients.end(), input.GetData());
        mfem::Vector result(inverse.action->Height());
        result = 17.0;
        inverse.action->Mult(input, result);
        return result;
    }

    /** @brief A custom strategy's reported success must still satisfy the physical block contract. */
    struct CustomInverse final
    {
        std::expected<Inverse, pc::Failure> result;

        /** @brief Return a chosen result so malformed successful declarations can be tested independently. */
        [[nodiscard]] std::expected<Inverse, pc::Failure> try_inverse(const physical::LinearizedBlock&) const
        {
            return result;
        }
    };
} // namespace

TEST_CASE("Shared physical layouts reject incorrect constituent sizes even when packed totals agree", "[unit][preconditioner][graph]")
{
    using Layout = serif::discretization::blocks::SystemLayout<dependencies::KeyList<Density, Temperature>, dependencies::KeyList<HeatBalance, Density>>;
    Layout layout{.state_sizes = {1, 2}, .residual_sizes = {2, 1}};
    const auto choice = pc::couple(pc::make_preconditioner<Temperature, HeatBalance>(pc::strategies::inverse::Diagonal{}),
        pc::make_preconditioner<Density>(pc::strategies::inverse::Diagonal{}), coupling::SecondThenFirst{});
    const auto context = dependencies::make_dependency_context(
        dependencies::own<HeatBlock>(physical::linearized_block(diagonal({2.0, 4.0}))),
        dependencies::own<DensityBlock>(physical::linearized_block(diagonal({5.0}))),
        dependencies::own<pc::keys::Block<HeatKeys, DensityKeys>>(physical::linearized_block(std::make_shared<::mfem::DenseMatrix>(2, 1))),
        dependencies::borrow<pc::keys::Layout>(layout));
    auto cache = dependencies::make_computation_cache<pc::Failure>(pc::make_preconditioner_graph(choice, layout), context);
    REQUIRE(cache.try_prepare(context));
    const auto accepted = cache.share<pc::keys::Preconditioner>();

    SECTION("Unknown block counts are not interchangeable")
    {
        layout.state_sizes = {2, 1};
    }
    SECTION("Equation block counts are checked independently")
    {
        layout.residual_sizes = {1, 2};
    }
    SECTION("Matching wrong row and column partitions do not hide a physical mismatch")
    {
        layout.state_sizes = {2, 1};
        layout.residual_sizes = {1, 2};
    }
    const auto rejected = cache.try_prepare(context);
    REQUIRE_FALSE(rejected);
    CHECK(rejected.error().error.code == Code::PRECONDITIONER_SIZE_MISMATCH);
    CHECK_FALSE(cache.ready());
    CHECK_THROWS(cache.share<pc::keys::Preconditioner>());
    const auto old = apply(*accepted, {6.0, 8.0, 20.0});
    CHECK(old[0] == Catch::Approx(4.0));
    CHECK(old[1] == Catch::Approx(3.0));
    CHECK(old[2] == Catch::Approx(2.0));

    layout.state_sizes = {1, 2};
    layout.residual_sizes = {2, 1};
    const auto recovered = cache.try_prepare(context);
    REQUIRE(recovered);
    CHECK(recovered->status<pc::keys::Inverse<HeatKeys, TemperatureKeys>>() == dependencies::PreparationStatus::reused);
    CHECK(recovered->status<pc::keys::Inverse<DensityKeys, DensityKeys>>() == dependencies::PreparationStatus::reused);
}

TEST_CASE("Triangular preconditioner graphs require and track only their used cross coupling", "[unit][preconditioner][graph]")
{
    using Required = pc::keys::Block<DensityKeys, TemperatureKeys>;
    using Unused = pc::keys::Block<HeatKeys, DensityKeys>;
    const auto choice = pc::couple(pc::make_preconditioner<Density>(pc::strategies::inverse::Diagonal{}),
        pc::make_preconditioner<Temperature, HeatBalance>(pc::strategies::inverse::Diagonal{}), coupling::SecondThenFirst{});
    const auto graph = pc::make_preconditioner_graph(choice);
    using Graph = decltype(graph);
    using Result = decltype(choice)::Result;
    static_assert(dependencies::ContainsKey<Required, typename Graph::Inputs>);
    static_assert(!dependencies::ContainsKey<Unused, typename Graph::Inputs>);
    static_assert(dependencies::ContainsKey<DensityBlock, typename Graph::Inputs>);
    static_assert(dependencies::ContainsKey<HeatBlock, typename Graph::Inputs>);

    auto used = physical::linearized_block(diagonal({4.0}));
    const auto context = dependencies::make_dependency_context(
        dependencies::own<DensityBlock>(physical::linearized_block(diagonal({2.0}))),
        dependencies::own<HeatBlock>(physical::linearized_block(diagonal({3.0}))), dependencies::borrow<Required>(used));
    auto cache = dependencies::make_computation_cache<pc::Failure>(graph, context);
    REQUIRE(cache.try_prepare(context));
    const auto result = apply(cache.get<Result>(), {6.0, 12.0});
    CHECK(result[0] == Catch::Approx(-5.0));
    CHECK(result[1] == Catch::Approx(4.0));

    auto ignored = physical::linearized_block(diagonal({99.0}));
    const auto extended = dependencies::make_dependency_context(
        dependencies::own<DensityBlock>(physical::linearized_block(diagonal({2.0}))),
        dependencies::own<HeatBlock>(physical::linearized_block(diagonal({3.0}))),
        dependencies::borrow<Required>(used), dependencies::borrow<Unused>(ignored));
    auto extended_cache = dependencies::make_computation_cache<pc::Failure>(graph, extended);
    REQUIRE(extended_cache.try_prepare(extended));
    const auto accepted = extended_cache.share<Result>();
    ignored = physical::linearized_block(diagonal({-13.0}));
    const auto reused = extended_cache.try_prepare(extended);
    REQUIRE(reused);
    CHECK(reused->computed_count() == 0);
    CHECK(extended_cache.share<Result>() == accepted);

    used = physical::linearized_block(diagonal({2.0}));
    const auto refreshed = extended_cache.try_prepare(extended);
    REQUIRE(refreshed);
    CHECK(refreshed->source_changed<Required>());
    CHECK(refreshed->status<Result>() == dependencies::PreparationStatus::computed);
    CHECK(refreshed->status<pc::keys::Inverse<DensityKeys, DensityKeys>>() == dependencies::PreparationStatus::reused);
    CHECK(refreshed->status<pc::keys::Inverse<HeatKeys, TemperatureKeys>>() == dependencies::PreparationStatus::reused);
    const auto updated = apply(extended_cache.get<Result>(), {6.0, 12.0});
    CHECK(updated[0] == Catch::Approx(-1.0));
    CHECK(updated[1] == Catch::Approx(4.0));
    CHECK(apply(*accepted, {6.0, 12.0})[0] == Catch::Approx(result[0]));
    CHECK(apply(*accepted, {6.0, 12.0})[1] == Catch::Approx(result[1]));
}

TEST_CASE("Custom inverse graph results must preserve the selected physical action contract", "[unit][preconditioner][graph]")
{
    std::expected<Inverse, pc::Failure> result = Inverse{.action = diagonal({0.5, 0.25})};
    Code expected_code = Code::PRECONDITIONER_INVALID_OPERATOR;
    SECTION("Successful preparation cannot publish a null inverse")
    {
        result->action.reset();
    }
    SECTION("An inverse must reverse the selected physical dimensions")
    {
        result->action = diagonal({0.5});
        expected_code = Code::PRECONDITIONER_SIZE_MISMATCH;
    }
    SECTION("The declared application contract must be valid")
    {
        result->contract = static_cast<pc::ApplicationContract>(99);
        expected_code = Code::PRECONDITIONER_INVALID_OPTIONS;
    }
    SECTION("An originating custom failure propagates unchanged")
    {
        expected_code = Code::PRECONDITIONER_INVALID_CALIBRATION;
        result = std::unexpected(pc::Failure{expected_code, "Custom calibration failed once."});
    }
    const auto choice = pc::make_preconditioner<Temperature, HeatBalance>(CustomInverse{result});
    const auto context = dependencies::make_dependency_context(dependencies::own<HeatBlock>(physical::linearized_block(diagonal({2.0, 4.0}))));
    auto cache = dependencies::make_computation_cache<pc::Failure>(pc::make_preconditioner_graph(choice), context);
    const auto rejected = cache.try_prepare(context);
    REQUIRE_FALSE(rejected);
    CHECK(rejected.error().error.code == expected_code);
    if (expected_code == Code::PRECONDITIONER_INVALID_CALIBRATION)
    {
        CHECK(rejected.error().error.message == "Custom calibration failed once.");
    }
}

TEST_CASE("Empty physical constraint lists reuse the field recipe and accepted inverse", "[unit][preconditioner][graph]")
{
    const auto fields = pc::make_preconditioner<Density>(pc::strategies::inverse::Diagonal{});
    using Empty = dependencies::KeyList<>;
    const auto choice = pc::with_constraints<pc::ConstraintCoordinates<Empty, Empty>>(fields);
    static_assert(std::same_as<decltype(fields), decltype(choice)>);
    const auto context = dependencies::make_dependency_context(dependencies::own<DensityBlock>(physical::linearized_block(diagonal({2.0}))));
    auto cache = dependencies::make_computation_cache<pc::Failure>(pc::make_preconditioner_graph(choice), context);
    REQUIRE(cache.try_prepare(context));
    CHECK(apply(cache.get<typename decltype(choice)::Result>(), {8.0})[0] == Catch::Approx(4.0));
}