#include <array>
#include <memory>
#include <span>

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

#include "serif/backend/mfem/operators/actions.hpp"
#include "serif/solver/preconditioner/preconditioner.hpp"
#include "serif/backend/mfem/solver/normalization/operators.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"

namespace
{
    namespace pc = serif::solver::preconditioner;
    namespace dependencies = serif::utils::dependencies;
    namespace operators = serif::operators::backend::mfem;
    namespace inverse = pc::strategies::inverse;
    namespace coupling = pc::strategies::coupling;
    using Inverse = pc::backend::mfem::Inverse;
    using Code = serif::utils::errors::SERiFErrorCode;

    struct Temperature
    {
    };
    struct HeatBalance
    {
    };
    struct Density
    {
    };
    struct MassConstraint
    {
    };
    struct MassParameter
    {
    };
    struct Potential
    {
    };
    struct Gradient
    {
    };

    /** @brief Construct owned matrices independently of the preconditioner factorizations. */
    std::shared_ptr<const mfem::DenseMatrix> matrix(const int rows, const int columns, std::initializer_list<double> values)
    {
        auto result = std::make_shared<mfem::DenseMatrix>(rows, columns);
        auto value = values.begin();
        for (int row = 0; row < rows; ++row)
        {
            for (int column = 0; column < columns; ++column)
            {
                (*result)(row, column) = *value++;
            }
        }
        return result;
    }

    /** @brief Apply using preallocated MFEM coordinates, retaining the factory's zero-guess contract. */
    mfem::Vector apply(const Inverse& inverse, std::initializer_list<double> values)
    {
        mfem::Vector input(values.size());
        std::copy(values.begin(), values.end(), input.GetData());
        mfem::Vector output(inverse.action->Height());
        output = 91.0;
        inverse.action->Mult(input, output);
        return output;
    }
} // namespace

TEST_CASE("A physical inverse graph reuses accepted results and preserves a frozen solve after failed preparation", "[unit][preconditioner][graph]")
{
    const auto choice = pc::make_preconditioner<Temperature, HeatBalance>(inverse::Diagonal{});
    using Recipe = decltype(choice);
    using Block = pc::keys::Block<Recipe::Rows, Recipe::Values>;
    auto physical = operators::linearized_block(matrix(2, 2, {2.0, 0.0, 0.0, -4.0}));
    const auto context = dependencies::make_dependency_context(dependencies::borrow<Block>(physical));
    auto cache = dependencies::make_computation_cache<pc::Failure>(pc::make_preconditioner_graph(choice), context);
    REQUIRE(cache.try_prepare(context));
    auto frozen = cache.share<Recipe::Result>();
    CHECK(apply(*frozen, {6.0, 8.0})[0] == Catch::Approx(3.0));
    CHECK(apply(*frozen, {6.0, 8.0})[1] == Catch::Approx(-2.0));
    const auto reused = cache.try_prepare(context);
    REQUIRE(reused);
    CHECK(reused->computed_count() == 0);
    CHECK(cache.share<Recipe::Result>() == frozen);

    physical = operators::linearized_block(matrix(2, 2, {0.0, 0.0, 0.0, -4.0}));
    const auto rejected = cache.try_prepare(context);
    REQUIRE_FALSE(rejected);
    CHECK(rejected.error().error.code == Code::PRECONDITIONER_SINGULAR_OPERATOR);
    CHECK_FALSE(cache.ready());
    CHECK_THROWS(cache.share<Recipe::Result>());
    CHECK(apply(*frozen, {6.0, 8.0})[0] == Catch::Approx(3.0));

    physical = operators::linearized_block(matrix(2, 2, {3.0, 0.0, 0.0, -8.0}));
    REQUIRE(cache.try_prepare(context));
    CHECK(cache.share<Recipe::Result>() != frozen);
    CHECK(apply(cache.get<Recipe::Result>(), {6.0, 8.0})[0] == Catch::Approx(2.0));
    cache.invalidate();
    CHECK(apply(*frozen, {6.0, 8.0})[1] == Catch::Approx(-2.0));
}

TEST_CASE("Composed inverse graphs obey separate physical state and equation orders", "[unit][preconditioner][graph]")
{
    using Values = dependencies::KeyList<Density, Temperature>;
    using Rows = dependencies::KeyList<HeatBalance, Density>;
    const serif::discretization::blocks::SystemLayout<Values, Rows> layout{.state_sizes = {1, 2}, .residual_sizes = {2, 1}};
    const auto heat = pc::make_preconditioner<Temperature, HeatBalance>(inverse::Diagonal{});
    const auto density = pc::make_preconditioner<Density>(inverse::Diagonal{});
    // A zero cross block makes the triangular coupling act blockwise, isolating the coordinate orders.
    const auto choice = pc::couple(heat, density, coupling::SecondThenFirst{});
    const auto context = dependencies::make_dependency_context(
        pc::physical_block<HeatBalance, Temperature>(operators::linearized_block(matrix(2, 2, {2.0, 0.0, 0.0, 4.0}))),
        pc::physical_block<Density, Density>(operators::linearized_block(matrix(1, 1, {5.0}))),
        pc::physical_block<HeatBalance, Density>(operators::linearized_block(matrix(2, 1, {0.0, 0.0}))),
        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 result = apply(cache.get<pc::keys::Preconditioner>(), {6.0, 8.0, 20.0});
    CHECK(result[0] == Catch::Approx(4.0));
    CHECK(result[1] == Catch::Approx(3.0));
    CHECK(result[2] == Catch::Approx(2.0));
}

TEST_CASE("Constraint graph refresh follows child inverse changes without requiring a new root Jacobian", "[unit][preconditioner][graph]")
{
    using D = dependencies::KeyList<Density>;
    using R = dependencies::KeyList<MassConstraint>;
    using V = dependencies::KeyList<MassParameter>;
    using Constraints = pc::ConstraintCoordinates<R, V>;
    const auto field_choice = pc::make_preconditioner<Density>(inverse::Diagonal{});
    const auto choice = pc::with_constraints<Constraints>(field_choice);
    using Recipe = decltype(choice);
    auto fields = operators::linearized_block(matrix(1, 1, {2.0}));
    const auto context = dependencies::make_dependency_context(
        dependencies::borrow<pc::keys::Block<D, D>>(fields),
        dependencies::own<pc::keys::Block<D, V>>(operators::linearized_block(matrix(1, 1, {1.0}))),
        dependencies::own<pc::keys::Block<R, D>>(operators::linearized_block(matrix(1, 1, {3.0}))),
        dependencies::own<pc::keys::Block<R, V>>(operators::linearized_block(matrix(1, 1, {4.0}))));
    auto cache = dependencies::make_computation_cache<pc::Failure>(pc::make_preconditioner_graph(choice), context);
    REQUIRE(cache.try_prepare(context));
    auto old = cache.share<Recipe::Result>();
    const auto old_result = apply(*old, {7.0, 15.0});
    CHECK(old_result[0] == Catch::Approx(2.6));
    CHECK(old_result[1] == Catch::Approx(1.8));
    fields = operators::linearized_block(matrix(1, 1, {4.0}));
    auto refreshed = cache.try_prepare(context);
    REQUIRE(refreshed);
    CHECK(refreshed->status<Recipe::Result>() == dependencies::PreparationStatus::computed);
    const auto result = apply(cache.get<Recipe::Result>(), {7.0, 15.0});
    CHECK(4.0 * result[0] + result[1] == Catch::Approx(7.0));
    CHECK(3.0 * result[0] + 4.0 * result[1] == Catch::Approx(15.0));
    CHECK(apply(*old, {7.0, 15.0})[0] == Catch::Approx(old_result[0]));
}

TEST_CASE("An incompatible physical representation returns a preparation error", "[unit][preconditioner][graph]")
{
    const auto choice = pc::make_preconditioner<Density>(inverse::Diagonal{});
    const auto context = dependencies::make_dependency_context(pc::physical_block<Density, Density>(operators::LinearizedBlock(matrix(1, 1, {2.0}))));
    auto cache = dependencies::make_computation_cache<pc::Failure>(pc::make_preconditioner_graph(choice), context);
    const auto result = cache.try_prepare(context);
    REQUIRE_FALSE(result);
    CHECK(result.error().error.code == Code::PRECONDITIONER_INCOMPATIBLE_STRATEGY);
}

TEST_CASE("MFEM adapters bind accepted Jacobians and reuse the existing normalization transform", "[unit][preconditioner][graph]")
{
    namespace normalization = serif::solver::normalization;
    auto physical = matrix(2, 2, {2.0, 0.0, 0.0, 4.0});
    const auto inverse = pc::backend::mfem::try_dense_inverse(*physical);
    REQUIRE(inverse);
    pc::backend::mfem::Solver solver(*inverse, physical);
    mfem::Vector input(2), output(2), wrong_size(1);
    input = 2.0;
    CHECK_THROWS_WITH(solver.Mult(wrong_size, output), Catch::Matchers::ContainsSubstring("PRECONDITIONER_SIZE_MISMATCH"));
    CHECK_THROWS_WITH(solver.Mult(input, wrong_size), Catch::Matchers::ContainsSubstring("PRECONDITIONER_SIZE_MISMATCH"));
    solver.Mult(input, output);
    CHECK(output[0] == Catch::Approx(1.0));
    CHECK(output[1] == Catch::Approx(0.5));
    CHECK_NOTHROW(solver.SetOperator(*physical));
    const auto other = matrix(2, 2, {2.0, 0.0, 0.0, 4.0});
    CHECK_THROWS(solver.SetOperator(*other));
    const auto scaling = normalization::make_diagonal_normalization(std::vector<double>{3.0, 5.0}, std::vector<double>{7.0, 11.0});
    REQUIRE(scaling);
    normalization::backend::mfem::ScaledJacobianOperator jacobian(*physical, *scaling);
    normalization::backend::mfem::ScaledPreconditioner scaled(solver, jacobian);
    mfem::Vector state(2), residual(2), recovered(2);
    state[0] = 2.0;
    state[1] = -3.0;
    jacobian.Mult(state, residual);
    scaled.Mult(residual, recovered);
    CHECK(recovered[0] == Catch::Approx(state[0]));
    CHECK(recovered[1] == Catch::Approx(state[1]));
}