module yaeos__newton_solver
   !! # newton_solver
   !! Robust Newton-Raphson solver for systems of nonlinear equations F(x) = 0.
   !!
   !! @note
   !! This module was generated by a LLM.
   !! @endnote
   !!
   !! Algorithm:
   !!   Pure Newton with Levenberg-Marquardt (LM) regularization and
   !!   Armijo backtracking line search used as a globalization strategy.
   !!   The solver starts as pure Newton (lambda = 0) and promotes to LM
   !!   whenever the Jacobian is detected to be singular or ill-conditioned,
   !!   or when the line search cannot find a sufficient decrease.
   !!
   !! Dependencies:
   !!   LAPACK  –  dgetrf (LU factorisation with partial pivoting)
   !!              dgetrs (triangular back-substitution)
   !!              dgecon (1-norm condition number estimator)
   !!
   !! ## Examples
   !! ### Minimal call:
   !!```fortran
   !!call newton(my_sub, x)
   !!```
   !!
   !! ### Full call:
   !!```fortran
   !!type(newton_settings) :: s
   !!type(newton_result)   :: r
   !!s%atol      = 1.0e-10_pr
   !!s%verbosity = 2
   !!call newton(my_sub, x, s, r)
   !!if (r%status /= NEWTON_SUCCESS) write(*,*) newton_status_msg(r%status)
   !!```
   use iso_fortran_env, only: pr => real64
   implicit none
   private

   ! Public API
   public :: newton             ! Main solver subroutine
   public :: newton_settings    ! Input parameter type
   public :: newton_result      ! Output diagnostic type
   public :: newton_status_msg  ! Status-code to string

   ! Public status codes (returned in newton_result%status)
   public :: NEWTON_SUCCESS          ! 0 – ||F|| and/or ||dX|| below tolerance
   public :: NEWTON_MAX_ITS          ! 1 – iteration limit reached without convergence
   public :: NEWTON_NAN_INF          ! 2 – NaN or Inf in F at initial point or after step
   public :: NEWTON_SINGULAR         ! 3 – Jacobian singular even with max regularization
   public :: NEWTON_STAGNATION       ! 4 – ||dX|| stopped decreasing; no progress
   public :: NEWTON_LINE_SEARCH_FAIL ! 5 – lambda reached lambda_max with no accepted step

   integer, parameter :: NEWTON_SUCCESS          = 0
   integer, parameter :: NEWTON_MAX_ITS          = 1
   integer, parameter :: NEWTON_NAN_INF          = 2
   integer, parameter :: NEWTON_SINGULAR         = 3
   integer, parameter :: NEWTON_STAGNATION       = 4
   integer, parameter :: NEWTON_LINE_SEARCH_FAIL = 5

   ! ===========================================================================
   ! newton_settings – all algorithmic knobs with documented defaults
   ! ===========================================================================
   type :: newton_settings

      ! ---- Convergence -------------------------------------------------------

      real(pr) :: atol = 1.0e-9_pr
      !! Absolute residual tolerance.
      !! Converged when  max_i |F_i(x)| < atol.
      !! Units: same as F. Default: 1e-8.

      real(pr) :: rtol = 1.0e-9_pr
      !! Relative step tolerance.
      !! Converged when  max_i |dX_i| < rtol * (max_i |x_i| + atol).
      !! Catches the case where the step becomes negligible compared to x.
      !! Dimensionless. Default: 1e-6.

      integer :: max_its = 100
      !! Maximum number of Newton iterations before returning NEWTON_MAX_ITS.
      !! Default: 100.

      ! ---- Levenberg-Marquardt regularization --------------------------------
      !
      ! When active, the Newton linear system  J*dX = -F  is replaced by the
      ! LM normal equations  (J^T J + lambda*I)*dX = -J^T F.
      ! lambda is maintained as an *absolute* value; lambda_min/lambda_max are
      ! dimensionless multipliers applied to jacobian_scale (= ||J||_F^2/n),
      ! which gives a problem-scale-aware floor and ceiling.

      real(pr) :: lambda0 = 0.0_pr
      !! Initial value of the LM damping parameter.
      !! 0.0 = start as pure Newton; the solver activates LM automatically
      !! if conditioning is poor. Set > 0 to force LM from the first iteration.
      !! Units: [J]^2 (scales with the Jacobian entries squared). Default: 0.

      real(pr) :: lambda_min = 1.0e-6_pr
      !! Minimum non-zero lambda, expressed as a multiplier of jacobian_scale.
      !! Effective floor = lambda_min * ||J||_F^2/n.
      !! Prevents lambda from decaying to numerical zero after a good step.
      !! Dimensionless multiplier. Default: 1e-6.

      real(pr) :: lambda_max = 1.0e8_pr
      !! Maximum lambda, expressed as a multiplier of jacobian_scale.
      !! If lambda exceeds this ceiling the solver returns NEWTON_LINE_SEARCH_FAIL.
      !! Dimensionless multiplier. Default: 1e8.

      real(pr) :: lambda_up = 10.0_pr
      !! Factor by which lambda is multiplied on a line-search failure.
      !! Larger values recover faster from bad Jacobians but may overshoot.
      !! Dimensionless. Default: 10.

      real(pr) :: lambda_down = 5.0_pr
      !! Factor by which lambda is divided after a successful step with t > 0.1.
      !! Drives the solver back toward pure Newton as the iterate improves.
      !! Dimensionless. Default: 5.

      ! ---- Armijo backtracking line search -----------------------------------
      !
      ! Merit function:  m(x) = 0.5 * ||F(x)||_2^2
      ! Acceptance test: m(x + t*dX) <= m(x) + armijo_c * t * slope
      !   where slope = (J^T F) . dX  (the directional derivative of m).

      real(pr) :: armijo_c = 1.0e-4_pr
      !! Sufficient decrease constant (Wolfe c1 condition).
      !! Smaller values accept more steps but give weaker convergence guarantees.
      !! Must satisfy 0 < armijo_c < 0.5. Dimensionless. Default: 1e-4.

      real(pr) :: armijo_tau = 0.5_pr
      !! Step-length reduction factor applied each backtracking iteration.
      !! t_new = armijo_tau * t_old  until Armijo or t < t_min.
      !! Must satisfy 0 < armijo_tau < 1. Dimensionless. Default: 0.5.

      integer :: armijo_max_its = 50
      !! Maximum number of backtracking halvings per Newton step.
      !! If reached without satisfying Armijo, ls_failed is set .true.
      !! Default: 50.

      real(pr) :: t_min = 1.0e-8_pr
      !! Minimum accepted step length.
      !! If t < t_min the line search declares failure instead of accepting a
      !! micro-step that satisfies Armijo trivially (m changes by ~machine eps).
      !! Dimensionless (fraction of the full Newton step). Default: 1e-8.

      ! ---- Jacobian conditioning ---------------------------------------------

      real(pr) :: cond_max = 1.0e10_pr
      !! Condition number threshold above which LM regularization is activated
      !! (or strengthened) to stabilise the linear solve.
      !! Estimated via LAPACK dgecon. Dimensionless. Default: 1e10.

      ! ---- Stagnation detection ----------------------------------------------

      integer :: stagnation_nits = 5
      !! Number of consecutive iterations with negligible change in ||dX||
      !! before the solver exits with NEWTON_STAGNATION.
      !! Default: 5.

      real(pr) :: stagnation_tol = 1.0e-12_pr
      !! Relative threshold for stagnation: iteration is counted as stagnant
      !! when  |||dX||_prev - ||dX||_curr| < stagnation_tol * (||dX|| + 1).
      !! Dimensionless. Default: 1e-12.

      ! ---- Diagnostics -------------------------------------------------------

      logical :: save_history = .false.
      !! If .true., allocate result%f_history(0:iterations) and store
      !! ||F||_inf at each iteration. Slightly increases memory use.
      !! Default: .false.

      integer :: verbosity = 0
      !! Controls stdout output.
      !!   0 – silent
      !!   1 – print header + one-line summary at exit
      !!   2 – also print one line per accepted iteration
      !! Default: 0.

   end type newton_settings

   ! ===========================================================================
   ! newton_result – diagnostic output from the solver
   ! ===========================================================================
   type :: newton_result

      integer :: status = 0
      !! Exit status code. Compare with NEWTON_SUCCESS, NEWTON_MAX_ITS, etc.
      !! Use newton_status_msg(status) for a human-readable string.

      integer :: iterations = 0
      !! Total number of Newton iterations performed (accepted + rejected).

      real(pr) :: f_norm = 0.0_pr
      !! ||F(x)||_inf at exit. Should be < atol on NEWTON_SUCCESS.
      !! Units: same as F.

      real(pr) :: dx_norm = 0.0_pr
      !! ||dX||_inf of the last computed Newton step at exit.
      !! Units: same as x.

      real(pr) :: lambda = 0.0_pr
      !! Value of the LM regularisation parameter at exit.
      !! 0.0 means the final solve was a pure Newton step.
      !! Units: [J]^2.

      real(pr) :: cond_est = 0.0_pr
      !! Estimated condition number of J (or J^T J + lambda I when LM active)
      !! at the final iteration, via LAPACK dgecon. Dimensionless.

      real(pr), allocatable :: f_history(:)
      !! ||F||_inf at iterations 0, 1, ..., result%iterations.
      !! Allocated only when settings%save_history = .true.
      !! Size: (iterations + 1).

   end type newton_result

   ! ---------------------------------------------------------------------------
   ! Abstract interface for the user-supplied residual/Jacobian subroutine.
   ! The user must implement a subroutine with this exact signature.
   !
   ! Arguments:
   !   X  – [in]  current iterate, size n
   !   F  – [out] residual vector F(X), size n
   !   J  – [out] Jacobian matrix dF/dX, shape (n, n), stored dense row-major
   !              J(i,j) = dF_i / dX_j
   ! ---------------------------------------------------------------------------
   abstract interface
      subroutine to_solve(X, F, J)
         import pr
         real(pr), intent(in)  :: X(:)     ! Current iterate
         real(pr), intent(out) :: F(:)     ! Residual vector F(X)
         real(pr), intent(out) :: J(:, :)  ! Jacobian dF_i/dX_j
      end subroutine to_solve
   end interface

   external :: dgetrf, dgetrs, dgecon

contains

   ! ===========================================================================
   ! PUBLIC: newton
   !
   ! Solves F(x) = 0 starting from the initial guess x using Newton's method
   ! with LM regularization and Armijo backtracking.
   !
   ! Arguments:
   !   sub      [in]     – user subroutine, must match the to_solve interface
   !   x        [inout]  – initial guess on entry; best iterate on exit (n)
   !   settings [in,opt] – algorithmic settings; all defaults used if absent
   !   result   [out,opt]– diagnostics; ignored if not present
   ! ===========================================================================
   subroutine newton(sub, x, settings, result)
      procedure(to_solve)                          :: sub
      ! User-supplied residual/Jacobian routine (matches to_solve interface)

      real(pr),              intent(inout)         :: x(:)
      ! [in]  Initial guess, length n
      ! [out] Solution estimate on exit (best iterate, even if not converged)

      type(newton_settings), intent(in),  optional :: settings
      ! Solver configuration; defaults are used for any absent fields

      type(newton_result),   intent(out), optional :: result
      ! Diagnostic output: status, iteration count, norms, history

      ! ---- local scalars -----------------------------------------------------
      type(newton_settings) :: s          ! Active settings (defaults or user)
      type(newton_result)   :: res        ! Internal result accumulator

      integer  :: n                       ! Problem size (number of unknowns)
      integer  :: its                     ! Current iteration counter
      integer  :: stag_count              ! Consecutive stagnant iterations

      real(pr) :: t                       ! Current line-search step length in (0,1]
      real(pr) :: lambda                  ! LM regularization parameter (absolute)
      real(pr) :: jscale                  ! ||J||_F^2/n; scale for lambda floor/ceiling
      real(pr) :: cond_est                ! Estimated cond(J) from dgecon

      real(pr) :: f_norm                  ! ||F(x)||_inf at current iterate
      real(pr) :: f_norm_new              ! ||F(x + t*dX)||_inf after line search
      real(pr) :: dx_norm                 ! ||dX||_inf of current Newton step
      real(pr) :: dx_norm_prev            ! ||dX||_inf from previous iteration

      logical  :: converged               ! .true. when a convergence criterion is met
      logical  :: stagnated               ! .true. when stag_count >= stagnation_nits
      logical  :: ls_failed               ! .true. when line search returned failed=.true.
      logical  :: small_step              ! .true. when accepted t is numerically tiny

      ! ---- local arrays ------------------------------------------------------
      real(pr), allocatable :: F(:)       ! Residual at current x, size n
      real(pr), allocatable :: J(:,:)     ! Jacobian at current x, shape (n,n)
      real(pr), allocatable :: dX(:)      ! Newton (or LM) step, size n
      real(pr), allocatable :: X_old(:)   ! x saved before applying the step
      real(pr), allocatable :: F_new(:)   ! Residual at trial point x+t*dX
      real(pr), allocatable :: J_new(:,:) ! Jacobian at trial point x+t*dX
      real(pr), allocatable :: hist_buf(:)! Scratch buffer for f_history

      ! ---- setup -------------------------------------------------------------
      if (present(settings)) then
         s = settings
      else
         s = newton_settings()   ! all defaults
      end if

      n = size(x)
      allocate(F(n), J(n,n), dX(n), X_old(n), F_new(n), J_new(n,n))
      if (s%save_history) allocate(hist_buf(s%max_its + 1))

      lambda       = s%lambda0
      stag_count   = 0
      converged    = .false.
      stagnated    = .false.
      ls_failed    = .false.
      small_step   = .false.
      dx_norm      = huge(1.0_pr)
      dx_norm_prev = huge(1.0_pr)
      cond_est     = 0.0_pr

      ! ---- initial evaluation ------------------------------------------------
      call sub(x, F, J)

      if (has_nan_inf(F)) then
         res%status = NEWTON_NAN_INF
         call pack_result(res, result, s, hist_buf, 0, &
            huge(1.0_pr), huge(1.0_pr), lambda, cond_est)
         return
      end if

      f_norm = norm_inf(F)
      if (s%save_history) hist_buf(1) = f_norm
      if (s%verbosity >= 1) call print_header(n)

      ! ========================================================================
      main_loop: do its = 1, s%max_its

         ! 1. Convergence check on ||F|| ---------------------------------------
         if (f_norm < s%atol) then
            converged = .true.
            exit main_loop
         end if

         ! 2. Jacobian scale for problem-aware lambda bounds -------------------
         !    jscale = ||J||_F^2 / n  (mean squared Jacobian entry).
         !    Used so lambda_min/lambda_max are dimensionless multipliers that
         !    automatically adapt to problem scaling; avoids the ~0 collapse
         !    that occurred with the old diagonal geometric-mean approach.
         jscale = jacobian_scale(J)

         ! 3. Linear solve (Newton or LM) --------------------------------------
         !    solve_lm attempts a pure Newton step first; if J is found to be
         !    singular or ill-conditioned it promotes to the LM normal equations.
         !    The recursion was removed (rev 3): promotion is now an explicit
         !    sequential branch inside solve_lm (see that subroutine).
         call solve_lm(J, F, lambda, jscale, s, dX, cond_est)

         dx_norm = norm_inf(dX)

         ! 4. Mixed convergence check on ||dX|| --------------------------------
         if (dx_norm < s%rtol * (norm_inf(x) + s%atol)) then
            converged = .true.
            exit main_loop
         end if

         ! 5. Stagnation detection ---------------------------------------------
         !    Use a relative threshold so it adapts to the scale of ||dX||.
         if (abs(dx_norm_prev - dx_norm) < &
            s%stagnation_tol * (dx_norm + 1.0_pr)) then
            stag_count = stag_count + 1
            if (stag_count >= s%stagnation_nits) then
               stagnated = .true.
               exit main_loop
            end if
         else
            stag_count = 0
         end if
         dx_norm_prev = dx_norm

         ! 6. Armijo backtracking line search ----------------------------------
         !    Evaluate at t=1 (full Newton step) first, then reduce t if
         !    needed until sufficient decrease is achieved or t < t_min.
         X_old = x
         t     = 1.0_pr
         call sub(X_old + t*dX, F_new, J_new)
         f_norm_new = norm_inf(F_new)

         call armijo_backtrack(sub, X_old, F, J, dX, s, t, F_new, J_new, f_norm_new, ls_failed)

         ! Detect a micro-step: t so small that t*||dX|| is negligible relative
         ! to the current iterate (Armijo satisfied trivially, no real progress).
         small_step = (t < s%t_min) .and. &
            (t * dx_norm < s%atol + s%rtol * norm_inf(x))

         ! 7. Handle line-search failure ---------------------------------------
         !    Increase lambda and retry the linear solve in the next iteration.
         !    Do NOT advance x; stay at X_old and try a more regularised step.
         if (ls_failed .or. f_norm_new >= f_norm .or. small_step) then

            if (lambda == 0.0_pr) lambda = jscale * s%lambda_min
            lambda = min(lambda * s%lambda_up, jscale * s%lambda_max)

            if (s%verbosity >= 2) &
               write(*,'(5X,"[ls fail] iter=",I4, &
               "  t=",ES8.2,"  lambda->",ES10.3)') its, t, lambda

            ! Give up if lambda has reached the ceiling
            if (lambda >= jscale * s%lambda_max) then
               res%status = NEWTON_LINE_SEARCH_FAIL
               x = X_old      ! return last valid iterate
               call pack_result(res, result, s, hist_buf, its, &
                  f_norm, dx_norm, lambda, cond_est)
               if (s%verbosity >= 1) call print_summary(res)
               return
            end if
            cycle main_loop
         end if

         ! 8. Accept the step --------------------------------------------------
         x      = X_old + t * dX
         F      = F_new
         J      = J_new
         f_norm = f_norm_new

         ! Relax lambda only when the step was taken near its full length
         ! (t > 0.1 implies the Newton direction was trustworthy).
         ! A tiny accepted t means the Jacobian is still unreliable.
         if (lambda > 0.0_pr .and. t > 0.1_pr) &
            lambda = max(lambda / s%lambda_down, s%lambda0)

         if (s%save_history) hist_buf(its + 1) = f_norm

         if (s%verbosity >= 2) &
            write(*,'(I6, 4ES14.5)') its, f_norm, dx_norm, t, lambda

      end do main_loop
      ! ========================================================================

      ! ---- classify exit condition -------------------------------------------
      if (converged) then
         res%status = NEWTON_SUCCESS
      else if (stagnated) then
         res%status = NEWTON_STAGNATION
      else
         res%status = NEWTON_MAX_ITS
      end if

      call pack_result(res, result, s, hist_buf, &
         min(its, s%max_its), f_norm, dx_norm, lambda, cond_est)
      if (s%verbosity >= 1) call print_summary(res)

   end subroutine newton


   ! ===========================================================================
   ! PRIVATE: armijo_backtrack
   !
   ! Reduces the step length t until the Armijo sufficient decrease condition
   ! is satisfied or t falls below t_min (in which case failed = .true.).
   !
   ! Merit function:  m(x) = 0.5 * ||F(x)||_2^2
   ! Armijo test:     m(x + t*dX) <= m(x) + armijo_c * t * slope
   !   where          slope = (J^T * F) . dX   [directional derivative of m]
   !
   ! Arguments:
   !   sub        [in]     – residual/Jacobian routine
   !   X_old      [in]     – current iterate before the step, size n
   !   F          [in]     – F(X_old), size n
   !   J          [in]     – J(X_old), shape (n,n)
   !   dX         [in]     – proposed Newton/LM step direction, size n
   !   s          [in]     – solver settings
   !   t          [inout]  – step length; 1.0 on entry, accepted value on exit
   !   F_new      [inout]  – F(X_old + t*dX) on entry; updated value on exit
   !   J_new      [inout]  – J(X_old + t*dX) on entry; updated value on exit
   !   f_norm_new [inout]  ��� ||F_new||_inf; updated on exit
   !   failed     [out]    – .true. if no acceptable t was found
   ! ===========================================================================
   subroutine armijo_backtrack(sub, X_old, F, J, dX, s, &
      t, F_new, J_new, f_norm_new, failed)
      procedure(to_solve)            :: sub
      real(pr),              intent(in)    :: X_old(:)   ! Iterate before step
      real(pr),              intent(in)    :: F(:)       ! F(X_old)
      real(pr),              intent(in)    :: J(:,:)     ! J(X_old)
      real(pr),              intent(in)    :: dX(:)      ! Newton step direction
      type(newton_settings), intent(in)    :: s          ! Solver settings
      real(pr),              intent(inout) :: t          ! Step length [0,1]
      real(pr),              intent(inout) :: F_new(:)   ! F at trial point
      real(pr),              intent(inout) :: J_new(:,:) ! J at trial point
      real(pr),              intent(inout) :: f_norm_new ! ||F_new||_inf
      logical,               intent(out)   :: failed     ! .true. = no step found

      integer  :: k       ! Backtracking iteration counter
      real(pr) :: m_old   ! Merit value at X_old:  0.5*||F||_2^2
      real(pr) :: m_new   ! Merit value at trial:  0.5*||F_new||_2^2
      real(pr) :: slope   ! Directional derivative of m at X_old along dX

      failed = .false.
      m_old  = 0.5_pr * dot_product(F, F)

      ! slope = nabla_m . dX = (J^T F) . dX
      ! For a Newton step dX = -J^{-1} F we have slope = -||F||_2^2 < 0,
      ! guaranteeing a descent direction. Under LM the step may not be a
      ! true descent; skip the Armijo check in that case and accept t=1.
      slope = dot_product(matmul(transpose(J), F), dX)
      if (slope >= 0.0_pr) return   ! non-descent direction; accept full step

      do k = 1, s%armijo_max_its

         ! Hard lower bound on t: below t_min the step is numerically zero
         ! and Armijo would be satisfied trivially with no real progress.
         if (t < s%t_min) then
            failed = .true.
            return
         end if

         ! Armijo sufficient decrease test
         m_new = 0.5_pr * dot_product(F_new, F_new)
         if (m_new <= m_old + s%armijo_c * t * slope) return

         ! Reduce step length and re-evaluate
         t = t * s%armijo_tau

         call sub(X_old + t*dX, F_new, J_new)

         if (has_nan_inf(F_new)) then
            ! NaN at this t: keep reducing; fail on last attempt
            if (k == s%armijo_max_its) then
               failed = .true.
               return
            end if
            cycle
         end if

         f_norm_new = norm_inf(F_new)
      end do

      ! Exhausted backtracking iterations without satisfying Armijo
      failed = .true.

   end subroutine armijo_backtrack


   ! ===========================================================================
   ! PRIVATE: solve_lm
   !
   ! Solves the Newton or Levenberg-Marquardt linear system for the step dX.
   !
   ! If lambda = 0:  attempts to solve  J * dX = -F  (pure Newton).
   !   If J is found to be singular (dgetrf info /= 0) or ill-conditioned
   !   (cond_est > cond_max), lambda is set to jscale * lambda_min and the
   !   LM branch is executed immediately in the same call (no recursion).
   !
   ! If lambda > 0:  solves the LM normal equations
   !                 (J^T J + lambda * I) * dX = -J^T F.
   !
   ! Note: the recursion present in rev 1-2 has been removed (rev 3).
   !   The Newton-fails-fallback-to-LM logic is now a plain sequential
   !   if/else so the subroutine is not recursive and compiles without
   !   the RECURSIVE keyword.
   !
   ! Arguments:
   !   J        [in]     – Jacobian matrix, shape (n,n)
   !   F        [in]     – Residual vector, size n
   !   lambda   [inout]  – LM damping parameter; may be updated on exit
   !   jscale   [in]     – ||J||_F^2/n; used to compute dimensional lambda bounds
   !   s        [in]     – Solver settings
   !   dX       [out]    – Computed step, size n;  0 if solve fails completely
   !   cond_est [out]    – Estimated condition number of the solved matrix
   ! ===========================================================================
   subroutine solve_lm(J, F, lambda, jscale, s, dX, cond_est)
      real(pr),              intent(in)    :: J(:,:)    !! Jacobian at x
      real(pr),              intent(in)    :: F(:)      !! Residual at x
      real(pr),              intent(inout) :: lambda    !! LM damping (absolute)
      real(pr),              intent(in)    :: jscale    !! ||J||_F^2/n
      type(newton_settings), intent(in)    :: s         !! Solver settings
      real(pr),              intent(out)   :: dX(:)     !! Computed Newton/LM step
      real(pr),              intent(out)   :: cond_est  !! cond(A) estimate

      integer  :: n      !! Problem size
      integer  :: info   !! LAPACK return code (0 = success)
      real(pr) :: anorm  !! 1-norm of A (input to dgecon)
      real(pr) :: rcond  !! Reciprocal condition number from dgecon

      real(pr), allocatable :: A(:,:)    !! Working copy of J (or J^T J + lam I)
      real(pr), allocatable :: rhs(:)    !! Right-hand side (-F or -J^T F)
      real(pr), allocatable :: work(:)   !! LAPACK workspace, size 4n
      integer,  allocatable :: ipiv(:)   !! Pivot indices from dgetrf, size n
      integer,  allocatable :: iwork(:)  !! Integer workspace for dgecon, size n

      logical :: use_lm   !! .true. when the LM branch is used

      n = size(F)
      allocate(A(n,n), rhs(n), ipiv(n), work(4*n), iwork(n))

      ! -----------------------------------------------------------------------
      ! Decide initial branch: lambda = 0 → try Newton; lambda > 0 → go LM
      ! -----------------------------------------------------------------------
      use_lm = (lambda > 0.0_pr)

      if (.not. use_lm) then
         ! ---- Attempt pure Newton: J * dX = -F ------------------------------
         A   = J
         rhs = -F

         anorm = mat1norm(A)
         call dgetrf(n, n, A, n, ipiv, info)

         if (info /= 0) then
            ! Exact singularity: promote to LM
            lambda = jscale * s%lambda_min
            use_lm = .true.
         else
            ! Estimate condition number of J
            call dgecon('1', n, A, n, anorm, rcond, work, iwork, info)
            cond_est = 1.0_pr / max(rcond, tiny(rcond))

            if (cond_est > s%cond_max) then
               ! Ill-conditioned: promote to LM
               lambda = jscale * s%lambda_min
               use_lm = .true.
            else
               ! Newton solve accepted: back-substitute
               call dgetrs('N', n, 1, A, n, ipiv, rhs, n, info)
               dX = rhs
            end if
         end if
      end if

      if (use_lm) then
         ! ---- LM normal equations: (J^T J + lambda I) * dX = -J^T F --------
         A   = matmul(transpose(J), J)   ! J^T J,  shape (n,n), symmetric PSD
         rhs = -matmul(transpose(J), F)  ! -J^T F, size n
         call add_diagonal(A, lambda)    ! A <- A + lambda * I

         anorm = mat1norm(A)
         call dgetrf(n, n, A, n, ipiv, info)

         if (info /= 0) then
            ! Singular even with regularization; zero step, ramp lambda
            dX       = 0.0_pr
            cond_est = huge(1.0_pr)
            lambda   = min(lambda * s%lambda_up, jscale * s%lambda_max)
            return
         end if

         call dgecon('1', n, A, n, anorm, rcond, work, iwork, info)
         cond_est = 1.0_pr / max(rcond, tiny(rcond))

         call dgetrs('N', n, 1, A, n, ipiv, rhs, n, info)
         dX = rhs
      end if

   end subroutine solve_lm


   ! ===========================================================================
   ! PRIVATE: jacobian_scale
   !
   ! Returns  ||J||_F^2 / n  (mean squared entry of the Jacobian).
   !
   ! This quantity has the same physical dimensions as the eigenvalues of J^T J,
   ! so multiplying it by the dimensionless lambda_min/lambda_max gives a
   ! dimensionally consistent shift for the LM regularization.
   !
   ! It is always strictly positive (guarded by tiny) and is well-defined even
   ! when all diagonal entries of J are zero, unlike the previous approach that
   ! used the geometric mean of |diag(J)| (which collapsed to ~0 in that case).
   !
   ! Arguments:
   !   J  [in] – Jacobian matrix, shape (n,n)
   !
   ! Returns:
   !   ||J||_F^2 / n  where ||J||_F^2 = sum_{i,j} J_{ij}^2
   ! ===========================================================================
   pure real(pr) function jacobian_scale(J)
      real(pr), intent(in) :: J(:,:)   !! Jacobian matrix, shape (n,n)
      integer  :: n                    !! Problem size (number of rows)
      real(pr) :: frob2                !! Frobenius norm squared: ||J||_F^2

      n     = size(J, 1)
      frob2 = sum(J**2)
      jacobian_scale = max(frob2 / real(n, pr), tiny(1.0_pr))
   end function jacobian_scale


   ! ===========================================================================
   ! PRIVATE: pack_result
   !
   ! Writes all diagnostic scalars into the internal result accumulator (res)
   ! and optionally copies it into the caller's result variable (res_out).
   ! Also copies the residual history buffer if save_history is active.
   !
   ! Arguments:
   !   res      [inout]     – internal accumulator; fields updated on exit
   !   res_out  [out, opt]  – caller's result variable; receives a copy of res
   !   s        [in]        – solver settings (needed for save_history flag)
   !   hist_buf [in, opt]   – scratch buffer containing ||F||_inf per iteration
   !   its      [in]        – number of iterations completed
   !   f_norm   [in]        – final ||F||_inf
   !   dx_norm  [in]        – final ||dX||_inf
   !   lambda   [in]        – final LM lambda
   !   cond_est [in]        – final condition number estimate
   ! ===========================================================================
   subroutine pack_result(res, res_out, s, hist_buf, its, &
      f_norm, dx_norm, lambda, cond_est)
      type(newton_result),   intent(inout)          :: res
      type(newton_result),   intent(out),  optional :: res_out
      type(newton_settings), intent(in)             :: s
      real(pr),              intent(in),   optional :: hist_buf(:)
      integer,               intent(in)             :: its       !! Iterations done
      real(pr),              intent(in)             :: f_norm    !! ||F||_inf at exit
      real(pr),              intent(in)             :: dx_norm   !! ||dX||_inf at exit
      real(pr),              intent(in)             :: lambda    !! LM lambda at exit
      real(pr),              intent(in)             :: cond_est  !! cond(J) at exit

      res%iterations = its
      res%f_norm     = f_norm
      res%dx_norm    = dx_norm
      res%lambda     = lambda
      res%cond_est   = cond_est

      if (present(res_out)) then
         res_out = res
         if (s%save_history .and. present(hist_buf) .and. its >= 0) then
            if (allocated(res_out%f_history)) deallocate(res_out%f_history)
            allocate(res_out%f_history(its + 1))
            res_out%f_history = hist_buf(1:its + 1)
         end if
      end if
   end subroutine pack_result


   ! ===========================================================================
   ! PUBLIC: newton_status_msg
   !
   ! Returns a human-readable string for a given NEWTON_* status code.
   !
   ! Arguments:
   !   status [in] – one of the NEWTON_* integer constants
   !
   ! Returns:
   !   48-character string describing the exit condition
   ! ===========================================================================
   function newton_status_msg(status) result(msg)
      integer, intent(in) :: status   ! One of the NEWTON_* constants
      character(len=48)   :: msg      ! Human-readable description

      select case (status)
       case (NEWTON_SUCCESS);          msg = "Converged"
       case (NEWTON_MAX_ITS);          msg = "Maximum iterations reached"
       case (NEWTON_NAN_INF);          msg = "NaN / Inf detected in residual"
       case (NEWTON_SINGULAR);         msg = "Singular Jacobian"
       case (NEWTON_STAGNATION);       msg = "Stagnation: no progress detected"
       case (NEWTON_LINE_SEARCH_FAIL); msg = "Line search failed"
       case default;                   msg = "Unknown status code"
      end select
   end function newton_status_msg


   ! ===========================================================================
   ! Private utility functions
   ! ===========================================================================

   ! norm_inf: infinity norm of a real vector (max absolute value)
   pure real(pr) function norm_inf(v)
      real(pr), intent(in) :: v(:)   ! Input vector, any length
      norm_inf = maxval(abs(v))
   end function norm_inf

   ! mat1norm: 1-norm of a matrix = max column sum of |entries|
   ! Required as the anorm argument to LAPACK dgecon when called with '1'.
   pure real(pr) function mat1norm(A)
      real(pr), intent(in) :: A(:,:)  ! Input matrix, shape (m,n)
      integer :: j                    ! Column index
      mat1norm = 0.0_pr
      do j = 1, size(A, 2)
         mat1norm = max(mat1norm, sum(abs(A(:, j))))
      end do
   end function mat1norm

   ! add_diagonal: adds a scalar to every diagonal entry of A in place
   ! Used to form  A + lambda * I  for the LM normal equations.
   subroutine add_diagonal(A, val)
      real(pr), intent(inout) :: A(:,:)  ! Square matrix modified in place
      real(pr), intent(in)    :: val     ! Scalar added to each A(i,i)
      integer :: i                       ! Diagonal index
      do i = 1, min(size(A,1), size(A,2))
         A(i,i) = A(i,i) + val
      end do
   end subroutine add_diagonal

   ! has_nan_inf: returns .true. if any element of v is NaN or +/-Inf
   pure logical function has_nan_inf(v)
      real(pr), intent(in) :: v(:)   ! Vector to check
      has_nan_inf = any(isnan(v)) .or. any(abs(v) > huge(1.0_pr))
   end function has_nan_inf

   ! print_header: formatted column header for verbosity >= 1
   subroutine print_header(n)
      integer, intent(in) :: n   ! Problem size (printed in header)
      write(*,'(A)')        "=================================================================="
      write(*,'(A,I0,A)')  "  Newton-LM solver  |  n = ", n, " unknowns"
      write(*,'(A)')        "=================================================================="
      write(*,'(A6,4A14)') "iter", "||F||_inf", "||dX||_inf", "step t", "lambda"
      write(*,'(A)')        "------------------------------------------------------------------"
   end subroutine print_header

   ! print_summary: one-line convergence report printed at exit
   subroutine print_summary(res)
      type(newton_result), intent(in) :: res   ! Final result to summarise
      write(*,'(A)')        "------------------------------------------------------------------"
      write(*,'(A,I6)')     "  Iterations : ", res%iterations
      write(*,'(A,ES12.4)') "  ||F||_inf  : ", res%f_norm
      write(*,'(A,ES12.4)') "  ||dX||_inf : ", res%dx_norm
      write(*,'(A,ES12.4)') "  cond(J)    : ", res%cond_est
      write(*,'(A,ES12.4)') "  lambda     : ", res%lambda
      write(*,'(A,I2,2X,A)') "  Status     : ", res%status, &
         trim(newton_status_msg(res%status))
      write(*,'(A)')        "=================================================================="
   end subroutine print_summary

end module yaeos__newton_solver
