// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// root_finding_nd.hpp
#ifndef SANGI_ROOT_FINDING_ND_HPP
#define SANGI_ROOT_FINDING_ND_HPP

#include "root_finding_base.hpp"
#include <math/core/vector.hpp>
#include <math/core/matrix.hpp>

namespace sangi {

/**
 * @brief Solve a nonlinear system of equations by Newton's method
 * @tparam T A field type with order structure (e.g. real numbers)
 * @tparam V Vector type
 * @tparam M Matrix type
 * @param F Vector-valued function of the equation system
 * @param J Function that computes the Jacobian matrix
 * @param x0 Initial guess vector
 * @param criteria Convergence criteria
 * @return Result object containing the root vector and convergence status
 */
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T> && requires(M m, V v) {
        { m.solve(v) } -> std::convertible_to<std::optional<V>>;
    }
RootFindingResult<V> newton_raphson_nd(
    const std::function<V(const V&)>& F,
    const std::function<M(const V&)>& J,
    const V& x0,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>()) {

    V x = x0;
    V x_prev = x0;
    size_t iterations = 0;

    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        // Compute the function value and Jacobian at the current point
        V f = F(x);

        // Convergence test based on the function value norm
        T f_norm = norm2(f);
        if (f_norm < criteria.abs_ftol) {
            return RootFindingResult<V>::success(x, iterations, f_norm);
        }

        M jacobian = J(x);
        auto dx_opt = jacobian.solve(-f);

        // Singular matrix countermeasure: regularize by adding a small value to the diagonal
        if (!dx_opt) {
            T reg_factor = criteria.abs_xtol * 10;
            M reg_jacobian = jacobian;
            for (size_t j = 0; j < reg_jacobian.rows(); ++j) {
                reg_jacobian(j, j) += reg_factor;
            }
            dx_opt = reg_jacobian.solve(-f);
            if (!dx_opt) {
                return RootFindingResult<V>::failure(iterations, f_norm);
            }
        }

        V dx = *dx_opt;

        // Divergence prevention: limit the step size
        T step_size = norm2(dx);
        const T max_step = T(10);  // Maximum step size
        if (step_size > max_step) {
            dx = dx * (max_step / step_size);
        }

        x_prev = x;
        V x_new = x + dx;
        x = x_new;

        // Convergence test based on vector values
        if (criteria.is_vector_converged(f, x, x_prev)) {
            return RootFindingResult<V>::success(x, iterations + 1, f_norm);
        }
    }

    // Maximum number of iterations reached, but return the last approximation
    V f_final = F(x);
    T error = norm2(f_final);
    bool converged = error < criteria.abs_ftol * 10; // Relaxed convergence test
    
    return RootFindingResult<V>::partial_success(x, converged, iterations, error);
}

/**
 * @brief Newton-Krylov method (Jacobian-free Newton-Krylov, JFNK)
 *
 * The user only needs to provide the equation system F (neither an analytic
 * Jacobian nor an explicit construction of a numerical Jacobian is required).
 * At each Newton step the Newton equation \f$J\,\delta = -F\f$ is solved with
 * GMRES (Krylov), but without forming the Jacobian as a matrix; it proceeds
 * using only the finite-difference approximation of the directional derivative
 *   \f$J\,v \approx (F(x + \varepsilon v) - F(x))/\varepsilon\f$
 * (matrix-free). Intended for large-scale nonlinear systems whose Jacobian
 * cannot be assembled explicitly / is too large. The small Hessenberg subproblem
 * of GMRES uses sangi's Matrix.
 * Backtracking damping supplements global convergence.
 *
 * @tparam T A field type with order structure (e.g. real numbers)
 * @param F Vector-valued function of the equation system \f$F:\mathbb{R}^n\to\mathbb{R}^n\f$
 * @param x0 Initial guess vector
 * @param criteria Convergence criteria
 * @param krylov_restart Maximum subspace dimension of GMRES (restart)
 * @param inner_tol Stopping threshold for the GMRES internal relative residual (per Newton step)
 * @return Result object containing the root vector and convergence status
 */
template<concepts::OrderedField T>
RootFindingResult<Vector<T>> newton_krylov(
    const std::function<Vector<T>(const Vector<T>&)>& F,
    const Vector<T>& x0,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>(),
    size_t krylov_restart = 30,
    T inner_tol = T(1e-3)) {

    using V = Vector<T>;
    const size_t n = x0.size();
    const T tiny = std::numeric_limits<T>::epsilon();
    V x = x0;
    size_t iterations = 0;

    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        V Fx = F(x);
        T fnorm = norm2(Fx);
        if (fnorm < criteria.abs_ftol)
            return RootFindingResult<V>::success(x, iterations, fnorm);

        // matrix-free operator J·v ≈ (F(x+εv) − F(x))/ε
        const T sqeps = std::sqrt(std::numeric_limits<T>::epsilon());
        const T xnorm = norm2(x);
        auto applyJ = [&](const V& v) -> V {
            T vn = norm2(v);
            if (!(vn > T(0))) return V(n, T(0));
            T eps = sqeps * (T(1) + xnorm) / vn;
            return (F(x + v * eps) - Fx) / eps;
        };

        // Solve J δ = −Fx with GMRES(m) (Arnoldi + Givens, 1 cycle)
        V b = Fx * T(-1);
        T beta = norm2(b);
        const size_t m = std::min(krylov_restart, n);
        std::vector<V> basis;
        basis.reserve(m + 1);
        basis.push_back(b / beta);
        Matrix<T> H(m + 1, m, T(0));           // Hessenberg (sangi's Matrix)
        std::vector<T> cs(m, T(0)), sn(m, T(0)), g(m + 1, T(0));
        g[0] = beta;
        size_t kdim = 0;

        for (size_t j = 0; j < m; ++j) {
            V w = applyJ(basis[j]);
            for (size_t i = 0; i <= j; ++i) {   // modified Gram-Schmidt
                H(i, j) = dot(w, basis[i]);
                w = w - basis[i] * H(i, j);
            }
            T hjj = norm2(w);
            H(j + 1, j) = hjj;
            if (hjj > tiny) basis.push_back(w / hjj);
            // Apply the existing Givens rotations to column j
            for (size_t i = 0; i < j; ++i) {
                T t        =  cs[i] * H(i, j) + sn[i] * H(i + 1, j);
                H(i + 1, j) = -sn[i] * H(i, j) + cs[i] * H(i + 1, j);
                H(i, j)    = t;
            }
            // New Givens rotation
            T denom = std::sqrt(H(j, j) * H(j, j) + H(j + 1, j) * H(j + 1, j));
            kdim = j + 1;
            if (denom < tiny) break;
            cs[j] = H(j, j) / denom;
            sn[j] = H(j + 1, j) / denom;
            H(j, j)     = cs[j] * H(j, j) + sn[j] * H(j + 1, j);
            H(j + 1, j) = T(0);
            g[j + 1] = -sn[j] * g[j];
            g[j]     =  cs[j] * g[j];
            if (hjj <= tiny) break;                          // lucky breakdown
            if (std::abs(g[j + 1]) <= inner_tol * beta) break;  // inner convergence
        }

        // Obtain the least-squares solution y by back substitution (kdim×kdim upper triangular)
        std::vector<T> y(kdim, T(0));
        for (size_t i = kdim; i-- > 0; ) {
            T s = g[i];
            for (size_t k = i + 1; k < kdim; ++k) s -= H(i, k) * y[k];
            y[i] = (std::abs(H(i, i)) > tiny) ? s / H(i, i) : T(0);
        }
        V delta(n, T(0));
        for (size_t i = 0; i < kdim; ++i) delta = delta + basis[i] * y[i];

        // Backtracking: adopt the α that reduces ‖F(x+αδ)‖
        T alpha = T(1);
        V x_new = x + delta;
        T new_norm = norm2(F(x_new));
        for (int ls = 0; ls < 20 && new_norm > fnorm; ++ls) {
            alpha *= T(0.5);
            x_new = x + delta * alpha;
            new_norm = norm2(F(x_new));
        }
        V x_prev = x;
        x = x_new;

        if (norm2(delta * alpha) < criteria.abs_xtol)
            return RootFindingResult<V>::success(x, iterations + 1, new_norm);
    }

    V f_final = F(x);
    T error = norm2(f_final);
    return RootFindingResult<V>::partial_success(
        x, error < criteria.abs_ftol * 10, iterations, error);
}

/**
 * @brief Numerical-Jacobian Newton method (finite difference, F only)
 *
 * The user only needs to provide F. At each Newton step the Jacobian is
 * <strong>built as an explicit n×n matrix (sangi's Matrix)</strong> by forward
 * differences
 *   \f$J[:,j] = (F(x + h_j e_j) - F(x)) / h_j,\quad h_j = h\,(1+|x_j|)\f$
 * and the Newton equation \f$J\,\delta = -F\f$ is solved exactly by a direct
 * method (LU). Internally it is a thin wrapper that supplies a numerical
 * Jacobian to newton_raphson_nd.
 *
 * Difference from the matrix-free newton_krylov: this one holds J explicitly
 * (memory O(n²)) and solves the linear system directly (exactly). Suited for
 * robustly solving small-to-medium problems. When J cannot be assembled for
 * large-scale problems, use newton_krylov.
 *
 * @tparam T A field type with order structure (e.g. real numbers)
 * @param F Vector-valued function of the equation system
 * @param x0 Initial guess vector
 * @param criteria Convergence criteria
 * @param fd_step Reference step width \f$h\f$ of the forward difference (actually \f$h(1+|x_j|)\f$)
 * @return Result object containing the root vector and convergence status
 */
template<concepts::OrderedField T>
RootFindingResult<Vector<T>> newton_fd(
    const std::function<Vector<T>(const Vector<T>&)>& F,
    const Vector<T>& x0,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>(),
    T fd_step = std::sqrt(std::numeric_limits<T>::epsilon())) {

    using V = Vector<T>;
    using M = Matrix<T>;

    std::function<M(const V&)> J = [&F, fd_step](const V& x) -> M {
        const size_t n = x.size();
        V fx = F(x);
        M jac(n, n, T(0));
        for (size_t j = 0; j < n; ++j) {
            T h = fd_step * (T(1) + std::abs(x[j]));
            if (h == T(0)) h = fd_step;
            V xp = x;
            xp[j] += h;
            V fp = F(xp);
            for (size_t i = 0; i < n; ++i) jac(i, j) = (fp[i] - fx[i]) / h;   // J[:,j]
        }
        return jac;
    };

    return newton_raphson_nd<T, V, M>(F, J, x0, criteria);
}

/**
 * @brief Root finding for a nonlinear equation system by Broyden's method
 * @tparam T A field type with order structure (e.g. real numbers)
 * @tparam V Vector type
 * @tparam M Matrix type
 * @param F Vector-valued function of the equation system
 * @param x0 Initial guess vector
 * @param J0 Initial Jacobian matrix (optional)
 * @param criteria Convergence criteria
 * @return Result object containing the root vector and convergence status
 */
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T> && requires(M m, V v) {
        { m.solve(v) } -> std::convertible_to<std::optional<V>>;
    }
RootFindingResult<V> broyden_method(
    const std::function<V(const V&)>& F,
    const V& x0,
    std::optional<M> J0 = std::nullopt,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>()) {

    const size_t n = x0.size();
    V x = x0;
    V f = F(x);
    T f_norm = norm2(f);
    
    // Set up the initial Jacobian
    M B;  // Approximate Jacobian
    if (J0) {
        B = *J0;
    } else {
        // Set the initial Jacobian to the identity matrix
        B = identity<T>(n);
    }
    
    size_t iterations = 0;
    
    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        // Convergence test
        if (f_norm < criteria.abs_ftol) {
            return RootFindingResult<V>::success(x, iterations, f_norm);
        }

        // Compute the step
        auto dx_opt = B.solve(-f);
        if (!dx_opt) {
            // Singular matrix countermeasure
            T reg_factor = criteria.abs_xtol * 10;
            M reg_B = B;
            for (size_t j = 0; j < n; ++j) {
                reg_B(j, j) += reg_factor;
            }
            dx_opt = reg_B.solve(-f);
            if (!dx_opt) {
                return RootFindingResult<V>::failure(iterations, f_norm);
            }
        }
        
        V dx = *dx_opt;
        
        // Limit the step size
        T step_size = norm2(dx);
        const T max_step = T(10);
        if (step_size > max_step) {
            dx = dx * (max_step / step_size);
        }
        
        V x_new = x + dx;
        V f_new = F(x_new);
        T f_new_norm = norm2(f_new);
        
        // Line-search-like adjustment
        T alpha = T(1);
        if (f_new_norm > f_norm) {
            // Backtracking
            const T min_alpha = T(0.1);
            const T reduce_factor = T(0.5);

            while (f_new_norm > f_norm && alpha > min_alpha) {
                alpha *= reduce_factor;
                x_new = x + alpha * dx;
                f_new = F(x_new);
                f_new_norm = norm2(f_new);
            }

            if (f_new_norm > f_norm) {
                // Case where no improvement is observed
                return RootFindingResult<V>::partial_success(x, false, iterations, f_norm);
            }
        }

        // Broyden update
        V df = f_new - f;
        V s = alpha * dx;  // Actual step

        // Denominator check of the Broyden update
        T s_dot_df = dot(s, df);
        if (std::abs(s_dot_df) > std::numeric_limits<T>::epsilon()) {
            V Bs = B * s;
            V temp = (df - Bs) / s_dot_df;
            
            // Rank-1 update
            for (size_t i = 0; i < n; ++i) {
                for (size_t j = 0; j < n; ++j) {
                    B(i, j) += temp[i] * s[j];
                }
            }
        }
        
        // Update
        x = x_new;
        f = f_new;
        f_norm = f_new_norm;

        // Convergence test - check again
        if (f_norm < criteria.abs_ftol) {
            return RootFindingResult<V>::success(x, iterations + 1, f_norm);
        }
    }
    
    // Case where the maximum number of iterations is reached
    bool converged = f_norm < criteria.abs_ftol * 10;
    return RootFindingResult<V>::partial_success(x, converged, iterations, f_norm);
}

/**
 * @brief Nonlinear least squares (Levenberg-Marquardt method)
 * @tparam T A field type with order structure (e.g. real numbers)
 * @tparam V Vector type
 * @tparam M Matrix type
 * @param F Function returning the residual vector
 * @param J Function that computes the Jacobian matrix
 * @param x0 Initial guess vector
 * @param criteria Convergence criteria
 * @return Result object containing the optimal solution vector and convergence status
 */
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T> && requires(M m, V v) {
        { m.solve(v) } -> std::convertible_to<std::optional<V>>;
    }
RootFindingResult<V> levenberg_marquardt(
    const std::function<V(const V&)>& F,
    const std::function<M(const V&)>& J,
    const V& x0,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>()) {

    V x = x0;
    V f = F(x);
    T f_norm = norm2(f);
    T f_squared = dot(f, f);  // Sum of squared residuals

    // Damping parameter of the LM method
    T lambda = T(0.01);
    const T lambda_down = T(0.1);  // λ decrease factor
    const T lambda_up = T(10);     // λ increase factor

    size_t iterations = 0;
    size_t consecutive_failures = 0;
    const size_t max_failures = 5;  // Maximum number of consecutive failures
    
    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        // Convergence test
        if (f_norm < criteria.abs_ftol) {
            return RootFindingResult<V>::success(x, iterations, f_norm);
        }

        // Jacobian computation
        M jac = J(x);
        M JtJ = transpose(jac) * jac;
        V JtF = transpose(jac) * f;

        // Construction of the LM equation
        M A = JtJ;
        for (size_t i = 0; i < A.rows(); ++i) {
            A(i, i) += lambda * (A(i, i) + T(1));  // Strengthen the diagonal elements
        }

        // Compute the step
        auto dx_opt = A.solve(-JtF);
        if (!dx_opt) {
            consecutive_failures++;
            if (consecutive_failures >= max_failures) {
                return RootFindingResult<V>::partial_success(x, false, iterations, f_norm);
            }
            
            // Increase λ and retry
            lambda *= lambda_up;
            continue;
        }

        V dx = *dx_opt;

        // Trial step
        V x_new = x + dx;
        V f_new = F(x_new);
        T f_new_squared = dot(f_new, f_new);

        // Evaluate the step
        if (f_new_squared < f_squared) {
            // Success: accept the step and decrease λ
            x = x_new;
            f = f_new;
            f_norm = norm2(f);
            f_squared = f_new_squared;
            lambda *= lambda_down;
            consecutive_failures = 0;
        } else {
            // Failure: reject the step and increase λ
            lambda *= lambda_up;
            consecutive_failures++;
            
            if (consecutive_failures >= max_failures) {
                return RootFindingResult<V>::partial_success(x, false, iterations, f_norm);
            }
        }
    }
    
    // Case where the maximum number of iterations is reached
    bool converged = f_norm < criteria.abs_ftol * 10;
    return RootFindingResult<V>::partial_success(x, converged, iterations, f_norm);
}

/**
 * @brief Powell hybrid method (equivalent to MINPACK hybrd)
 *
 * A robust nonlinear equation system solver combining a finite-difference
 * Jacobian + dogleg trust region + Broyden rank-1 update. No analytic
 * derivative of the Jacobian is required.
 *
 * When Newton's method is applicable it uses the Newton step, and when the
 * Jacobian is ill-conditioned or singular it degrades to a scaled gradient
 * (Cauchy) step. The trust region ensures global convergence.
 *
 * @tparam T  Ordered field type
 * @tparam V  Vector type (operator[], size(), arithmetic operations)
 * @tparam M  Matrix type (operator(i,j), rows(), cols())
 * @param F        Equation system F(x) = 0
 * @param x0       Initial guess vector
 * @param criteria Convergence criteria
 * @return Root vector and convergence status
 */
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T>
RootFindingResult<V> powell_hybrid(
    const std::function<V(const V&)>& F,
    const V& x0,
    const ConvergenceCriteria<T>& criteria = ConvergenceCriteria<T>())
{
    const size_t n = x0.size();
    V x = x0;
    V f = F(x);
    T f_norm = norm2(f);

    // Internal Gaussian elimination solver: A*x = b (partial pivoting)
    auto solve_linear = [&](M A, V b) -> std::optional<V> {
        for (size_t k = 0; k < n; ++k) {
            // Pivot selection
            size_t max_row = k;
            T max_val = std::abs(A(k, k));
            for (size_t i = k + 1; i < n; ++i) {
                T v = std::abs(A(i, k));
                if (v > max_val) { max_val = v; max_row = i; }
            }
            if (max_val < std::numeric_limits<T>::epsilon() * T(100))
                return std::nullopt;
            if (max_row != k) {
                for (size_t j = k; j < n; ++j) std::swap(A(max_row, j), A(k, j));
                std::swap(b[max_row], b[k]);
            }
            // Forward elimination
            for (size_t i = k + 1; i < n; ++i) {
                T factor = A(i, k) / A(k, k);
                for (size_t j = k + 1; j < n; ++j)
                    A(i, j) -= factor * A(k, j);
                b[i] -= factor * b[k];
            }
        }
        // Back substitution
        V result = b;
        for (int k = static_cast<int>(n) - 1; k >= 0; --k) {
            for (size_t j = static_cast<size_t>(k) + 1; j < n; ++j)
                result[k] -= A(k, j) * result[j];
            result[k] /= A(k, k);
        }
        return result;
    };

    // Finite-difference Jacobian computation
    auto compute_jacobian = [&](const V& xc, const V& fc) -> M {
        M J = identity<T>(n);
        T eps_fd = std::sqrt(std::numeric_limits<T>::epsilon());
        for (size_t j = 0; j < n; ++j) {
            V xp = xc;
            T h = eps_fd * std::max(std::abs(xc[j]), T(1));
            xp[j] += h;
            V fp = F(xp);
            for (size_t i = 0; i < n; ++i) {
                J(i, j) = (fp[i] - fc[i]) / h;
            }
        }
        return J;
    };

    // Jacobian computation
    M J = compute_jacobian(x, f);

    // Trust region radius
    T delta = T(100) * norm2(x);
    if (delta < T(1)) delta = T(100);
    const T delta_max = T(1e8);

    // Diagonal scaling vector
    V diag = x;
    for (size_t j = 0; j < n; ++j) {
        T col_norm = T(0);
        for (size_t i = 0; i < n; ++i)
            col_norm += J(i, j) * J(i, j);
        diag[j] = std::max(std::sqrt(col_norm), T(1));
    }

    size_t iterations = 0;
    size_t jacobian_age = 0;
    const size_t jacobian_refresh = n * 2;

    for (iterations = 0; iterations < criteria.max_iterations; ++iterations) {
        // Convergence test
        if (f_norm < criteria.abs_ftol) {
            return RootFindingResult<V>::success(x, iterations, f_norm);
        }

        // Newton step: J * p_newton = -f
        V neg_f = f * T(-1);
        auto p_newton_opt = solve_linear(J, neg_f);
        V p_newton;
        bool newton_ok = false;
        T p_newton_norm = T(0);

        if (p_newton_opt) {
            p_newton = *p_newton_opt;
            p_newton_norm = norm2(p_newton);
            newton_ok = std::isfinite(p_newton_norm) && p_newton_norm < T(1e15);
        }

        // Cauchy step (scaled gradient): p_cauchy = -α * J^T * f
        V jtf = x;
        for (size_t j = 0; j < n; ++j) {
            T s = T(0);
            for (size_t i = 0; i < n; ++i) s += J(i, j) * f[i];
            jtf[j] = s;
        }
        T jtf_norm = norm2(jtf);

        V jjtf = f;
        for (size_t i = 0; i < n; ++i) {
            T s = T(0);
            for (size_t j = 0; j < n; ++j) s += J(i, j) * jtf[j];
            jjtf[i] = s;
        }
        T jjtf_norm = norm2(jjtf);

        T alpha = (jjtf_norm > T(0))
            ? (jtf_norm * jtf_norm) / (jjtf_norm * jjtf_norm)
            : T(1);
        V p_cauchy = jtf * (-alpha);
        T p_cauchy_norm = norm2(p_cauchy);

        // Dogleg step selection
        V step = p_cauchy;
        if (newton_ok && p_newton_norm <= delta) {
            // Newton step is inside the trust region → use it as is
            step = p_newton;
        } else if (p_cauchy_norm >= delta) {
            // Even the Cauchy step is outside the trust region → scale it
            step = p_cauchy * (delta / p_cauchy_norm);
        } else if (newton_ok) {
            // Dogleg interpolation: the point on the Cauchy → Newton line that intersects delta
            V diff = p_newton - p_cauchy;
            T diff_norm_sq = dot(diff, diff);
            T pc_dot_diff = dot(p_cauchy, diff);
            T pc_norm_sq = p_cauchy_norm * p_cauchy_norm;

            // Solve ||p_cauchy + tau * diff||² = delta²
            T discriminant = pc_dot_diff * pc_dot_diff
                - diff_norm_sq * (pc_norm_sq - delta * delta);
            if (discriminant >= T(0)) {
                T tau = (-pc_dot_diff + std::sqrt(discriminant)) / diff_norm_sq;
                tau = std::max(T(0), std::min(T(1), tau));
                step = p_cauchy + diff * tau;
            } else {
                step = p_cauchy * (delta / p_cauchy_norm);
            }
        }

        // Trial step
        V x_new = x + step;
        V f_new = F(x_new);
        T f_new_norm = norm2(f_new);

        // Actual reduction vs predicted reduction
        T actual_reduction = f_norm * f_norm - f_new_norm * f_new_norm;
        // Prediction: ||f||² - ||f + J*step||²
        V jstep = f;
        for (size_t i = 0; i < n; ++i) {
            T s = T(0);
            for (size_t j = 0; j < n; ++j) s += J(i, j) * step[j];
            jstep[i] = f[i] + s;
        }
        T predicted_norm_sq = dot(jstep, jstep);
        T predicted_reduction = f_norm * f_norm - predicted_norm_sq;

        T ratio = (std::abs(predicted_reduction) > std::numeric_limits<T>::epsilon() * f_norm * f_norm)
            ? actual_reduction / predicted_reduction
            : T(0);

        // Trust region update
        if (ratio < T(0.1)) {
            delta *= T(0.5);
            if (delta < criteria.abs_xtol) {
                // Trust region is tiny → regard as converged
                return RootFindingResult<V>::partial_success(
                    x, f_norm < criteria.abs_ftol * T(10), iterations, f_norm);
            }
        } else if (ratio > T(0.75)) {
            delta = std::min(T(2) * delta, delta_max);
        }

        // Step acceptance test
        if (ratio > T(0.0001)) {
            // Broyden rank-1 update: J_new = J + (df - J*s)*s^T / (s^T*s)
            V s = step;
            V df = f_new - f;
            T sts = dot(s, s);

            if (sts > std::numeric_limits<T>::epsilon()) {
                // J*s
                V js = f;
                for (size_t i = 0; i < n; ++i) {
                    T sum = T(0);
                    for (size_t j = 0; j < n; ++j) sum += J(i, j) * s[j];
                    js[i] = sum;
                }
                // (df - J*s) / (s^T*s)
                V update = (df - js) * (T(1) / sts);
                for (size_t i = 0; i < n; ++i)
                    for (size_t j = 0; j < n; ++j)
                        J(i, j) += update[i] * s[j];
            }

            x = x_new;
            f = f_new;
            f_norm = f_new_norm;
            ++jacobian_age;

            // Diagonal scaling update
            for (size_t j = 0; j < n; ++j) {
                T col_norm = T(0);
                for (size_t i = 0; i < n; ++i)
                    col_norm += J(i, j) * J(i, j);
                diag[j] = std::max(diag[j], std::sqrt(col_norm));
            }

            // Jacobian refresh
            if (jacobian_age >= jacobian_refresh) {
                J = compute_jacobian(x, f);
                jacobian_age = 0;
            }

            // Convergence test
            if (f_norm < criteria.abs_ftol) {
                return RootFindingResult<V>::success(x, iterations + 1, f_norm);
            }
        }
        // Case where ratio is too small: reject the step, delta has already been shrunk
    }

    bool converged = f_norm < criteria.abs_ftol * T(10);
    return RootFindingResult<V>::partial_success(x, converged, iterations, f_norm);
}

// Wrapper function for compatibility with the old interface
template<concepts::OrderedField T, typename V, typename M>
    requires concepts::VectorOf<V, T> && concepts::MatrixOf<M, T>
std::optional<V> newton_raphson_nd_legacy(
    const std::function<V(const V&)>& F,
    const std::function<M(const V&)>& J,
    const V& x0,
    T tolerance = std::numeric_limits<T>::epsilon() * 100,
    size_t max_iterations = 50) {

    ConvergenceCriteria<T> criteria;
    criteria.abs_ftol = tolerance;
    criteria.abs_xtol = tolerance;
    criteria.max_iterations = max_iterations;

    auto result = newton_raphson_nd<T, V, M>(F, J, x0, criteria);
    return result.root;
}

#if SANGI_HAS_MKL
// Add the MKL implementation here
#endif

} // namespace sangi

#endif // SANGI_ROOT_FINDING_ND_HPP