#pragma once

#include "serif/solver/continuation/driver.hpp"
#include "serif/backend/mfem/solver/continuation/detail/physical_transfer.hpp"
#include "serif/solver/equilibrium.hpp"

namespace serif::solver::continuation {
    /** @brief Continuation uses one absolute qualification threshold across all family members. */
    [[nodiscard]] inline SolverOptions<> default_solver_options() {
        return {.method = {.relative_tolerance = 0.0, .absolute_tolerance = 1.0e-10}};
    }

    namespace detail {
        /** @brief Adapt one fresh configured discretization per stage to the same public equilibrium solve. */
        template <typename Family, typename Factory, typename SolverOptions, typename Solve>
        [[nodiscard]] auto corrector(const Family& family, Factory& configure, SolverOptions options, Solve solve) {
            options.method.validate();
            utils::errors::serif_ensure<utils::errors::SERiFErrorCode::INVALID_ARGUMENT, std::invalid_argument>(
                options.method.relative_tolerance == 0.0 && options.method.absolute_tolerance > 0.0,
                "Continuation requires zero relative tolerance and a fixed positive absolute Newton tolerance.");
            using Configuration = std::invoke_result_t<Factory&, typename Family::ModelType>;
            using FiniteElements = std::remove_cvref_t<decltype(std::declval<Configuration>().finite_elements)>;
            using TransferGuard = decltype(backend::mfem::physical_transfer_guard(std::declval<const FiniteElements&>()));
            return [&family, &configure, options = std::move(options), solve = std::move(solve), transfer = std::optional<TransferGuard>{}]
                (auto model, auto state, CorrectionControls controls) mutable {
                const auto expected = family.model_signature(model);
                auto configuration = std::invoke(configure, std::move(model));
                utils::errors::serif_ensure<utils::errors::SERiFErrorCode::INVALID_ARGUMENT, std::invalid_argument>(
                    family.model_signature(configuration.model) == expected,
                    "The continuation configuration factory changed the requested physical model. Preserve its prescriptions and equation identity.");
                if (!transfer) {
                    transfer.emplace(backend::mfem::physical_transfer_guard(configuration.finite_elements));
                }
                utils::errors::serif_ensure<utils::errors::SERiFErrorCode::INVALID_ARGUMENT, std::invalid_argument>(
                    (*transfer)(configuration.finite_elements),
                    "Continuation cannot transfer physical coefficients to changed finite-element bindings, reference geometry, or deformation coordinates.");
                auto stage_options = options;
                stage_options.method.absolute_tolerance = controls.absolute_tolerance;
                stage_options.method.maximum_iterations = static_cast<int>(std::min<std::uint64_t>(
                    stage_options.method.maximum_iterations, controls.maximum_work));
                return std::invoke(solve, std::move(configuration), std::move(state), std::move(stage_options));
            };
        }
    }
}

namespace serif::solver {
    /**
     * @brief Continue moving-star equilibria through an explicitly declared physical family.
     * @param configure Own or borrow a callable that returns a configured discretization for its supplied model.
     * It must preserve the physical coordinate layout and finite-element reference of transferred states; fresh evaluators are implicit.
     * @note The explicit factory supports move-only resource policies without inventing a mutable clone/context hierarchy.
     */
    template <typename Family, typename Factory, typename State, typename Options = SolverOptions<>>
        requires std::copy_constructible<Options>
    [[nodiscard]] auto continue_equilibrium(const Family& family, Factory configure, const State& initial_state,
        continuation::Options<Family> continuation_options, Options solver_options = continuation::default_solver_options()) {
        const double tolerance = solver_options.method.absolute_tolerance;
        auto correct = continuation::detail::corrector(family, configure, std::move(solver_options), [](auto configuration, auto state, auto options) {
            return solve_equilibrium(std::move(configuration), std::move(state), std::move(options));
        });
        return continuation::run(family, initial_state, std::move(correct), std::move(continuation_options), tolerance);
    }

    /** @brief Start with the family's physics-owned default coordinates and bounded fallback route. */
    template <typename Family, typename Factory, typename State>
    [[nodiscard]] auto continue_equilibrium(const Family& family, Factory configure, const State& initial_state) {
        return continue_equilibrium(family, std::move(configure), initial_state, serif::solver::continuation::make_options(family));
    }

    /** @brief Continue a fixed-geometry model while preserving its prescribed displacement. */
    template <typename Family, typename Factory, typename State, typename Options = SolverOptions<>>
        requires std::copy_constructible<Options>
    [[nodiscard]] auto continue_equilibrium(const Family& family, Factory configure, const State& initial_state, const Vector& displacement,
        continuation::Options<Family> continuation_options, Options solver_options = continuation::default_solver_options()) {
        const double tolerance = solver_options.method.absolute_tolerance;
        auto correct = continuation::detail::corrector(family, configure, std::move(solver_options), [&displacement](auto configuration, auto state, auto options) {
            return solve_equilibrium(std::move(configuration), std::move(state), displacement, std::move(options));
        });
        return continuation::run(family, initial_state, std::move(correct), std::move(continuation_options), tolerance);
    }

    /** @brief Use the same default route and absolute qualification tolerance for fixed geometry. */
    template <typename Family, typename Factory, typename State>
    [[nodiscard]] auto continue_equilibrium(const Family& family, Factory configure, const State& initial_state, const Vector& displacement) {
        return continue_equilibrium(family, std::move(configure), initial_state, displacement, serif::solver::continuation::make_options(family));
    }
}