#include <algorithm>
#include <limits>
#include <memory>
#include <type_traits>
#include <variant>

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

#include "mfem/linalg/densemat.hpp"
#include "mfem/general/communication.hpp"
#include "serif/tests/equations/test_equations.hpp"
#include "serif/backend/mfem/model/equilibrium.hpp"
#include "serif/backend/mfem/solver/seed/state.hpp"
#include "serif/solver/seed/polytrope.hpp"
#include "serif/tests/solver/stellar_fixture.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"
#include "serif/utils/dependencies/timing.hpp"
#include "serif/backend/mfem/operators/actions.hpp"
#include "serif/backend/mfem/model/preconditioned_system.hpp"
#include "serif/backend/mfem/model/computations/physical_blocks.hpp"
#include "serif/backend/mfem/model/computations/solver_layout.hpp"
#include "serif/solver/normalization/normalization.hpp"
#include "serif/solver/preconditioner/preconditioner.hpp"

namespace {
    namespace model = serif::model;
    namespace backend = model::backend::mfem;
    namespace dependencies = serif::utils::dependencies;
    namespace blocks = serif::discretization::blocks;
    namespace normalization = serif::solver::normalization;
    namespace preconditioner = serif::solver::preconditioner;
    using Density = blocks::density::mass;
    using Enthalpy = ::serif::tests::equations::Enthalpy;
    using Layout = blocks::SystemLayout<dependencies::KeyList<Enthalpy, Density>, dependencies::KeyList<Density, Enthalpy>>;
    using DensityBlock = serif::operators::PhysicalBlock<dependencies::KeyList<Density>, dependencies::KeyList<Density>>;
    using ReorderedBlock = serif::operators::PhysicalBlock<dependencies::KeyList<Enthalpy, Density>, dependencies::KeyList<Density, Enthalpy>>;
    using Rules = dependencies::Bindings<dependencies::Binding<Density, normalization::IdentityCoordinate>, dependencies::Binding<Enthalpy, normalization::IdentityCoordinate>>;
    using NormalizationPlan = normalization::NormalizationPlan<Rules, Rules>;
    struct DensityResponse {
    };
    struct EvaluationService {};
    struct LinearService {};

    /** @brief Supply a nonsymmetric physical Jacobian with an independently known two-coordinate inverse. */
    [[nodiscard]] std::shared_ptr<const backend::PhysicalSystem> small_physical_system() {
        auto matrix = std::make_shared<::mfem::DenseMatrix>(2);
        (*matrix)(0, 0) = 4.0;
        (*matrix)(0, 1) = 1.0;
        (*matrix)(1, 0) = 2.0;
        (*matrix)(1, 1) = 3.0;
        const auto block = serif::operators::backend::mfem::linearized_block(std::shared_ptr<const ::mfem::DenseMatrix>(matrix));
        return backend::make_physical_system(2, 2, {{.row_offset = 0, .values = {9.0, 7.0}}},
            {{.row_offset = 0, .column_offset = 0, .block = block}});
    }

    /** @brief Keep inverse ownership independent from the physical system so both lifetimes are tested. */
    [[nodiscard]] std::shared_ptr<const ::mfem::Operator> small_physical_inverse() {
        auto matrix = std::make_shared<::mfem::DenseMatrix>(2);
        (*matrix)(0, 0) = 0.3;
        (*matrix)(0, 1) = -0.1;
        (*matrix)(1, 0) = -0.2;
        (*matrix)(1, 1) = 0.4;
        return matrix;
    }
}

TEST_CASE("Prepared stellar linear systems retain matching physical and normalized inverse owners", "[unit][model][preconditioner][normalization][solver_services]") {
    const auto check_coordinates = []<bool Normalize>() {
        using System = std::conditional_t<Normalize, backend::NormalizedSystem, backend::PhysicalSystem>;
        using Prepared = backend::PreconditionedSystem<System>;
        std::weak_ptr<const backend::PhysicalSystem> physical_owner;
        std::weak_ptr<const ::mfem::Operator> inverse_owner;
        std::shared_ptr<const Prepared> frozen;
        {
            const auto physical = small_physical_system();
            const auto inverse = small_physical_inverse();
            physical_owner = physical;
            inverse_owner = inverse;
            const auto system = [&]() {
                if constexpr (Normalize) {
                    const auto coordinates = normalization::make_diagonal_normalization({2.0, 5.0}, {3.0, 7.0});
                    const auto scales = normalization::prepare_stellar_reference_scales(
                        serif::dimensions::MassValue{2.0}, serif::dimensions::LengthValue{3.0});
                    REQUIRE(coordinates);
                    REQUIRE(scales);
                    const auto normalized = backend::make_normalized_system(physical, *coordinates, *scales);
                    REQUIRE(normalized);
                    return *normalized;
                } else {
                    return physical;
                }
            }();
            frozen = std::make_shared<const Prepared>(system, preconditioner::backend::mfem::Inverse{.action = inverse});
        }
        REQUIRE_FALSE(physical_owner.expired());
        REQUIRE_FALSE(inverse_owner.expired());
        CHECK(frozen->contract() == preconditioner::ApplicationContract::fixed_linear);
        const ::mfem::Operator& jacobian = [&]() -> const ::mfem::Operator& {
            if constexpr (Normalize) {
                return frozen->system()->jacobian();
            } else {
                return *frozen->system()->jacobian().action();
            }
        }();
        CHECK_NOTHROW(frozen->preconditioner().SetOperator(jacobian));
        ::mfem::Vector right_hand_side(2), correction(2), reconstructed(2);
        right_hand_side[0] = Normalize ? 27.0 : 9.0;
        right_hand_side[1] = Normalize ? 49.0 : 7.0;
        correction = 123.0;
        frozen->preconditioner().Mult(right_hand_side, correction);
        CHECK(correction[0] == Catch::Approx(Normalize ? 4.0 : 2.0));
        CHECK(correction[1] == Catch::Approx(Normalize ? 5.0 : 1.0));
        jacobian.Mult(correction, reconstructed);
        reconstructed -= right_hand_side;
        CHECK(reconstructed.Norml2() < 1.0e-13);

        // Equal dimensions and identical entries do not establish a matching preparation.
        const auto unrelated = small_physical_system();
        CHECK_THROWS(frozen->preconditioner().SetOperator(*unrelated->jacobian().action()));
        if constexpr (Normalize) {
            const auto replacement = backend::make_normalized_system(unrelated, frozen->system()->coordinates(), frozen->system()->reference_scales());
            REQUIRE(replacement);
            CHECK_THROWS(frozen->preconditioner().SetOperator((*replacement)->jacobian()));
        }
        CHECK_NOTHROW(frozen->preconditioner().SetOperator(jacobian));
        correction = -91.0;
        frozen->preconditioner().Mult(right_hand_side, correction);
        CHECK(correction[0] == Catch::Approx(Normalize ? 4.0 : 2.0));
        CHECK(correction[1] == Catch::Approx(Normalize ? 5.0 : 1.0));

        right_hand_side[0] = 9.0;
        right_hand_side[1] = 7.0;
        frozen->physical_preconditioner().Mult(right_hand_side, correction);
        CHECK(correction[0] == Catch::Approx(2.0));
        CHECK(correction[1] == Catch::Approx(1.0));
        frozen.reset();
        CHECK(physical_owner.expired());
        CHECK(inverse_owner.expired());
    };
    SECTION("Physical coordinates") {
        check_coordinates.template operator()<false>();
    }
    SECTION("Nonunit normalized coordinates") {
        check_coordinates.template operator()<true>();
    }
}

TEST_CASE("Prepared stellar systems preserve application contracts and reject incompatible owners", "[unit][model][preconditioner][solver_services]") {
    using Prepared = backend::PreconditionedSystem<backend::PhysicalSystem>;
    const preconditioner::backend::mfem::Inverse inverse{.action = small_physical_inverse(), .contract = preconditioner::ApplicationContract::variable};
    const auto physical = small_physical_system();
    const Prepared prepared(physical, inverse);
    CHECK(prepared.contract() == preconditioner::ApplicationContract::variable);
    CHECK_THROWS_WITH(Prepared({}, inverse), Catch::Matchers::ContainsSubstring("evaluated star"));
    CHECK_THROWS_WITH(Prepared(physical, {}), Catch::Matchers::ContainsSubstring("PRECONDITIONER_INVALID_OPERATOR"));
    const preconditioner::backend::mfem::Inverse wrong_size{.action = std::make_shared<const ::mfem::DenseMatrix>(1)};
    CHECK_THROWS_WITH(Prepared(physical, wrong_size), Catch::Matchers::ContainsSubstring("PRECONDITIONER_SIZE_MISMATCH"));
}

TEST_CASE("Model normalization and inverse computations share one physical evaluation", "[model][integration][preconditioner][normalization]") {
    using namespace dependencies;

    const Layout layout{.state_sizes = {2, 1}, .residual_sizes = {1, 2}};
    double response = 4.0;
    int physical_preparations = 0;
    int physical_actions = 0;
    const auto inputs = make_dependency_context(borrow<DensityResponse>(response));
    const auto physical = make_computation<model::keys::PhysicalSystem, InputList<DensityResponse>>(
        [&](double density_response) -> std::expected<std::shared_ptr<const model::backend::mfem::PhysicalSystem>, std::string> {
            ++physical_preparations;
            if (density_response < 0.0) {
                return std::unexpected("Synthetic physical rejection for transaction testing.");
            }
            auto matrix = std::make_shared<::mfem::DenseMatrix>(3);
            for (int row = 0; row < 3; ++row) {
                for (int column = 0; column < 3; ++column) {
                    (*matrix)(row, column) = 2.0 + row * 3.0 + column;
                }
            }
            (*matrix)(0, 2) = density_response;
            const auto original = serif::operators::backend::mfem::linearized_block(std::shared_ptr<const ::mfem::DenseMatrix>(matrix));
            const auto action = serif::operators::backend::mfem::linearized_action(3, 3, [original, &physical_actions](const auto& input, auto& output) {
                ++physical_actions;
                original.action()->Mult(input, output);
            }, [original, &physical_actions](const auto& input, auto& output) {
                ++physical_actions;
                original.action()->MultTranspose(input, output);
            });
            const serif::operators::backend::mfem::LinearizedBlock block(action.action(), [original]() {
                return original.diagonal();
            }, [original]() {
                return original.assemble();
            });
            return model::backend::mfem::make_physical_system(3, 3, {}, {{0, 0, block}});
        });
    const auto inverse = preconditioner::make_preconditioner<Density>(preconditioner::strategies::inverse::Diagonal{});
    using InverseKey = decltype(inverse)::Result;
    auto graph = std::apply([](auto... nodes) {
        return make_computation_graph(std::move(nodes)...);
    }, std::tuple_cat(std::tuple{physical}, model::computations::make_solver_layout_computations(layout),
        model::computations::make_physical_block_computations(layout, KeyList<DensityBlock, ReorderedBlock>{}),
        normalization::make_normalization_computations<NormalizationPlan>(layout),
        preconditioner::make_preconditioner_computations(inverse)));
    auto cache = make_computation_cache<std::variant<std::string, preconditioner::Failure>>(std::move(graph), inputs);
    using Requested = KeyList<model::keys::PhysicalSystem, InverseKey, ReorderedBlock, normalization::keys::Normalization>;
    REQUIRE(cache.try_prepare(inputs, Requested{}));
    CHECK(physical_preparations == 1);
    const auto frozen = cache.get<model::keys::PhysicalSystem>();
    const auto frozen_inverse = cache.template share<InverseKey>();
    const auto& factors = cache.template get<normalization::keys::Normalization>();
    REQUIRE(factors.state_factors().size() == 3);
    REQUIRE(factors.residual_factors().size() == 3);
    for (double factor : factors.state_factors()) {
        CHECK(factor == 1.0);
    }
    ::mfem::Vector source(1), result(1);
    source = 8.0;
    frozen_inverse->action->Mult(source, result);
    CHECK(result[0] == Catch::Approx(2.0));

    const auto reordered = cache.template get<ReorderedBlock>();
    ::mfem::Vector direction(3), dual(3), reordered_result, transpose;
    direction[0] = 2.0;
    direction[1] = 3.0;
    direction[2] = 5.0;
    dual[0] = -0.4;
    dual[1] = 0.7;
    dual[2] = 1.2;
    reordered.action()->Mult(direction, reordered_result);
    CHECK(reordered_result[0] == Catch::Approx(7.0 * 2.0 + 5.0 * 3.0 + 6.0 * 5.0));
    CHECK(reordered_result[1] == Catch::Approx(10.0 * 2.0 + 8.0 * 3.0 + 9.0 * 5.0));
    CHECK(reordered_result[2] == Catch::Approx(4.0 * 2.0 + 2.0 * 3.0 + 3.0 * 5.0));
    reordered.action()->MultTranspose(dual, transpose);
    CHECK(dual * reordered_result == Catch::Approx(direction * transpose));
    CHECK(physical_actions == 2);
    CHECK(reordered.diagonal() == std::vector<double>{7, 8, 3});
    const auto reordered_matrix = reordered.assemble();
    ::mfem::Vector assembled_result(3);
    reordered_matrix->Mult(direction, assembled_result);
    assembled_result -= reordered_result;
    CHECK(assembled_result.Norml2() == 0.0);
    CHECK(physical_actions == 2);

    REQUIRE(cache.try_prepare(inputs, Requested{}));
    CHECK(physical_preparations == 1);
    response = 8.0;
    REQUIRE(cache.try_prepare(inputs, Requested{}));
    CHECK(physical_preparations == 2);
    cache.template get<InverseKey>().action->Mult(source, result);
    CHECK(result[0] == Catch::Approx(1.0));
    frozen_inverse->action->Mult(source, result);
    CHECK(result[0] == Catch::Approx(2.0));
    const auto last = cache.get<model::keys::PhysicalSystem>();

    // A completed physical candidate must not be published if a requested inverse rejects it.
    response = 0.0;
    const auto preparations = cache.preparation_count();
    const auto singular = cache.try_prepare(inputs, Requested{});
    REQUIRE_FALSE(singular);
    CHECK(std::holds_alternative<preconditioner::Failure>(singular.error().error));
    last->block(0, 1, 2, 1).action()->Mult(source, result);
    CHECK(result[0] == Catch::Approx(64.0));
    CHECK_FALSE(cache.ready());
    CHECK_FALSE(cache.template contains<model::keys::PhysicalSystem>());
    CHECK_THROWS_AS(cache.template get<InverseKey>(), std::logic_error);
    CHECK(cache.report().computed_count() == 0);
    CHECK(cache.report().reused_count() == 0);
    CHECK(cache.preparation_count() == preparations);
    response = 8.0;
    REQUIRE(cache.try_prepare(inputs, Requested{}));
    CHECK(cache.get<model::keys::PhysicalSystem>() == last);
    CHECK(cache.report().computed_count() == 0);

    response = -1.0;
    CHECK_FALSE(cache.try_prepare(inputs, Requested{}));
    CHECK_FALSE(cache.contains<model::keys::PhysicalSystem>());
    last->block(0, 1, 2, 1).action()->Mult(source, result);
    CHECK(result[0] == Catch::Approx(64.0));
    frozen->block(0, 1, 2, 1).action()->Mult(source, result);
    CHECK(result[0] == Catch::Approx(32.0));
}

TEST_CASE("Physical solver blocks reject unrepresentable or unrelated coordinate layouts", "[unit][model][solver_services]") {
    using namespace dependencies;

    const std::size_t too_large = static_cast<std::size_t>(std::numeric_limits<int>::max()) + 1;
    CHECK_THROWS_WITH(model::computations::physical_block<DensityBlock>(Layout{.state_sizes = {too_large, 1}, .residual_sizes = {1, 2}}),
                      Catch::Matchers::ContainsSubstring("INVALID_DOF_COUNT"));
    CHECK_THROWS_WITH(model::computations::physical_block<DensityBlock>(Layout{.state_sizes = {2, 1}, .residual_sizes = {too_large, 2}}),
                      Catch::Matchers::ContainsSubstring("INVALID_DOF_COUNT"));

    const Layout layout{.state_sizes = {2, 1}, .residual_sizes = {1, 2}};
    std::shared_ptr<const model::backend::mfem::PhysicalSystem> system;
    const auto inputs = make_dependency_context(borrow<model::keys::PhysicalSystem>(system));
    auto cache = make_computation_cache(inputs, model::computations::physical_block<DensityBlock>(layout));
    CHECK_THROWS_WITH(cache.try_prepare(inputs), Catch::Matchers::ContainsSubstring("INVALID_ARGUMENT"));
    for (const auto dimensions : {std::pair{2, 3}, std::pair{3, 2}}) {
        system = model::backend::mfem::make_physical_system(dimensions.first, dimensions.second, {}, {});
        CHECK_THROWS_WITH(cache.try_prepare(inputs), Catch::Matchers::ContainsSubstring("MATRIX_SHAPE_ERROR"));
    }
    system = model::backend::mfem::make_physical_system(3, 3, {}, {});
    REQUIRE(cache.try_prepare(inputs));
    CHECK(cache.get<DensityBlock>().action()->Height() == 1);
    CHECK(cache.get<DensityBlock>().action()->Width() == 1);
}

TEST_CASE("Configured model services share cache transactions and defer linear-only work", "[model][integration][solver_services]") {
    using namespace dependencies;
    namespace fixture = serif::tests::solver;
    using Constant = blocks::barotropic_constant::mass_normalization;
    if (!::mfem::Mpi::IsInitialized()) {
        ::mfem::Mpi::Init();
    }
    const auto physics = fixture::make_model();
    const auto seed = serif::solver::seed::make_polytropic_seed(physics);
    REQUIRE(seed);
    const auto resources = fixture::make_resources(seed->scales.radius.value());
    const auto finite_elements = fixture::make_discretization(resources, 0);
    const auto state = serif::solver::seed::backend::mfem::make_polytropic_state(physics, finite_elements, *seed);
    REQUIRE(state);
    REQUIRE(state->template get<Constant>()[0] < 0.5);
    const auto configuration = ::serif::model::make_discretization(::serif::model::DiscretizationOptions{
        .model = physics, .finite_elements = finite_elements,
        .normalization = normalization::NormalizationOptions{.scheme = *normalization::prepare_stellar_reference_scales(seed->scales.mass, seed->scales.radius)},
        .preconditioner = preconditioner::make_primal_preconditioner(state->layout())});
    int evaluations = 0;
    int linear_preparations = 0;
    const auto services = [&](const auto& bound, const auto& layout, auto, const auto& existing) {
        REQUIRE(bound->model);
        CHECK(layout == state->layout());
        using Existing = decltype(backend::detail::computation_keys(existing));
        static_assert(ContainsKey<model::preconditioning::PreparedSystem, Existing>);
        // The surface potential level is a tracked state input, so changing it invalidates both extensions.
        const auto evaluation = make_computation<EvaluationService, InputList<Constant>>([&](const auto& constant) -> std::expected<double, std::string> {
            ++evaluations;
            if (constant.values[0] > 10.0) {
                return std::unexpected("Synthetic evaluation-service rejection.");
            }
            return constant.values[0];
        });
        const auto linear = make_computation<LinearService, InputList<EvaluationService, model::preconditioning::PreparedSystem>>(
            [&](double constant, const auto& prepared) -> std::expected<std::size_t, std::string> {
                ++linear_preparations;
                if (constant > 0.5) {
                    return std::unexpected("Synthetic linear-service rejection.");
                }
                return prepared->system()->residual().size();
            });
        return std::pair{std::tuple{evaluation}, std::tuple{linear}};
    };
    GraphTiming timing;
    auto evaluator = backend::bind_stellar_model(configuration, *state, {}, make_dependency_context(), services, time_graph_nodes(timing));
    REQUIRE(timing.runs().size() == 1);
    const auto recorded = [&]<typename Key>() -> const NodeTiming& {
        const auto& nodes = timing.runs().front()->nodes;
        const auto node = std::ranges::find(nodes, std::string{nameof::nameof_type<Key>()}, &NodeTiming::key_name);
        REQUIRE(node != nodes.end());
        return *node;
    };
    const auto evaluated = evaluator.evaluate(*state);
    REQUIRE(evaluated);
    CHECK(evaluations == 1);
    CHECK(linear_preparations == 0);
    CHECK(recorded.template operator()<EvaluationService>().completed == 1);
    CHECK(recorded.template operator()<LinearService>().calls == 0);
    CHECK(evaluator.cache().report().template status<LinearService>() == PreparationStatus::not_required);
    REQUIRE(evaluator.evaluate(*state));
    CHECK(evaluations == 1);
    const auto prepared = evaluator.prepare_linear_system(*state);
    REQUIRE(prepared);
    CHECK(evaluations == 1);
    CHECK(linear_preparations == 1);
    CHECK((*prepared)->system() == *evaluated);
    CHECK(evaluator.cache().template get<LinearService>() == (*evaluated)->residual().size());
    CHECK(recorded.template operator()<EvaluationService>().calls == 1);
    CHECK(recorded.template operator()<LinearService>().completed == 1);

    auto changed = *state;
    changed.template get<Constant>()[0] = 1.0;
    REQUIRE(evaluator.evaluate(changed));
    CHECK(evaluations == 2);
    CHECK(linear_preparations == 1);
    const auto rejected_linear = evaluator.prepare_linear_system(changed);
    REQUIRE_FALSE(rejected_linear);
    CHECK(std::get<std::string>(rejected_linear.error().error) == "Synthetic linear-service rejection.");
    CHECK(recorded.template operator()<LinearService>().rejected == 1);
    CHECK(recorded.template operator()<LinearService>().calls == 2);
    CHECK(linear_preparations == 2);
    CHECK_FALSE(evaluator.cache().ready());
    CHECK((*prepared)->system() == *evaluated);
    CHECK_FALSE((*prepared)->system()->residual().empty());

    changed.template get<Constant>()[0] = 11.0;
    const auto rejected_evaluation = evaluator.evaluate(changed);
    REQUIRE_FALSE(rejected_evaluation);
    CHECK(std::get<std::string>(rejected_evaluation.error().error) == "Synthetic evaluation-service rejection.");
    CHECK(recorded.template operator()<EvaluationService>().rejected == 1);
    CHECK(linear_preparations == 2);
    CHECK_FALSE(evaluator.cache().ready());
    REQUIRE(evaluator.evaluate(*state));
    CHECK(evaluator.cache().template get<EvaluationService>() == state->template get<Constant>()[0]);
}