#include "serif/tests/solver/native_problem.hpp"
#include <array>
#include <cmath>
#include <expected>
#include <functional>
#include <limits>
#include <memory>
#include <vector>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include "mfem/linalg/densemat.hpp"
#include "serif/solver/nonlinear/newton.hpp"
#include "serif/solver/normalization/diagonal.hpp"

namespace {
    namespace solver = serif::solver;
    /** @brief Keep one independent physical value for orchestration contracts. */
    struct State {
        std::array<double, 1> coordinates;
        /** @brief Edit only a local candidate's physical coordinate. */
        std::span<double> values() { return coordinates; }
        /** @brief Observe accepted coordinates without mutating them. */
        std::span<const double> values() const { return coordinates; }
    };
    /** @brief Freeze one identity equation with dimensionless coordinates. */
    struct System {
        std::array<double, 1> residual_values;
        ::mfem::DenseMatrix derivative{1};
        solver::normalization::DiagonalNormalization normalization = *solver::normalization::make_diagonal_normalization({1.0}, {1.0});
        /** @brief Retain residual values independently of subsequent evaluations. */
        explicit System(double value) : residual_values{value} { derivative(0, 0) = 1.0; }
        /** @brief Expose the complete residual. */
        const auto& residual() const { return residual_values; }
        /** @brief Expose its matching linearization. */
        const ::mfem::Operator& jacobian() const { return derivative; }
        /** @brief Transfer normalized corrections back to physical coordinates. */
        const auto& coordinates() const { return normalization; }
    };
    struct Identity : ::mfem::Solver {
        /** @brief The test derivative is exactly one. */
        Identity() : ::mfem::Solver(1) {}
        /** @brief A frozen identity does not require operator rebinding. */
        void SetOperator(const ::mfem::Operator&) override {}
        /** @brief Apply the exact identity inverse. */
        void Mult(const ::mfem::Vector& input, ::mfem::Vector& output) const override { output = input; }
    };
    struct Prepared {
        std::shared_ptr<const System> owner;
        mutable Identity inverse;
        /** @brief Retain the system before borrowing its inverse. */
        explicit Prepared(std::shared_ptr<const System> system) : owner(std::move(system)) {}
        /** @brief Expose the matching residual and derivative owner. */
        const auto& system() const { return owner; }
        /** @brief Borrow the retained exact inverse. */
        ::mfem::Solver& preconditioner() const { return inverse; }
        /** @brief Identity is a fixed linear operation. */
        auto contract() const { return solver::preconditioner::ApplicationContract::fixed_linear; }
    };
    struct Problem : serif::tests::solver::NativeProblem {
        using System = ::System;
        using Model = solver::NoModel;
        int evaluations{0};
        int preparations{0};
        /** @brief This algebraic regression has no physical model. */
        std::shared_ptr<const Model> model_owner() const { return {}; }
        /** @brief Publish immutable snapshots for the identity residual. */
        std::expected<std::shared_ptr<const System>, std::string> evaluate(const State& state) {
            ++evaluations;
            return std::make_shared<const System>(state.values()[0]);
        }
        /** @brief Count preparation independently of residual evaluation. */
        std::expected<std::shared_ptr<const Prepared>, std::string> prepare(const State& state) {
            ++preparations;
            return std::make_shared<const Prepared>(*evaluate(state));
        }
        /** @brief Preserve a numerical rejection's diagnostic. */
        std::string failure_message(const std::string& error) const { return error; }
        /** @brief The scalar fixture reports only numerical failures. */
        bool recoverable(const std::string&) const { return true; }
    };
    struct Scripted {
        std::function<void(::mfem::Vector&)> correction;
        /** @brief Supply a known correction to test orchestration rather than Krylov algebra. */
        solver::linear::Report solve(const ::mfem::Operator&, ::mfem::Solver&, solver::preconditioner::ApplicationContract,
            const ::mfem::Vector&, ::mfem::Vector& output, solver::linear::Control) const {
            correction(output);
            return {.status = solver::linear::Status::converged, .iterations = 3};
        }
    };
}

TEST_CASE("Newton numerical options and default metrics retain their migrated contracts", "[unit][solver][nonlinear][parity]") {
    solver::Newton options{.relative_tolerance = 0.25, .absolute_tolerance = 0.5};
    CHECK(options.convergence_threshold(4.0) == 1.0);
    options.absolute_tolerance = 2.0;
    CHECK(options.convergence_threshold(4.0) == 2.0);
    CHECK_THROWS_AS(options.convergence_threshold(-1.0), std::logic_error);
    CHECK_THROWS_AS(options.convergence_threshold(std::numeric_limits<double>::infinity()), std::logic_error);
    for (const auto invalid : {solver::Newton{.relative_tolerance = -1.0},
        solver::Newton{.absolute_tolerance = std::numeric_limits<double>::infinity()}, solver::Newton{.maximum_iterations = -1},
        solver::Newton{.linear_solve = {.maximum_iterations = 0}}}) {
        CHECK_THROWS_AS(invalid.validate(), std::logic_error);
    }
    for (const auto invalid : {solver::Backtracking{.initial_step = 0.0}, solver::Backtracking{.contraction = 1.0},
        solver::Backtracking{.fraction_to_boundary = 1.0}, solver::Backtracking{.sufficient_decrease = 0.0},
        solver::Backtracking{.minimum_step = 2.0}, solver::Backtracking{.maximum_trials = 0}}) {
        CHECK_THROWS_AS(invalid.validate(), std::logic_error);
    }
    const std::array residual{3.0, 4.0};
    const auto metric = solver::NormalizedResidualMetric{}(residual);
    CHECK(metric.residual_norm == 5.0);
    CHECK(metric.merit == 12.5);
    CHECK_FALSE(solver::MetricEvaluation{.residual_norm = -1.0}.valid());
    CHECK_FALSE(solver::MetricEvaluation{.merit = std::numeric_limits<double>::quiet_NaN()}.valid());
}

TEST_CASE("A zero Newton work budget evaluates its checkpoint without preparing any inverse", "[unit][solver][nonlinear][parity]") {
    Problem problem;
    const auto result = solver::solve_newton(problem, State{{2.0}}, solver::SolverOptions{.method = {.maximum_iterations = 0}});
    CHECK(result.status() == solver::SolveStatus::maximum_iterations);
    CHECK(result.has_state());
    CHECK(result.state().values()[0] == 2.0);
    CHECK(problem.evaluations == 1);
    CHECK(problem.preparations == 0);
    CHECK(result.diagnostics().attempted_iterations == 0);
    CHECK(result.diagnostics().line_search_trials == 0);
}

TEST_CASE("Newton events and counters distinguish exhausted globalization from an accepted final iteration", "[unit][solver][nonlinear][parity]") {
    for (const bool decreasing : {false, true}) {
        Problem problem;
        std::vector<solver::EventKind> events;
        const auto observer = [&](const solver::SolveEvent& event, const State& state, const auto&) {
            events.push_back(event.kind);
            if (event.kind == solver::EventKind::finished) {
                CHECK(event.diagnostics.attempted_iterations == 1);
                CHECK(event.diagnostics.accepted_iterations == static_cast<int>(decreasing));
                CHECK(event.diagnostics.total_linear_iterations == 3);
                CHECK(state.values()[0] == (decreasing ? 0.5 : 1.0));
            }
        };
        const auto result = solver::solve_newton(problem, State{{1.0}}, solver::SolverOptions{
            .method = {.relative_tolerance = 0.0, .absolute_tolerance = 0.0, .maximum_iterations = 1},
            .linear_solver = Scripted{[decreasing](auto& output) { output[0] = decreasing ? -0.5 : 0.0; }},
            .line_search = {.maximum_trials = 1}}, observer);
        CHECK(result.status() == (decreasing ? solver::SolveStatus::maximum_iterations : solver::SolveStatus::globalization_failure));
        CHECK(result.diagnostics().line_search_trials == 1);
        CHECK(result.diagnostics().rejected_trials == static_cast<int>(!decreasing));
        std::vector expected{solver::EventKind::initial, solver::EventKind::before_iteration};
        if (decreasing) {
            expected.push_back(solver::EventKind::before_accept);
            expected.push_back(solver::EventKind::accepted);
        } else {
            expected.push_back(solver::EventKind::trial_rejected);
        }
        expected.push_back(solver::EventKind::finished);
        CHECK(events == expected);
    }
}

TEST_CASE("A trial metric exception preserves published checkpoints and a fresh solve starts independently", "[unit][solver][nonlinear][parity]") {
    Problem problem;
    const State original{{1.0}};
    std::shared_ptr<const System> checkpoint;
    int metrics = 0;
    const auto metric = [&](std::span<const double> residual) {
        serif::utils::errors::serif_ensure<serif::utils::errors::SERiFErrorCode::INVALID_ARGUMENT, std::runtime_error>(
            ++metrics != 2, "injected trial metric failure");
        return solver::NormalizedResidualMetric{}(residual);
    };
    const auto observer = [&](const solver::SolveEvent& event, const State&, const auto& system) {
        if (event.kind == solver::EventKind::initial) { checkpoint = system; }
    };
    CHECK_THROWS_AS(solver::solve_newton(problem, original, solver::SolverOptions{.metric = metric}, observer), std::runtime_error);
    REQUIRE(checkpoint);
    CHECK(checkpoint->residual()[0] == 1.0);
    CHECK(original.values()[0] == 1.0);
    const auto recovered = solver::solve_newton(problem, original);
    CHECK(recovered.converged());
    CHECK(recovered.state().values()[0] == 0.0);
    CHECK(checkpoint->residual()[0] == 1.0);
}

TEST_CASE("A nonfinite Newton correction cannot become the next solve's incoming guess", "[unit][solver][nonlinear][parity]") {
    Problem problem;
    int calls = 0;
    auto options = solver::SolverOptions{.linear_solver = Scripted{[&](auto& output) {
        CHECK(output[0] == 0.0);
        output[0] = ++calls == 1 ? std::numeric_limits<double>::quiet_NaN() : -1.0;
    }}};
    CHECK(solver::solve_newton(problem, State{{1.0}}, options).status() == solver::SolveStatus::linear_solve_failure);
    CHECK(solver::solve_newton(problem, State{{1.0}}, options).converged());
    CHECK(calls == 2);
}