#include <algorithm>
#include <array>
#include <memory>
#include <string_view>

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

#include "serif/solver/continuation/continuation.hpp"
#include "serif/solver/continuation/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/timing.hpp"

namespace {
    namespace family_api = serif::solver::continuation;
    namespace solver = serif::solver;
    namespace model = serif::model;
    namespace backend = model::backend::mfem;
    namespace fixture = serif::tests::solver;
    namespace dependencies = serif::utils::dependencies;
    namespace dimensions = serif::dimensions;
    namespace normalization = solver::normalization;
    using Density = fixture::Density;

    /** @brief A known density field isolates public continuation orchestration from stellar conditioning. */
    struct DensityEquation : serif::equations::Contribution<serif::discretization::forms::Operand<Density>> {
        using Value = serif::discretization::forms::Operand<Density>;
        using Operands = serif::utils::types::TypeList<Value>;
        using ResidualQuantity = dimensions::MassDensity;
        static constexpr std::string_view name = "continuation density regression";
        double target;

        /** @brief Bind a physical density target read from the fresh stage's EOS. */
        explicit DensityEquation(double density) : target(density) {}
        /** @brief Integrate the difference from the exact representable constant field. */
        [[nodiscard]] serif::equations::PointLinearization<Operands> evaluate(const serif::equations::EquationPoint<Operands>& point) const {
            serif::equations::PointLinearization<Operands> result;
            result.residual = point.fields.template get<Density>() - target;
            result.field_derivatives.template get<Density>() = 1.0;
            return result;
        }
        /** @brief A constant target adds no polynomial order. */
        [[nodiscard]] int additional_order(Value, int) const { return 0; }
    };

    /** @brief Rebuild an explicitly simple equation from the same EOS parameter family as a full star. */
    [[nodiscard]] auto density_model() {
        return model::make_stellar_model(serif::eos::models::Polytrope{1.5, 1.0}, [](const auto& physics) {
            return model::EquilibriumEquations{.interior_equations = serif::equations::EquationRegistry{
                DensityEquation{physics.equation_of_state().polytropic_index()}}};
        });
    }

    /** @brief Bind only density and prescribed displacement, retaining the real stellar reference mesh. */
    [[nodiscard]] auto density_discretization(const std::shared_ptr<fixture::StellarResources>& resources) {
        return backend::make_stellar_discretization(resources, dependencies::make_dependency_context(
            dependencies::borrow<Density>(resources->density), dependencies::borrow<fixture::Displacement>(resources->displacement)),
            *resources->mesh.exterior_coordinate->values, resources->mapper);
    }

    /** @brief Change numerical normalization deliberately while preserving physical coefficient meanings. */
    [[nodiscard]] auto numerical_coordinates(double radius) {
        return normalization::NormalizationOptions{.scheme = *normalization::prepare_stellar_reference_scales(
            dimensions::MassValue{1.0}, dimensions::LengthValue{radius}, dimensions::QuantityValue<dimensions::GravitationalConstant>{1.0})};
    }

}

TEST_CASE("Public continuation solves fresh models and transfers physical finite-element coefficients", "[integration][solver][continuation]") {
    const auto resources = fixture::make_resources();
    const auto finite_elements = density_discretization(resources);
    const auto family = family_api::in(family_api::PolytropicIndex{}).with(
        family_api::preserve_pressure_at_reference_density(dimensions::DensityValue{1.0})).bind(density_model());
    using Description = decltype(family.target_model().equations()(family.target_model()));
    using Plan = model::EquilibriumPlan<serif::eos::models::Polytrope, Description>;
    ::mfem::Vector displacement(resources->displacement.size());
    displacement = 0.0;
    auto state = model::StellarState(backend::equilibrium_layout<Plan>(finite_elements, displacement));
    std::ranges::fill(state.values(), 0.8);
    int configurations = 0;
    const auto configure = [&](auto physical) {
        ++configurations;
        const double n = physical.equation_of_state().polytropic_index();
        return ::serif::model::make_discretization(::serif::model::DiscretizationOptions{
            .model = std::move(physical), .finite_elements = finite_elements, .normalization = numerical_coordinates(n),
            .preconditioner = solver::preconditioner::make_preconditioner<Density>(solver::preconditioner::strategies::inverse::Diagonal{})});
    };
    auto options = family_api::make_options(family);
    options.steps.initial_step_fraction = {1.0};
    const auto result = solver::continue_equilibrium(family, configure, state, displacement, options);
    REQUIRE(result.converged());
    REQUIRE(result.last_converged() != nullptr);
    CHECK(configurations == 2);
    CHECK(result.last_converged()->model().equation_of_state().polytropic_index() == 1.5);
    CHECK(result.last_converged()->model().equation_of_state().polytropic_constant() == 1.0);
    CHECK(*result.last_converged()->diagnostics().convergence_threshold == 1.0e-10);
    CHECK(*result.last_converged()->diagnostics().final_residual_norm <= 1.0e-10);
    for (const double density : result.last_converged()->state().values()) {
        CHECK(density == Catch::Approx(1.5).margin(1.0e-8));
    }
    CHECK(std::ranges::all_of(state.values(), [](double value) { return value == 0.8; }));

    SECTION("a same-size different mesh cannot silently reinterpret physical coefficients") {
        const auto other_resources = fixture::make_resources();
        int calls = 0;
        const auto changed_resources = [&](auto physical) {
            auto configuration = configure(std::move(physical));
            return ::serif::model::DiscretizationOptions{.model = std::move(configuration.model),
                .finite_elements = ++calls > 1 ? density_discretization(other_resources) : finite_elements,
                .normalization = std::move(configuration.normalization), .preconditioner = std::move(configuration.preconditioner)};
        };
        CHECK_THROWS_AS(solver::continue_equilibrium(family, changed_resources, state, displacement, options), std::invalid_argument);
        CHECK(calls == 2);
    }
    SECTION("the configuration factory cannot change an exact physical target") {
        const auto wrong_model = [&](auto physical) {
            return configure(physical.transform_specifications([](const auto&) { return serif::eos::models::Polytrope{1.7, 1.0}; }));
        };
        CHECK_THROWS_AS(solver::continue_equilibrium(family, wrong_model, state, displacement, options), std::invalid_argument);
    }
    SECTION("configuration callback failures do not become continuation retries") {
        int calls = 0;
        const auto rejected = [&](auto physical) {
            ++calls;
            serif::utils::errors::serif_ensure<serif::utils::errors::SERiFErrorCode::INVALID_ARGUMENT, std::invalid_argument>(false, "Deliberately invalid configuration.");
            return configure(std::move(physical));
        };
        CHECK_THROWS_AS(solver::continue_equilibrium(family, rejected, state, displacement, options), std::invalid_argument);
        CHECK(calls == 1);
    }
    SECTION("fixed-geometry solves accept reused and consumed configurations") {
        auto configuration = configure(family.model_at(family.target_coordinates()));
        auto options = solver::continuation::default_solver_options();
        CHECK(solver::solve_equilibrium(configuration, result.last_converged()->state(), displacement, options).converged());
        CHECK(solver::solve_equilibrium(std::move(configuration), result.last_converged()->state(), displacement, std::move(options)).converged());
    }
}

TEST_CASE("Public continuation qualifies the complete moving stellar assembly before declaring an anchor", "[integration][solver][continuation]") {
    const auto original = fixture::make_model();
    const auto target = model::make_stellar_model(original.equation_of_state(), family_api::identified_equations(original.equations(),
        dimensions::QuantityValue<dimensions::GravitationalConstant>{1.0}),
        original.specification<model::FixedTotalMass>(), original.specification<model::FixedAngularMomentum>());
    const auto family = family_api::in(family_api::AngularMomentum{}).bind(target);
    const auto radial = solver::seed::make_polytropic_seed(target);
    REQUIRE(radial);
    const auto resources = fixture::make_resources(radial->scales.radius.value());
    const auto finite_elements = fixture::make_discretization(resources, 0);
    const auto state = solver::seed::backend::mfem::make_polytropic_state(target, finite_elements, *radial);
    REQUIRE(state);
    const auto configure = [&](auto physical) {
        return ::serif::model::make_discretization(::serif::model::DiscretizationOptions{.model = std::move(physical),
            .finite_elements = finite_elements, .normalization = numerical_coordinates(radial->scales.radius.value())});
    };
    auto solver_options = solver::continuation::default_solver_options();
    solver_options.method.maximum_iterations = 0;
    const auto result = solver::continue_equilibrium(family, configure, *state, family_api::make_options(family), solver_options);
    CHECK_FALSE(result.converged());
    CHECK(result.last_converged() == nullptr);
    REQUIRE(result.failed.get() != nullptr);
    CHECK(result.failed.get()->has_state());
    CHECK(result.failed.get()->status() == solver::SolveStatus::maximum_iterations);
    CHECK(*result.failed.get()->diagnostics().convergence_threshold == 1.0e-10);
    auto configuration = configure(family.model_at(family.target_coordinates()));
    CHECK(solver::solve_equilibrium(configuration, *state, solver_options).status() == solver::SolveStatus::maximum_iterations);
    CHECK(solver::solve_equilibrium(std::move(configuration), *state, std::move(solver_options)).status() == solver::SolveStatus::maximum_iterations);
}

TEST_CASE("Graph timing forwards through fixed bindings and preserves solve and restart results", "[integration][solver][graph_timing]") {
    const auto physical = density_model();
    const auto resources = fixture::make_resources();
    const auto finite_elements = density_discretization(resources);
    using Description = decltype(physical.equations()(physical));
    using Plan = model::EquilibriumPlan<serif::eos::models::Polytrope, Description>;
    ::mfem::Vector displacement(resources->displacement.size());
    displacement = 0.0;
    auto state = model::StellarState(backend::equilibrium_layout<Plan>(finite_elements, displacement));
    std::ranges::fill(state.values(), 0.8);
    const auto configuration = model::make_discretization(model::DiscretizationOptions{
        .model = physical, .finite_elements = finite_elements, .normalization = numerical_coordinates(1.0),
        .preconditioner = solver::preconditioner::make_preconditioner<Density>(solver::preconditioner::strategies::inverse::Diagonal{})});
    const auto prescribed = dependencies::make_dependency_context();

    SECTION("physical and configured bindings accept owning states and explicit field views") {
        dependencies::GraphTiming timing;
        const auto observer_factory = dependencies::time_graph_nodes(timing);
        auto physical_state = backend::bind_stellar_model(physical, finite_elements, state, displacement, {}, prescribed, observer_factory);
        auto normalized_state = backend::bind_stellar_model(configuration, state, displacement, {}, prescribed,
            backend::detail::NoAdditionalComputations{}, observer_factory);
        const auto fields = backend::detail::state_coefficients(finite_elements, state, prescribed);
        auto physical_views = backend::bind_stellar_model(physical, finite_elements, fields.context(), displacement, {}, observer_factory);
        auto normalized_views = backend::bind_stellar_model(configuration, fields.context(), displacement, {}, observer_factory);
        REQUIRE(timing.runs().size() == 4);
        REQUIRE(physical_state.evaluate(state, displacement));
        REQUIRE(normalized_state.evaluate(state, displacement));
        REQUIRE(physical_views.evaluate(fields.context(), displacement));
        REQUIRE(normalized_views.evaluate(fields.context(), displacement));
        for (const auto& run : timing.runs()) {
            CHECK(std::ranges::any_of(run->nodes, [](const auto& node) { return node.completed != 0; }));
            CHECK(std::ranges::all_of(run->nodes, [](const auto& node) { return node.rejected == 0 && node.exceptions == 0; }));
        }
        const auto first = timing.runs().front()->nodes;
        REQUIRE(physical_state.evaluate(state, displacement));
        for (const auto& node : timing.runs().front()->nodes) {
            CHECK(node.calls == first[node.index].calls);
            CHECK(node.seconds == first[node.index].seconds);
        }
    }

    SECTION("timing does not change physical results, iteration counts, or checkpoint history") {
        auto options = solver::continuation::default_solver_options();
        const auto plain = solver::solve_equilibrium(configuration, state, displacement, options);
        dependencies::GraphTiming timing;
        const auto timed = solver::solve_equilibrium(configuration, state, displacement, options, solver::NoObserver{}, {}, prescribed,
            dependencies::time_graph_nodes(timing));
        REQUIRE(plain.converged());
        REQUIRE(timed.converged());
        CHECK(timed.diagnostics().accepted_iterations == plain.diagnostics().accepted_iterations);
        CHECK(timed.diagnostics().total_linear_iterations == plain.diagnostics().total_linear_iterations);
        CHECK(timed.diagnostics().line_search_trials == plain.diagnostics().line_search_trials);
        CHECK(timed.diagnostics().final_residual_norm == plain.diagnostics().final_residual_norm);
        CHECK(std::ranges::equal(timed.state().values(), plain.state().values()));
        REQUIRE(timing.runs().size() == 1);
        CHECK(std::ranges::any_of(timing.runs().front()->nodes, [](const auto& node) { return node.completed > 1; }));

        options.method.maximum_iterations = 0;
        const auto stopped = solver::solve_equilibrium(configuration, state, displacement, options);
        REQUIRE(stopped.has_checkpoint());
        CHECK(stopped.status() == solver::SolveStatus::maximum_iterations);
        options = solver::continuation::default_solver_options();
        const auto resumed = solver::resume_equilibrium(configuration, stopped.checkpoint(), displacement, options, solver::NoObserver{}, {}, prescribed,
            dependencies::time_graph_nodes(timing));
        REQUIRE(resumed.converged());
        CHECK(std::ranges::equal(resumed.state().values(), plain.state().values()));
        CHECK(resumed.diagnostics().accepted_iterations == plain.diagnostics().accepted_iterations);
        CHECK(resumed.diagnostics().final_residual_norm == plain.diagnostics().final_residual_norm);
        REQUIRE(timing.runs().size() == 2);
        CHECK(timing.runs()[0] != timing.runs()[1]);
    }
}

TEST_CASE("Moving-star node timings remain separate across a rebuilt checkpoint", "[integration][solver][graph_timing]") {
    const auto physical = fixture::make_model();
    const auto radial = solver::seed::make_polytropic_seed(physical);
    REQUIRE(radial);
    const auto resources = fixture::make_resources(radial->scales.radius.value());
    const auto finite_elements = fixture::make_discretization(resources, 0);
    const auto state = solver::seed::backend::mfem::make_polytropic_state(physical, finite_elements, *radial);
    REQUIRE(state);
    const auto configuration = model::make_discretization(model::DiscretizationOptions{
        .model = physical, .finite_elements = finite_elements, .normalization = numerical_coordinates(radial->scales.radius.value())});
    const auto prescribed = dependencies::make_dependency_context();
    dependencies::GraphTiming physical_timing;
    auto physical_evaluator = backend::bind_stellar_model(physical, finite_elements, *state, {}, prescribed,
        dependencies::time_graph_nodes(physical_timing));
    REQUIRE(physical_timing.runs().size() == 1);
    CHECK(std::ranges::all_of(physical_timing.runs().front()->nodes, [](const auto& node) { return node.calls == 0; }));
    CHECK(physical_evaluator.layout() == state->layout());

    auto options = solver::continuation::default_solver_options();
    options.method.maximum_iterations = 0;
    const auto plain = solver::solve_equilibrium(configuration, *state, options);
    dependencies::GraphTiming timing;
    const auto timed = solver::solve_equilibrium(configuration, *state, options, solver::NoObserver{}, {}, prescribed,
        dependencies::time_graph_nodes(timing));
    REQUIRE(plain.has_checkpoint());
    REQUIRE(timed.has_checkpoint());
    CHECK(timed.status() == plain.status());
    CHECK(timed.diagnostics().final_residual_norm == plain.diagnostics().final_residual_norm);
    CHECK(std::ranges::equal(timed.state().values(), plain.state().values()));
    REQUIRE(timing.runs().size() == 1);
    const auto before_resume = timing.runs().front()->nodes;
    const auto resumed = solver::resume_equilibrium(configuration, timed.checkpoint(), options, solver::NoObserver{}, {}, prescribed,
        dependencies::time_graph_nodes(timing));
    REQUIRE(resumed.has_checkpoint());
    CHECK(resumed.status() == plain.status());
    CHECK(resumed.diagnostics().final_residual_norm == plain.diagnostics().final_residual_norm);
    CHECK(std::ranges::equal(resumed.state().values(), plain.state().values()));
    REQUIRE(timing.runs().size() == 2);
    CHECK(timing.runs()[0] != timing.runs()[1]);
    for (const auto& node : timing.runs().front()->nodes) {
        CHECK(node.calls == before_resume[node.index].calls);
        CHECK(node.seconds == before_resume[node.index].seconds);
        CHECK(node.calls == timing.runs()[1]->nodes[node.index].calls);
    }
}