#include "serif/tests/equations/test_equations.hpp"
#include "serif/tests/solver/native_problem.hpp"
#include <cmath>
#include <expected>
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include "mfem/linalg/densemat.hpp"
#include "serif/backend/mfem/model/preconditioned_system.hpp"
#include "serif/discretization/blocks/fields.hpp"
#include "serif/model/state.hpp"
#include "serif/backend/mfem/solver/preconditioner/inverse/dense.hpp"
#include "serif/dimensions/type_alias.hpp"
#include "serif/solver/nonlinear/newton.hpp"

namespace {
    namespace solver = serif::solver;
    namespace model = serif::model;
    namespace blocks = serif::discretization::blocks;
    namespace backend = model::backend::mfem;
    namespace keys = serif::utils::dependencies;
    using Layout = blocks::SystemLayout<keys::KeyList<blocks::gravity::potential, blocks::density::mass, ::serif::tests::equations::Enthalpy>,
        keys::KeyList<blocks::density::mass, ::serif::tests::equations::Enthalpy, blocks::gravity::potential>>;
    using State = model::StellarState<Layout>;
    using System = backend::NormalizedSystem;
    using Prepared = backend::PreconditionedSystem<System>;

    /** @brief A curved gravity manifold with reordered rows and nonunit coordinates. */
    struct CurvedProblem : serif::tests::solver::NativeProblem {
        using System = ::System;
        using Layout = ::Layout;
        using Model = solver::NoModel;
        Layout coordinates{.state_sizes = {1, 1, 1}, .residual_sizes = {1, 1, 1}};
        int preparations{0};

        /** @brief The controlled problem has no descriptive stellar physics owner. */
        [[nodiscard]] std::shared_ptr<const Model> model_owner() const { return {}; }
        /** @brief Expose distinct physical value and residual identity orderings. */
        [[nodiscard]] const Layout& layout() const { return coordinates; }
        /** @brief Form the full derivative of the nonlinear residual. */
        [[nodiscard]] std::shared_ptr<const ::mfem::DenseMatrix> jacobian(const State& state) const {
            const double x = state.values()[1];
            auto matrix = std::make_shared<::mfem::DenseMatrix>(3);
            *matrix = 0.0;
            (*matrix)(0, 0) = 1.0; (*matrix)(0, 1) = 2.0 * x; (*matrix)(0, 2) = 1.0;
            (*matrix)(1, 1) = -2.0 * x; (*matrix)(1, 2) = 1.0;
            (*matrix)(2, 0) = 2.0; (*matrix)(2, 1) = -2.0 * std::cos(x);
            return matrix;
        }
        /** @brief Preserve the nonlinear residual while publishing a frozen physical and normalized system. */
        [[nodiscard]] std::expected<std::shared_ptr<const System>, std::string> evaluate(const State& state) {
            const double x = state.values()[1];
            const double potential = state.values()[0];
            const double gradient = state.values()[2];
            if (!std::isfinite(x) || x <= 0.0) { return std::unexpected("outside test domain"); }
            const double target = 0.5 + std::sin(0.5);
            const std::vector<double> residual{x * x + gradient + potential - target, gradient - x * x, 2.0 * (potential - std::sin(x))};
            const auto physical = backend::make_physical_system(3, 3,
                {{.row_offset = 0, .values = residual}},
                {{.row_offset = 0, .column_offset = 0, .block = serif::operators::backend::mfem::linearized_block(jacobian(state))}});
            const auto normalization = solver::normalization::make_diagonal_normalization({2.0, 3.0, 5.0}, {7.0, 11.0, 13.0});
            const auto scales = solver::normalization::prepare_stellar_reference_scales(serif::dimensions::MassValue{1.0}, serif::dimensions::LengthValue{1.0});
            return std::make_shared<const System>(physical, *normalization, *scales);
        }
        /** @brief Give the full Newton solve an exact physical inverse. */
        [[nodiscard]] std::expected<std::shared_ptr<const Prepared>, std::string> prepare(const State& state) {
            ++preparations;
            const auto system = evaluate(state);
            if (!system) { return std::unexpected(system.error()); }
            auto inverse = solver::preconditioner::backend::mfem::try_dense_inverse(*jacobian(state));
            if (!inverse) { return std::unexpected(inverse.error().message); }
            return std::make_shared<const Prepared>(*system, std::move(*inverse));
        }
        /** @brief Keep test-domain rejection distinguishable from accepted results. */
        [[nodiscard]] static bool recoverable(const std::string&) { return true; }
        /** @brief Forward the original failure diagnostic. */
        [[nodiscard]] static std::string failure_message(const std::string& failure) { return failure; }
    };
}

namespace {
    /** @brief Isolate the policy gate from equation convergence and linear work. */
    struct RejectedInitial : solver::StandardNewton {
        /** @brief Reject the initial physical candidate even when its residual is exactly zero. */
        [[nodiscard]] solver::TrialAssessment admissible(CurvedProblem&, const State&) const { return {.usable = false, .message = "Test policy rejects this state."}; }
    };
}

TEST_CASE("A zero residual cannot bypass the initial physical admissibility policy", "[solver][strategies]") {
    CurvedProblem problem;
    const State equilibrium(problem.layout(), {std::sin(0.5), 0.5, 0.25});
    const auto evaluation = problem.evaluate(equilibrium);
    REQUIRE(evaluation);
    CHECK(std::ranges::all_of((*evaluation)->residual(), [](double value) { return value == 0.0; }));
    const auto result = solver::solve_newton(problem, equilibrium, solver::SolverOptions{.strategies = RejectedInitial{}});
    CHECK(result.status() == solver::SolveStatus::initial_state_rejected);
    CHECK_FALSE(result.has_state());
    CHECK_FALSE(result.converged());
    CHECK(problem.preparations == 0);
}

namespace {
    using Border = blocks::rotation::angular_velocity;
    using BorderedLayout = blocks::SystemLayout<keys::KeyList<blocks::gravity::potential, blocks::density::mass, Border, ::serif::tests::equations::Enthalpy>,
        keys::KeyList<Border, blocks::density::mass, ::serif::tests::equations::Enthalpy, blocks::gravity::potential>>;
    using BorderedState = model::StellarState<BorderedLayout>;

    /** @brief An independently ordered generated border must remain in the retained Newton equation. */
    struct BorderedProblem : serif::tests::solver::NativeProblem {
        using System = ::System;
        using Layout = BorderedLayout;
        using Model = solver::NoModel;
        Layout coordinates{.state_sizes = {1, 1, 1, 1}, .residual_sizes = {1, 1, 1, 1}};
        CurvedProblem fields;

        /** @brief Expose the complete layout including its generated scalar coordinate. */
        const Layout& layout() const { return coordinates; }
        /** @brief Algebraic regression models have no descriptive physical owner. */
        std::shared_ptr<const Model> model_owner() const { return {}; }
        /** @brief Remove only the generated scalar when evaluating the independent three-field equations. */
        State field_state(const BorderedState& state) const { return State(fields.layout(), {state.values()[0], state.values()[1], state.values()[3]}); }
        /** @brief Add the independent scalar constraint while retaining different row and column orderings. */
        std::shared_ptr<const ::mfem::DenseMatrix> jacobian(const BorderedState& state) const {
            const auto physical = fields.jacobian(field_state(state));
            auto result = std::make_shared<::mfem::DenseMatrix>(4);
            *result = 0.0;
            (*result)(0, 2) = 1.0;
            constexpr int columns[]{0, 1, 3};
            for (int row = 0; row < 3; ++row) {
                for (int column = 0; column < 3; ++column) { (*result)(row + 1, columns[column]) = (*physical)(row, column); }
            }
            return result;
        }
        /** @brief Prepare immutable physical and normalized rows including the scalar target. */
        std::expected<std::shared_ptr<const System>, std::string> evaluate(const BorderedState& state) {
            auto evaluation = fields.evaluate(field_state(state));
            if (!evaluation) { return std::unexpected(evaluation.error()); }
            const auto original = (*evaluation)->physical()->residual();
            const auto physical = backend::make_physical_system(4, 4,
                {{.row_offset = 0, .values = {state.values()[2] - 0.2, original[0], original[1], original[2]}}},
                {{.row_offset = 0, .column_offset = 0, .block = serif::operators::backend::mfem::linearized_block(jacobian(state))}});
            const auto normalization = solver::normalization::make_diagonal_normalization({2.0, 3.0, 17.0, 5.0}, {19.0, 7.0, 11.0, 13.0});
            return std::make_shared<const System>(physical, *normalization, (*evaluation)->reference_scales());
        }
        /** @brief Apply an exact full inverse. */
        std::expected<std::shared_ptr<const Prepared>, std::string> prepare(const BorderedState& state) {
            auto evaluation = evaluate(state);
            if (!evaluation) { return std::unexpected(evaluation.error()); }
            auto inverse = solver::preconditioner::backend::mfem::try_dense_inverse(*jacobian(state));
            if (!inverse) { return std::unexpected(inverse.error().message); }
            return std::make_shared<const Prepared>(*evaluation, std::move(*inverse));
        }
        /** @brief Preserve a numerical rejection's original message. */
        static std::string failure_message(const std::string& failure) { return failure; }
        /** @brief Only synthetic domain rejection is recoverable in this regression. */
        static bool recoverable(const std::string&) { return true; }
    };
}

TEST_CASE("Newton retains independently ordered optional scalar borders", "[solver][strategies]") {
    BorderedProblem problem;
    const auto result = solver::solve_newton(problem, BorderedState(problem.layout(), {0.0, 0.8, -0.7, 0.0}),
        solver::SolverOptions{.method = {.relative_tolerance = 0.0, .absolute_tolerance = 1.0e-9}});
    INFO(result.message());
    REQUIRE(result.converged());
    CHECK(result.state().values()[0] == Catch::Approx(std::sin(0.5)).margin(1.0e-8));
    CHECK(result.state().values()[1] == Catch::Approx(0.5).margin(1.0e-8));
    CHECK(result.state().values()[2] == Catch::Approx(0.2).margin(1.0e-8));
    CHECK(result.state().values()[3] == Catch::Approx(0.25).margin(1.0e-8));
}