#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <limits>
#include <memory>
#include <stdexcept>
#include <vector>
#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

#include "serif/solver/continuation/controller.hpp"

namespace continuation_controller_test {
    using namespace serif::solver::continuation;
    using Catch::Approx;

    /** @brief Hold growth fixed to isolate the selected controller behavior. */
    template <std::size_t N> CoordinateOptions<N> fixed_steps(double initial, double minimum, double maximum) {
        CoordinateOptions<N> options;
        options.initial_step_fraction.fill(initial);
        options.minimum_step_fraction.fill(minimum);
        options.maximum_step_fraction.fill(maximum);
        options.growth_factor = 1.;
        return options;
    }

    struct FakeContext {
        std::array<double, 2> coordinates;
        std::array<double, 3> fields;
        int prepared_generation;
    };
}

TEST_CASE("Continuation proposals preserve accepted coordinates and external contexts until promotion",
          "[unit][solver][continuation_controller][transaction]") {
    using namespace continuation_controller_test;
    auto options = fixed_steps<2>(.5, .125, 1.);
    options.order = CoordinateOrder::preferred;
    CoordinateController<2> controller({1., 0.}, {2., .1}, options);
    auto accepted = std::make_unique<FakeContext>(FakeContext{{1., 0.}, {2., 3., 5.}, 7});
    const auto originalAddress = accepted.get();
    const auto p = controller.propose();
    REQUIRE(p);
    CHECK(controller.accepted() == controller.start());
    auto trial = std::make_unique<FakeContext>(FakeContext{p->coordinates, {97., 98., 99.}, 8});
    controller.resolve(p->id, AttemptOutcome::recoverable_rejection, 3);
    trial.reset();
    CHECK(accepted.get() == originalAddress);
    CHECK(accepted->fields == std::array<double, 3>{2., 3., 5.});
    CHECK(accepted->prepared_generation == 7);
    CHECK(controller.accepted() == controller.start());
    const auto retry = controller.propose();
    REQUIRE(retry);
    CHECK(retry->coordinates[0] == Approx(1.25));
    trial = std::make_unique<FakeContext>(FakeContext{retry->coordinates, {11., 13., 17.}, 9});
    controller.validate_resolution(retry->id, AttemptOutcome::success, 2);
    CHECK(controller.accepted() == controller.start());
    controller.resolve(retry->id, AttemptOutcome::success, 2);
    accepted.swap(trial); // No fallible user code between resolution and swap.
    CHECK(controller.accepted() == accepted->coordinates);
    CHECK(accepted->prepared_generation == 9);
    CHECK(controller.start() == std::array<double, 2>{1., 0.});
    CHECK(controller.target() == std::array<double, 2>{2., .1});
    CHECK(controller.counters().attempts == 2);
    CHECK(controller.counters().successes == 1);
    CHECK(controller.counters().rejections == 1);
    CHECK(controller.counters().work == 5);
}

TEST_CASE("A preferred coordinate can yield to another axis while strict stages cannot",
          "[unit][solver][continuation_controller][order]") {
    using namespace continuation_controller_test;
    auto options = fixed_steps<2>(1., .125, 1.);
    options.maximum_rejections_per_axis = 1;
    options.order = CoordinateOrder::preferred;
    // This corrector's feasible branch requires the second coordinate first.
    const auto correct = [](const auto &p) { return p.coordinates[0] <= p.coordinates[1]
        ? AttemptOutcome::success : AttemptOutcome::recoverable_rejection; };
    CoordinateController<2> preferred({0., 0.}, {1., 1.}, options);
    std::vector<std::size_t> attempted;
    while (const auto p = preferred.propose()) {
        attempted.push_back(p->axis);
        preferred.resolve(p->id, correct(*p), 1);
    }
    CHECK(attempted == std::vector<std::size_t>{0, 1, 0});
    CHECK(preferred.status() == ControllerStatus::complete);
    CHECK(preferred.accepted() == preferred.target());
    options.order = CoordinateOrder::strict;
    CoordinateController<2> strict({0., 0.}, {1., 1.}, options);
    const auto blocked = strict.propose();
    REQUIRE(blocked);
    strict.resolve(blocked->id, correct(*blocked), 1);
    CHECK(strict.status() == ControllerStatus::no_progress);
    CHECK_FALSE(strict.propose());
    CHECK(strict.accepted() == strict.start());
    CHECK(strict.counters().attempts == 1);
}

TEST_CASE("Automatic progress compares fractions while strict order honors the supplied permutation",
          "[unit][solver][continuation_controller][order]") {
    using namespace continuation_controller_test;
    auto options = fixed_steps<2>(.5, .125, .5);
    options.axis_order = {1, 0};
    CoordinateController<2> automatic({0., 1000.}, {1., -1000.}, options);
    const auto first = automatic.propose();
    REQUIRE(first);
    CHECK(first->axis == 1);
    CHECK(first->coordinates == std::array<double, 2>{0., 0.});
    automatic.resolve(first->id, AttemptOutcome::success);
    const auto second = automatic.propose();
    REQUIRE(second);
    CHECK(second->axis == 0); // Largest remaining fraction, independent of units.
    options.order = CoordinateOrder::strict;
    CoordinateController<2> strict({0., 1000.}, {1., -1000.}, options);
    auto p = strict.propose();
    REQUIRE(p);
    strict.resolve(p->id, AttemptOutcome::success);
    p = strict.propose();
    REQUIRE(p);
    CHECK(p->axis == 1);
    CHECK(p->coordinates == std::array<double, 2>{0., -1000.});
}

TEST_CASE("Adaptive retries shrink actual clipped progress and successful steps grow within limits",
          "[unit][solver][continuation_controller][adaptation]") {
    using namespace continuation_controller_test;
    auto options = fixed_steps<1>(.5, .05, .5);
    options.growth_factor = 2.;
    CoordinateController<1> controller({0.}, {1.}, options);
    auto p = controller.propose();
    REQUIRE(p);
    CHECK(p->coordinates[0] == .5);
    controller.resolve(p->id, AttemptOutcome::recoverable_rejection, 4);
    p = controller.propose();
    REQUIRE(p);
    CHECK(p->coordinates[0] == .25);
    controller.resolve(p->id, AttemptOutcome::success, 2);
    p = controller.propose();
    REQUIRE(p);
    CHECK(p->coordinates[0] == .75);
    controller.resolve(p->id, AttemptOutcome::success, 2);
    p = controller.propose();
    REQUIRE(p);
    CHECK(p->coordinates[0] == 1.);
    CHECK(p->step_fraction == .25);
    controller.resolve(p->id, AttemptOutcome::recoverable_rejection, 3);
    p = controller.propose();
    REQUIRE(p);
    CHECK(p->coordinates[0] == .875); // Half the clipped .25, not half the nominal .5.
    controller.resolve(p->id, AttemptOutcome::success, 1);
    p = controller.propose();
    REQUIRE(p);
    CHECK(p->coordinates[0] == 1.);
    controller.resolve(p->id, AttemptOutcome::success, 1);
    CHECK(controller.status() == ControllerStatus::complete);
    CHECK(controller.accepted()[0] == controller.target()[0]);
    CHECK(controller.counters().work == 13);
}

TEST_CASE("Continuation handles zero spans exact endpoints and budgets without extra attempts",
          "[unit][solver][continuation_controller][bounds]") {
    using namespace continuation_controller_test;
    CoordinateController<2> zero({2., .1}, {2., .1});
    CHECK(zero.status() == ControllerStatus::complete);
    CHECK_FALSE(zero.propose());
    CHECK(zero.counters().attempts == 0);
    auto options = fixed_steps<2>(1., .1, 1.);
    options.maximum_attempts = 1;
    CoordinateController<2> one({2., .1}, {2., -.2}, options);
    auto p = one.propose();
    REQUIRE(p);
    CHECK(p->axis == 1);
    one.resolve(p->id, AttemptOutcome::success);
    CHECK(one.status() == ControllerStatus::complete); // Success wins on the last allowed attempt.
    options.initial_step_fraction.fill(.5);
    CoordinateController<2> bounded({0., 0.}, {1., 1.}, options);
    p = bounded.propose();
    REQUIRE(p);
    bounded.resolve(p->id, AttemptOutcome::success);
    CHECK(bounded.status() == ControllerStatus::attempt_limit);
    CHECK_FALSE(bounded.propose());
    CHECK(bounded.accepted() == std::array<double, 2>{.5, 0.});
}

TEST_CASE("Repeated rejection terminates at minimum progress and cannot spin on unrepresentable steps",
          "[unit][solver][continuation_controller][bounds]") {
    using namespace continuation_controller_test;
    auto options = fixed_steps<1>(.5, .125, 1.);
    CoordinateController<1> reject({0.}, {1.}, options);
    std::vector<double> attempted;
    while (const auto p = reject.propose()) {
        attempted.push_back(p->coordinates[0]);
        reject.resolve(p->id, AttemptOutcome::recoverable_rejection, 1);
    }
    CHECK(attempted == std::vector<double>{.5, .25, .125});
    CHECK(reject.status() == ControllerStatus::no_progress);
    CHECK(reject.counters().work == 3);
    CHECK(reject.accepted() == reject.start());
    options.initial_step_fraction[0] = .25;
    const double start = 1.;
    const double target = std::nextafter(start, 2.);
    CoordinateController<1> rounded({start}, {target}, options);
    CHECK_FALSE(rounded.propose());
    CHECK(rounded.status() == ControllerStatus::no_progress);
    CHECK(rounded.counters().attempts == 0);
}

TEST_CASE("Decimal continuation endpoints remain exact with reverse and zero-length coordinates",
          "[unit][solver][continuation_controller][bounds]") {
    using namespace continuation_controller_test;
    auto options = fixed_steps<3>(.17, .01, .17);
    CoordinateController<3> controller({1.1, -.07, 3.}, {3.3, -.13, 3.}, options);
    while (const auto p = controller.propose()) {
        const auto previous = controller.accepted();
        for (std::size_t i = 0; i < 3; ++i) {
            if (i != p->axis) {
                CHECK(p->coordinates[i] == previous[i]);
            } else {
                CHECK(p->coordinates[i] != previous[i]);
                CHECK(p->coordinates[i] >= std::min(controller.start()[i], controller.target()[i]));
                CHECK(p->coordinates[i] <= std::max(controller.start()[i], controller.target()[i]));
            }
        }
        controller.resolve(p->id, AttemptOutcome::success, 1);
    }
    CHECK(controller.status() == ControllerStatus::complete);
    CHECK(controller.accepted() == controller.target());
    CHECK(controller.counters().attempts <= 14);
}

TEST_CASE("Invalid continuation resolutions leave pending and accepted state unchanged",
          "[unit][solver][continuation_controller][transaction][invalid]") {
    using namespace continuation_controller_test;
    auto options = fixed_steps<1>(.25, .125, 1.);
    CoordinateController<1> controller({0.}, {1.}, options);
    CHECK_THROWS_AS(controller.resolve(1, AttemptOutcome::success), std::logic_error);
    const auto p = controller.propose();
    REQUIRE(p);
    CHECK_THROWS_AS(controller.propose(), std::logic_error);
    CHECK_THROWS_AS(controller.resolve(p->id + 1, AttemptOutcome::success), std::logic_error);
    CHECK_THROWS_AS(controller.resolve(p->id, static_cast<AttemptOutcome>(99)), std::invalid_argument);
    CHECK(controller.pending()->id == p->id);
    CHECK(controller.accepted() == controller.start());
    CHECK(controller.counters().work == 0);
    controller.resolve(p->id, AttemptOutcome::success, std::numeric_limits<std::uint64_t>::max());
    const auto next = controller.propose();
    REQUIRE(next);
    const auto accepted = controller.accepted();
    CHECK_THROWS_AS(controller.validate_resolution(next->id, AttemptOutcome::success, 1), std::overflow_error);
    CHECK_THROWS_AS(controller.resolve(next->id, AttemptOutcome::success, 1), std::overflow_error);
    CHECK(controller.accepted() == accepted);
    CHECK(controller.pending()->id == next->id);
    CHECK(controller.counters().successes == 1);
    CHECK(controller.counters().work == std::numeric_limits<std::uint64_t>::max());
    controller.resolve(next->id, AttemptOutcome::fatal_failure, 0);
    CHECK(controller.status() == ControllerStatus::fatal_failure);
    CHECK(controller.accepted() == accepted);
    CHECK(controller.counters().fatal_failures == 1);
    CHECK_FALSE(controller.propose());
    CHECK_THROWS_AS(controller.resolve(next->id, AttemptOutcome::success), std::logic_error);
}

TEST_CASE("Continuation rejects unsupported controls before proposing or doing work",
          "[unit][solver][continuation_controller][invalid]") {
    using namespace continuation_controller_test;
    const double nan = std::numeric_limits<double>::quiet_NaN();
    const double inf = std::numeric_limits<double>::infinity();
    for (const double bad : {nan, inf, -inf}) {
        CHECK_THROWS_AS((CoordinateController<1>({bad}, {1.})), std::invalid_argument);
        CHECK_THROWS_AS((CoordinateController<1>({0.}, {bad})), std::invalid_argument);
    }
    CHECK_THROWS_AS((CoordinateController<1>({-std::numeric_limits<double>::max()},
        {std::numeric_limits<double>::max()})), std::invalid_argument);
    const auto invalid = [](auto change) {
        CoordinateOptions<2> options;
        change(options);
        CHECK_THROWS_AS((CoordinateController<2>({0., 0.}, {1., 1.}, options)), std::invalid_argument);
    };
    invalid([&](auto &o) { o.initial_step_fraction[0] = nan; });
    invalid([](auto &o) { o.minimum_step_fraction[1] = 0.; });
    invalid([](auto &o) { o.maximum_step_fraction[0] = 1.1; });
    invalid([](auto &o) { o.minimum_step_fraction[0] = .5; });
    invalid([](auto &o) { o.maximum_step_fraction[0] = .1; });
    invalid([](auto &o) { o.growth_factor = .9; });
    invalid([&](auto &o) { o.growth_factor = inf; });
    invalid([](auto &o) { o.shrink_factor = 0.; });
    invalid([](auto &o) { o.shrink_factor = 1.; });
    invalid([](auto &o) { o.maximum_attempts = 0; });
    invalid([](auto &o) { o.maximum_rejections_per_axis = 0; });
    invalid([](auto &o) { o.axis_order = {0, 0}; });
    invalid([](auto &o) { o.axis_order = {0, 2}; });
    invalid([](auto &o) { o.order = static_cast<CoordinateOrder>(99); });
}