#include <array>
#include <limits>
#include <optional>
#include <vector>

#include <catch2/catch_test_macros.hpp>

#include "serif/solver/continuation/polytropic.hpp"
#include "serif/solver/continuation/driver.hpp"

namespace {
    namespace family = serif::solver::continuation;
    namespace driver = serif::solver::continuation;
    using Momentum = serif::model::FixedAngularMomentum;
    using State = std::array<double, 3>;

    /** @brief Define physical continuation without finite-element or nonlinear-solver dependencies. */
    [[nodiscard]] auto make_family() {
        const auto target = serif::model::make_stellar_model(serif::eos::models::Polytrope{2.0, 0.73}, []() {},
            Momentum{serif::dimensions::AngularMomentumValue{0.1}}, serif::model::FixedTotalMass{serif::dimensions::MassValue{1.0}});
        return family::in(family::PolytropicIndex{}, family::AngularMomentum{})
            .with(family::preserve_pressure_at_reference_density(serif::dimensions::DensityValue{0.2})).bind(target);
    }

    struct Diagnostics {
        std::uint64_t attempted_iterations{1};
        std::optional<double> final_residual_norm{0.0};
    };
    struct Report {
        bool success{true};
        bool current{true};
        State physical{};
        Diagnostics diagnostic{};
        /** @brief Test the same certificate interface as an owning nonlinear result. */
        [[nodiscard]] bool converged() const { return success; }
        /** @brief Rejected initial states have no prepared physical snapshot. */
        [[nodiscard]] bool has_state() const { return current; }
        /** @brief Transfer physical rather than normalized coordinates. */
        [[nodiscard]] const State& state() const { return physical; }
        /** @brief Expose attempted work even when correction fails. */
        [[nodiscard]] const Diagnostics& diagnostics() const { return diagnostic; }
        /** @brief Keep driver diagnostic composition independent of solver-specific status enums. */
        [[nodiscard]] std::string message() const { return success ? "converged" : "corrector rejected"; }
    };
}

TEST_CASE("Continuation qualifies starting models and transfers only successful physical states", "[unit][solver][continuation][driver]") {
    const auto bound = make_family();
    auto options = family::make_options(bound, family::starting_at(family::PolytropicIndex{}, 1.0), family::prefer(family::PolytropicIndex{}));
    options.steps.initial_step_fraction.fill(1.0);
    options.steps.maximum_rejections_per_axis = 1;
    const State initial{2.0, 3.0, 5.0};
    std::vector<State> inputs;
    std::vector<std::array<double, 2>> attempted;
    const auto result = driver::run(bound, initial, [&](auto model, State state, driver::CorrectionControls controls) {
        CHECK(controls.absolute_tolerance == 1.0e-8);
        inputs.push_back(state);
        const double index = model.equation_of_state().polytropic_index();
        const double momentum = model.template specification<Momentum>().target().value();
        attempted.push_back({index, momentum});
        if (index == 2.0 && momentum == 0.0) {
            return Report{.success = false, .physical = {999.0, 999.0, 999.0}, .diagnostic = {.attempted_iterations = 3, .final_residual_norm = 2.0}};
        }
        state[0] += 1.0;
        return Report{.physical = state, .diagnostic = {.attempted_iterations = 2, .final_residual_norm = 1.0e-10}};
    }, options, 1.0e-8);
    REQUIRE(result.converged());
    REQUIRE(result.last_converged());
    CHECK(attempted == std::vector<std::array<double, 2>>{{1.0, 0.0}, {2.0, 0.0}, {1.0, 0.1}, {2.0, 0.1}});
    CHECK(inputs[1] == inputs[2]);
    CHECK(inputs[2][0] == 3.0);
    CHECK(initial == State{2.0, 3.0, 5.0});
    CHECK(result.last_converged()->state() == State{5.0, 3.0, 5.0});
    CHECK(result.total_work() == 9);
    CHECK(result.total_attempts() == 4);
    CHECK(result.counters.rejections == 1);
    CHECK(result.accepted_coordinates == result.target);
}

TEST_CASE("Continuation starting fallbacks are bounded and explicit starts disable them", "[unit][solver][continuation][driver]") {
    const auto bound = make_family();
    std::vector<std::array<double, 2>> attempted;
    const auto correct = [&](auto model, State state, driver::CorrectionControls) {
        const auto coordinates = bound.coordinates_of(model);
        attempted.push_back(coordinates);
        return Report{.success = coordinates == bound.target_coordinates(), .physical = state,
            .diagnostic = {.attempted_iterations = 1, .final_residual_norm = coordinates == bound.target_coordinates() ? 0.0 : 1.0}};
    };
    const auto result = driver::run(bound, State{}, correct, family::make_options(bound), 1.0e-8);
    CHECK(result.converged());
    CHECK(attempted == std::vector<std::array<double, 2>>{{1.0, 0.0}, {2.0, 0.0}, {1.0, 0.1}, {2.0, 0.1}});
    CHECK(result.starting_attempts == 4);
    CHECK(result.counters.attempts == 0);
    attempted.clear();
    const auto stopped = driver::run(bound, State{}, correct, family::make_options(bound, family::starting_at(family::PolytropicIndex{}, 1.0)), 1.0e-8);
    CHECK_FALSE(stopped.converged());
    CHECK(stopped.last_converged() == nullptr);
    CHECK(stopped.status == driver::Status::starting_star_failed);
    CHECK(attempted.size() == 1);
}

TEST_CASE("Continuation preserves intermediate ownership when attempt work or progress limits stop it", "[unit][solver][continuation][driver]") {
    const auto bound = make_family();
    const auto correct = [](auto, State state, driver::CorrectionControls controls) {
        CHECK(controls.maximum_work >= 1);
        state[1] += 2.0;
        return Report{.physical = state};
    };
    auto options = family::make_options(bound);
    options.steps.maximum_attempts = 2;
    options.maximum_records = 1;
    const auto attempts = driver::run(bound, State{}, correct, options, 1.0e-8);
    CHECK_FALSE(attempts.converged());
    CHECK(attempts.status == driver::Status::attempt_limit);
    REQUIRE(attempts.last_converged());
    CHECK(attempts.last_converged()->state()[1] == 4.0);
    CHECK(attempts.stages.size() == 1);
    CHECK(attempts.omitted_records == 1);
    CHECK(attempts.total_attempts() == 2);
    options.steps.maximum_attempts = 128;
    options.maximum_work = 2;
    const auto work = driver::run(bound, State{}, correct, options, 1.0e-8);
    CHECK(work.status == driver::Status::work_limit);
    CHECK(work.total_work() == 2);
    CHECK(work.last_converged()->state()[1] == 4.0);
    CHECK(attempts.last_converged()->state()[1] == 4.0);
    options.maximum_work = 100;
    options.steps.maximum_rejections_per_axis = 1;
    int calls = 0;
    const auto blocked = driver::run(bound, State{}, [&](auto, State state, driver::CorrectionControls) {
        return Report{.success = calls++ == 0, .physical = state};
    }, options, 1.0e-8);
    CHECK(blocked.status == driver::Status::no_progress);
    REQUIRE(blocked.last_converged());
    CHECK(blocked.accepted_coordinates == bound.starting_coordinates());
}

TEST_CASE("A supplied starting star must belong to the family and qualify at the fixed absolute tolerance", "[unit][solver][continuation][driver]") {
    const auto bound = make_family();
    const driver::StartingStar anchor{bound.model_at({1.5, 0.02}), State{7.0, 11.0, 13.0}};
    int calls = 0;
    const auto correct = [&](auto model, State state, driver::CorrectionControls controls) {
        if (calls++ == 0) {
            CHECK(bound.coordinates_of(model) == std::array<double, 2>{1.5, 0.02});
            CHECK(state == anchor.state);
            CHECK(controls.absolute_tolerance == 1.0e-9);
        }
        return Report{.physical = state, .diagnostic = {.attempted_iterations = 1, .final_residual_norm = 1.0e-10}};
    };
    const auto result = driver::run(bound, anchor, correct, family::make_options(bound), 1.0e-9);
    CHECK(result.converged());
    CHECK(calls >= 2);
    CHECK(anchor.state == State{7.0, 11.0, 13.0});
    CHECK_THROWS_AS(driver::run(bound, anchor, correct, family::make_options(bound,
        family::starting_at(family::PolytropicIndex{}, 1.2)), 1.0e-9), std::invalid_argument);
    const auto false_certificate = driver::run(bound, anchor, [](auto, State state, driver::CorrectionControls) {
        return Report{.physical = state, .diagnostic = {.final_residual_norm = 1.0}};
    }, family::make_options(bound), 1.0e-9);
    CHECK_FALSE(false_certificate.converged());
    CHECK_FALSE(false_certificate.last_converged());
    CHECK(false_certificate.starting_attempts == 1);
    const auto unrelated = anchor.model.transform_specifications([](const auto& value) {
        if constexpr (std::same_as<std::remove_cvref_t<decltype(value)>, serif::model::FixedTotalMass>) {
            return serif::model::FixedTotalMass{serif::dimensions::MassValue{2.0}};
        } else {
            return value;
        }
    });
    CHECK_THROWS_AS(driver::run(bound, driver::StartingStar{unrelated, anchor.state}, correct, family::make_options(bound), 1.0e-9), std::invalid_argument);
}

TEST_CASE("Invalid continuation laws stop without replacing the last converged star", "[unit][solver][continuation][driver]") {
    const auto prototype = make_family();
    const auto bound = family::in(family::PolytropicIndex{}).with(family::law(family::PolytropicConstant{}, family::depends_on(family::PolytropicIndex{}), [](double index) {
        return index == 1.0 || index == 2.0 ? 0.73 : std::numeric_limits<double>::quiet_NaN();
    })).bind(prototype.target_model());
    auto options = family::make_options(bound);
    int calls = 0;
    const auto result = driver::run(bound, State{}, [&](auto, State state, driver::CorrectionControls) {
        ++calls;
        return Report{.physical = state};
    }, options, 1.0e-8);
    CHECK(result.status == driver::Status::invalid_model);
    CHECK(calls == 1);
    REQUIRE(result.last_converged());
    CHECK(result.accepted_coordinates == bound.starting_coordinates());
    CHECK(result.counters.fatal_failures == 1);
    CHECK(result.total_attempts() == 2);
    CHECK_THROWS_AS(driver::run(bound, State{}, [](auto, State state, driver::CorrectionControls) { return Report{.physical = state}; }, options, 0.0), std::invalid_argument);
}

TEST_CASE("Continuation rejects invalid residual certificates before promoting any physical state", "[unit][solver][continuation][driver]") {
    const auto bound = make_family();
    const State initial{2.0, 3.0, 5.0};
    const std::array<std::optional<double>, 4> invalid_norms{
        std::nullopt, -1.0, std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::infinity()};
    for (const auto invalid : invalid_norms) {
        for (const bool reject_initial : {false, true}) {
            auto options = family::make_options(bound, family::starting_at(family::PolytropicIndex{}, 1.0));
            options.steps.maximum_attempts = 2;
            int calls = 0;
            const auto result = driver::run(bound, initial, [&](auto, State state, driver::CorrectionControls) {
                if (++calls == 1 && !reject_initial) {
                    state[0] += 1.0;
                    return Report{.physical = state};
                }
                return Report{.physical = {999.0, 999.0, 999.0}, .diagnostic = {.final_residual_norm = invalid}};
            }, options, 1.0e-8);
            CHECK_FALSE(result.converged());
            REQUIRE(result.failed);
            REQUIRE_FALSE(result.stages.empty());
            CHECK(result.stages.back().disposition == driver::StageDisposition::not_converged);
            CHECK(initial == State{2.0, 3.0, 5.0});
            if (reject_initial) {
                CHECK(result.status == driver::Status::starting_star_failed);
                CHECK_FALSE(result.last_converged());
                CHECK(calls == 1);
            } else {
                CHECK(result.status == driver::Status::attempt_limit);
                REQUIRE(result.last_converged());
                CHECK(result.last_converged()->state() == State{3.0, 3.0, 5.0});
                CHECK(result.accepted_coordinates == options.start);
                CHECK(calls == 2);
            }
        }
    }
}