#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <vector>

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

#include "mfem/linalg/densemat.hpp"
#include "serif/backend/mfem/model/normalized_system.hpp"
#include "serif/utils/dependencies/computation_cache.hpp"
#include "serif/backend/mfem/model/normalization/computations.hpp"

namespace {
    namespace model = serif::model;
    namespace backend = model::backend::mfem;
    namespace dependencies = serif::utils::dependencies;
    namespace normalization = serif::solver::normalization;
    namespace dimensions = serif::dimensions;
    struct PhysicalResidual {};
    struct ResidualFactor {};

    /** @brief Give the frozen normalization adapter an independently owned physical action. */
    [[nodiscard]] std::shared_ptr<const backend::PhysicalSystem> physical_system(double value) {
        auto matrix = std::make_shared<mfem::DenseMatrix>(1);
        (*matrix)(0, 0) = 3.0 + value;
        const auto block = serif::operators::backend::mfem::linearized_block(std::shared_ptr<const mfem::DenseMatrix>(matrix));
        return backend::make_physical_system(1, 1, {{.row_offset = 0, .values = {value}}}, {{.row_offset = 0, .column_offset = 0, .block = block}});
    }

    /** @brief Keep a valid reference fixed while failure is induced only in coordinate conversion. */
    [[nodiscard]] normalization::StellarReferenceScales reference_scales() {
        const auto scales = normalization::prepare_stellar_reference_scales(dimensions::MassValue{2.0}, dimensions::LengthValue{3.0});
        REQUIRE(scales);
        return *scales;
    }
}

TEST_CASE("A normalization overflow rejects the new physical preparation without losing frozen solver coordinates", "[model][normalization][rejection]") {
    using namespace dependencies;

    double value = 1.0;
    double factor = 2.0;
    int physical_preparations = 0;
    const auto scales = reference_scales();
    const auto inputs = make_dependency_context(borrow<PhysicalResidual>(value), borrow<ResidualFactor>(factor), borrow<normalization::keys::Scales>(scales));
    const auto physical = make_computation<model::keys::PhysicalSystem, InputList<PhysicalResidual>>([&](double residual) {
        ++physical_preparations;
        return physical_system(residual);
    });
    const auto coordinates = make_computation<normalization::keys::Normalization, InputList<ResidualFactor>>([](double residual_factor) {
        return normalization::make_diagonal_normalization({3.0}, {residual_factor});
    });
    const auto normalized = make_computation<model::normalization::PreparedSystem,
        InputList<model::keys::PhysicalSystem, normalization::keys::Normalization, normalization::keys::Scales>>(backend::make_normalized_system);
    auto graph = make_computation_graph(physical, coordinates, normalized);
    auto cache = make_computation_cache(std::move(graph), inputs);
    using Requested = KeyList<model::normalization::PreparedSystem>;
    REQUIRE(cache.try_prepare(inputs, Requested{}));
    const auto physical_snapshot = cache.get<model::keys::PhysicalSystem>();
    const auto frozen = cache.template get<model::normalization::PreparedSystem>();
    REQUIRE(frozen->physical() == physical_snapshot);
    CHECK(frozen->residual() == std::vector<double>{2.0});
    mfem::Vector direction(1), original_action, repeated_action;
    direction[0] = 1.0;
    frozen->jacobian().Mult(direction, original_action);
    CHECK(original_action[0] == Catch::Approx(8.0 / 3.0));
    const auto preparations = cache.preparation_count();

    value = 2.0;
    factor = std::numeric_limits<double>::max();
    const auto rejected = cache.try_prepare(inputs, Requested{});
    REQUIRE_FALSE(rejected);
    CHECK_THAT(rejected.error().error, Catch::Matchers::ContainsSubstring("normalized stellar residual is not finite"));
    CHECK(physical_preparations == 2);
    CHECK(frozen->physical() == physical_snapshot);
    CHECK_FALSE(cache.ready());
    CHECK_FALSE(cache.template contains<model::keys::PhysicalSystem>());
    CHECK_FALSE(cache.template contains<model::normalization::PreparedSystem>());
    CHECK(cache.report().computed_count() == 0);
    CHECK(cache.report().reused_count() == 0);
    CHECK(cache.preparation_count() == preparations);
    CHECK(frozen->residual() == std::vector<double>{2.0});
    frozen->jacobian().Mult(direction, repeated_action);
    CHECK(repeated_action[0] == original_action[0]);

    value = 1.0;
    factor = 2.0;
    REQUIRE(cache.try_prepare(inputs, Requested{}));
    CHECK(cache.get<model::keys::PhysicalSystem>() == physical_snapshot);
    CHECK(cache.template get<model::normalization::PreparedSystem>() == frozen);
    CHECK(cache.report().computed_count() == 0);
    value = 2.0;
    REQUIRE(cache.try_prepare(inputs, Requested{}));
    CHECK(cache.get<model::keys::PhysicalSystem>() != physical_snapshot);
    CHECK(cache.template get<model::normalization::PreparedSystem>()->residual() == std::vector<double>{4.0});
    CHECK(frozen->residual() == std::vector<double>{2.0});
}

TEST_CASE("Normalized system construction diagnoses missing owners and incompatible coordinates", "[model][normalization][validation]") {
    const auto factors = normalization::make_diagonal_normalization({1.0}, {2.0});
    REQUIRE(factors);
    CHECK_THROWS_WITH(backend::make_normalized_system({}, *factors, reference_scales()), Catch::Matchers::ContainsSubstring("physical system"));
    const auto wrong_state = normalization::make_diagonal_normalization({1.0, 1.0}, {2.0});
    const auto wrong_residual = normalization::make_diagonal_normalization({1.0}, {2.0, 2.0});
    REQUIRE(wrong_state);
    REQUIRE(wrong_residual);
    CHECK_THROWS_AS(backend::make_normalized_system(physical_system(1.0), *wrong_state, reference_scales()), std::invalid_argument);
    CHECK_THROWS_AS(backend::make_normalized_system(physical_system(1.0), *wrong_residual, reference_scales()), std::invalid_argument);
    const auto excessive = normalization::make_diagonal_normalization({1.0}, {std::numeric_limits<double>::max()});
    REQUIRE(excessive);
    CHECK_FALSE(backend::make_normalized_system(physical_system(2.0), *excessive, reference_scales()));
}