module marq_mod
      implicit none
      logical, parameter :: dbg = .false.

      !Fabian: Declare visibility of subroutines in this module for other modules
      private
      public :: mrqmin

      contains


!> Routine for non linear least-squares fitting by the Marquardt-Levenberg method.
!! The current implementation is based on Numerical Recipies pp. 526ff.
!! This subroutine is only called by the parent subroutine fit and and itself only calls funcs as a subroutine, which is the interface to the user-specific model routines.
!!
!! @author  Fabian Fritsch
!! @version 0.2
!! @date 15.03.2022
!!
!! @todo Implement different versions of the Marquardt-Levenberg method, e.g. MINPACK.

      subroutine mrqmin(par,npar,ma,mfit,rms,micit,set)
      use dim_parameter, only: log_convergence
      implicit none
! Input variables (not changed within this subroutine).
      integer npar              !< number of parameters
      integer mfit              !< number of active parameters
      integer ma(npar)             !< array that contains info if a parameter is active or not
      integer micit             !< number of microiterations / optimization cycles for the Marquardt Levenberg algorithm
      integer set !number of current set
! Input/output variables (changed/updated within this subroutine)
      double precision par(npar) !< parameter vector (on input: guess for parameters, on output: optimized parameters)
! Output variables

      double precision rms      !< root mean square error for the optimized parameter set

! Internal variables (used at most by subordinated subroutines)
!> @param alpha weighted approximation to the Hesse matrix (wt**2 * J^T * J)
!> @param covar damped alpha, ie. (wt**2 * J^T * J + alamda * 1)
      double precision alpha(mfit,mfit) !< weighted approximation to the Hesse matrix (wt**2 * J^T * J) / curvature matrix
      double precision covar(mfit,mfit) !< damped alpha, ie. (wt**2 * J^T * J + alamda * diag(J^T * J))
      double precision da(mfit)       !< J^T * (difference between y and fitted y)
      double precision beta(mfit)     !< @param[in] beta J^T * (difference between y and fitted y)

      logical skip              !< logical: true, if a parameter set must be skipped

      double precision chisq        !< chi-squared error (current parameter set)
      double precision ochisq       !< chi-squared error (previous best parameter set)
      double precision alamda       !< Marquardt-Levenberg parameter

      double precision atry(npar)     !< work array for temporal storage of the changed parameter vector
      double precision asave(npar)    !< work array for temporal storage of the changed parameter vector

      double precision trust        !< trust region parameter
      double precision ilamda       !< initial value for the Marquardt Levenberg parameter alamda

      logical quit

! Internal variables

      integer i,j,lauf          !< iteration variables
      double precision incr_alamda !< increase factor of alamda
      double precision decr_alamda !< decrease facotr of alamda

      character(len=80) file_name,fmt
      integer id_plot
      integer counter
      integer rejected

! Lapack variables (for details,see lapack documentation for the called subroutine)

      integer info
      double precision work(mfit)
      integer*8 lwork !important that its 64bit!
      integer*8 ipiv(mfit)

!> Open logging files
      if (log_convergence) then
         write (file_name,"('.conv_marq_set',i0,'.dat')") set
         id_plot=6666+set
         open(id_plot,file=trim(file_name),action='write')
         write (id_plot,69)
 69      format ('#Levenberg-Marquardt',/,
     +        '#it    = iteration number',/,
     +        '#nf    = number of function evaluations',/,
     +        '#stepl = step length used',/,
     +        '#tstep = norm of the displacement (total step)',/,
     +        '#projg = norm of the projected gradient',/,
     +        '#f     = function value')
         write (id_plot,123)
 123     format ('#',3x,'it',3x,'nf',3x,'stepl',4x,'pstep',5x,
     $        'projg',8x,'f')
      endif
!> Initialize skip, Marquardt parameter, error variables and parameter work arrays

      skip=.false.
      !ilamda=0.1d0  !Initial Marquardt parameter
      ilamda=0.05d0  !Initial Marquardt parameter
      alamda=ilamda !Current Marquardt parameter
      rms=1.d6
      atry(1:npar)=par(1:npar)
      asave(1:npar)=par(1:npar)

      !Fabian: Adjusted values (TO-DO: Make this Input-Cards)
      trust=2.0d+1
      incr_alamda=1.5d0
      decr_alamda=1.d0/1.5d0

      counter=0
      rejected=0
      alpha=0.d0
      beta=0.d0
      chisq=0.d0

!> Calculate RMS error for the given start parameters

      call mrqcof(par,npar,ma,mfit,
     $     alpha,beta,chisq,skip)
      ochisq=chisq

      !Check initial Hessian for (trivial) rank-deficiency
      !This check reveals for which parameter no data is available, i.e no optimization will occur although the parameter is active
      !Note: In general, rank deficiency might lead to ambiguous results when solving the normal equations. Here this will not occur due to the regularization in the LM-Algorithm
      if (dbg) then
         do i=1,mfit
            if(all(abs(alpha(i,:)).lt.1E-16)) then
               write(6,*)
     $          'Warning: Rank deficiency of J^T*J for active param',i            
            endif
         enddo
      endif

!> Termination if micit=0, i.e. only genetic algorithm and no LM optimization
      if (micit.eq.0) then
         rms=dsqrt(chisq)
         if (log_convergence) close(id_plot)
         return
      endif

!> Write warning and return if error for initial parameters is large
      if (skip) then
        write(6,*) 'WARNING: initial parameter set skipped'
!        call flush
!        flush(6)
        rms=1.e6
        if (log_convergence) close(id_plot)
        return
      endif

!-------------------------------------------------------------------------

!> Start optimization of parameters using the Marquardt-Levenberg algorihm

      do lauf=1,micit

!-------------------------------------------------------------------------o

         !> Scale approximated Hessian (alpha = wt**2 * J^T * J) if necessary
         call mrqrescale(mfit,alpha)

         !-------------------------------------------------------------------------

         !> Calculate covariance matrix and "gradient": wt**2 * J^T * J + alamda * diag(J^T * J) == alpha + alamda * diag(J^T * J)

         !Copy alpha to covar and beta to da
         covar(1:mfit,1:mfit)=alpha(1:mfit,1:mfit)
         da(1:mfit)=beta(1:mfit)
         !Adjust diagonal elements of covar by alamda
         do i=1,mfit
            covar(i,i)=covar(i,i)*(1.d0 + alamda)
            if (dabs(covar(i,i)).lt.1.d-12) then
               covar(i,i)=1.d-8
            endif
         enddo

         !-------------------------------------------------------------------------        

         !> Solve set of equations, i.e. covar * vec = da

         !Lapack version (on output: da contains solution the equation, called shift vector)
         ipiv=0
         lwork = max(1,mfit)
         call dsysv('U',mfit,1,covar,mfit,ipiv,da,mfit, !qr decomposition 
     $        work,lwork,info)
!         call dposv('U',mfit,1,covar,mfit,da,mfit, !cholesky decomposition 
!        $        info)
         !-------------------------------------------------------------------------        

         !> Calculate trust region of the shift vector, if necessary rescale the entire shift vector

         call mrqtrustregion(trust,npar,ma,par,mfit,da)

         !-------------------------------------------------------------------------

         !> Check if the new (rescaled) parameters are better then the previous best (micro)iteration

         !Calculate the new parameters and write them on atry
         j=0
         do i=1,npar
!Nicole: added flexible value of ma
            if (ma(i).ge.1) then
!            if (abs(ma(i)).eq.1) then
               j=j+1
               atry(i)=par(i)+da(j)
            endif
         enddo

         !Calculate RMS with the new parameters atry
         call mrqcof(atry,npar,ma,mfit,
     $        covar,da,chisq,skip)

         ! Write warning, if mrqcof (more precisely funcs within mrqcof) yields the skip message
         if (skip) then
            write(6,*) 'WARNING: parameter set skipped'
!            call flush
!            flush(6)
            rms=1.e6
            if (log_convergence) close(id_plot)
            return
         endif

         !Compare the new RMS with the RMS of the previous best parameters, if yes: save the parameters
         if(chisq.lt.ochisq) then
            counter=counter+1 !number of accepted steps
            asave(1:npar)=atry(1:npar)
            !Write logging information
            if(log_convergence) then
               write(id_plot,124) counter,lauf,1.d0,sum((atry-par)**2),
     $              sum(da**2),chisq
 124           format(1x,2(1x,i4),1p,2(2x,e8.1),1p,2(1x,e10.3))
            endif
         endif

         !-------------------------------------------------------------------------

         !Perform convergence / error checks
         quit=.false.
         call mrqconvergence(lauf,micit,npar,par,mfit,da,asave,
     $        chisq,ochisq,rms,rejected,quit)
         if(quit) then
            if (log_convergence) close(id_plot)
            return
         endif

         !-------------------------------------------------------------------------

         !Increase counter of consecutive rejected steps by one
         rejected=rejected+1

         !Adjust the marquardt parameter (alamda) for the next iteration
         !If chisq has been reduced: Update alpha, beta, par and ochisq
         if (chisq.lt.ochisq) then
            rejected=0 !reset counter of consecutive rejected step
            alamda=alamda*decr_alamda
            if(alamda.lt.1E-8) alamda=1E-8
            alpha(1:mfit,1:mfit)=covar(1:mfit,1:mfit)
            beta(1:mfit)=da(1:mfit)
            par(1:npar)=atry(1:npar)
            ochisq=chisq
         else
            alamda=alamda*incr_alamda
            !If after a certain number of iterations in which the rms is not reduced or convergence occurs,
            !alamda takes a certain value, then take the result of this iteration as a new input guess !
            if (alamda.gt.1.d5) then
               write(6,*) 'Warning: Large alamda, try new parameters'
               alamda=ilamda
               par(1:npar)=atry(1:npar)
               ochisq=chisq
            endif
         endif

      enddo

      if(log_convergence) close(id_plot)

      end subroutine


c###############################################################

!> Routine for calculating the residuals, gradients, approximated Hessian for the Marquardt-Levenberg-Algorithm

      subroutine mrqcof(par,npar,ma,mfit,
     $                  alpha,beta,chisq,skip)

      use dim_parameter, only: ntot,qn,numdatpt,nstat,hybrid

      use data_module, only: x1_m,x2_m,y_m,wt_m,ny_m
      use funcs_mod,only: funcs
      implicit none

! Input variables (not changed within this subroutine).
!      double precision q(*),x1(*),x2(*),y(*),ny(*),wt(*) !< coordinates,data and weights
      integer npar              !< number of parameters
      double precision par(npar) !< parameter vector (on input: guess for parameters, on output: optimized parameters)
      integer mfit              !< number of active parameters
      integer ma(npar)             !< array that contains info if a parameter is active or not

! Output variables
      double precision alpha(mfit,mfit) !< weighted approximation to the Hesse matrix (wt**2 * J^T * J) / curvature matrix
      double precision beta(mfit)       !< weighted J^T * (difference between y and fitted y)
      double precision chisq         !< chisq error
      logical skip              !< logical: true, if a parameter set must be skipped

! Internal variables (used at most by subordinated subroutines)
      double precision ymod(ntot) !< fitted datapoints (for one geometry)
      double precision dy(ntot)   !< difference between ab-initio and fitted datapoints
      double precision dyda(ntot,npar) !gradient of datapoints (for one geometry) with respect to the parameters

! Internal variables
      integer i,j,k,l,m,n !< iteration variables
      integer nloop
      integer active(mfit)

!-------------------------------------------------------------------------

!> Initialize skip, alpha, beta and chisq

      skip=.false.
      alpha(1:mfit,1:mfit)=0.d0
      beta(1:mfit)=0.d0
      chisq=0.d0

      nloop=ntot
      !if(.not.hybrid) nloop=nstat ! commented out by JP on 08.10.2025
      ! In fitting dipole nstat is meaningless, I always go with ntot
      ! and send to zero the unused parts of y_m and ymod
      j=0
      do i=1,npar
         if(ma(i).ge.1) then 
            j=j+1
            active(j)=i
         endif
      enddo 

      do i=1,numdatpt
    
         call funcs(i,par,ymod,dyda,npar,ma,skip)
         !write(58,*) "ymod ",i,nloop,ymod(1:ntot)
         if (skip) return

         !Idea: Since the quantities dyda,dy and wt_m are rather small, one might consider scaling them
         !Idea: and then rescale the final quantities alpha,beta,chisq accordingly
         !Idea: Scale dyda,dy and wt_m by 1D+5; final rescale of alpha,beta and chisq by 1D-10

         do n=1,nloop
            dy(n)=(y_m(n,i)-ymod(n))*wt_m(n,i)
            dyda(n,:) = dyda(n,:)*wt_m(n,i) 
            !write(58,"(A3,2I0,3ES18.9)")"dy ",n,i,ymod(n),y_m(n,i),dy(n)
            do l=1,mfit 
                do m=1,l 
                   alpha(m,l)=alpha(m,l)+
     &                (dyda(n,active(l))*dyda(n,active(m)))!*wt_m(n,i) !JP
                enddo 
                beta(l) = beta(l) + (dy(n) * dyda(n,active(l)))
            enddo 
            chisq = chisq + (dy(n)**2)
        enddo 
      enddo 
      
!-------------------------------------------------------------------------

      !Fill in missing parts of the symmetric matrix alpha
      do i=2,mfit
        do j=1,i-1
          alpha(i,j)=alpha(j,i)
        enddo
      enddo

      end subroutine

c###############################################################

!> Scale approximated Hessian (alpha = wt**2 * J^T * J) if necessary
      subroutine mrqrescale(mfit,alpha)

      implicit none

!> Input variables
      integer mfit

!> Input/output variables
      double precision alpha(mfit,mfit)

!> Internal variables
      double precision dum !Maik changed to double from int
      integer i,j

!> Find largest value of the approximated Hessian alpha
      dum=0.d0
      do i=1,mfit
         do j=1,mfit
            if (abs(alpha(i,j)).gt.dum) dum=abs(alpha(i,j)) !find largest value of the approximated Hessian
         enddo
      enddo

!> Rescale approximated Hessian if largest value is greater then a threshold (hardcoded)
      if (dum.gt.1.d30) then
         dum=1.d3/dum
         write(6,'(''Warning: Hessian scaled by'',d12.3)') dum
         write(6,*)
         do i=1,mfit
            do j=1,mfit
               alpha(i,j)=alpha(i,j)*dum
            enddo
         enddo
      endif

      end subroutine

c###############################################################

      subroutine mrqtrustregion(trust,npar,ma,par,mfit,da)

      implicit none

!> Input variables
      double precision trust
      integer npar
      integer ma(npar)
      double precision par(npar)
      integer mfit

!> Input/output variables
      double precision da(mfit)

!> Internal variables
      integer i,j
      double precision dum

      !Init
      dum=0.d0
      j=0

      !Find the largest relative (magnitude of gradient / magnitude of parameter = da / par ) shift of a parameter
      !Explanation: For parameters of high scale, their gradients are also of high scale even if the relative change is small
      do i=1,npar
! Nicole: values of ma (active parameter) changed
         if (ma(i).ge.1) then
            j=j+1
            if (abs(par(i)).gt.1.d-4) then
               if (abs(da(j)/par(i)).gt.dum) then
                  dum=abs(da(j)/par(i))
               endif
            endif
         endif
      enddo
      !If maximum relative shift exceeds a threshold, scale shift vector
      if (dum.gt.trust) then
         dum=trust/dum
         j=0
         do i=1,npar
! Nicole: values of ma (active parameter) changed
            if (ma(i).ge.1) then
               j=j+1
               da(j)=da(j)*dum
            endif
         enddo
      endif

      end subroutine

c###############################################################

!> Perform convergence and error checks
      subroutine mrqconvergence(lauf,micit,npar,par,mfit,da,asave,
     $     chisq,ochisq,rms,rejected,quit)

      implicit none

!> Input variables
      integer lauf
      integer micit
      integer npar
      integer mfit
      double precision da(mfit)
      double precision asave(npar)
      double precision chisq
      double precision ochisq
      integer rejected !number of consecutive rejected steps


!> Input/Output variables
      double precision par(npar)
      double precision rms

!> Output variable
      logical quit

!> Internal variables
      integer i
      double precision dum
      double precision check


      quit=.false.

      !Negative termination, if rms is too large, quit this iteration
      if (chisq.gt.1.d3) then
         write(6,*) 'chi^2 unreasonable!', chisq, lauf
         rms=dsqrt(ochisq)
         par(1:npar)=asave(1:npar)
         quit=.true.
         return
      endif

!     Negative-neutral termination, if maximum number of consecutive not accepted microiterations reached
      if (rejected.ge.25) then
         write(6,*) 'Warning: 25 consecutive non accepted steps!'
         if (chisq.lt.ochisq) rms=dsqrt(chisq)
         if (chisq.ge.ochisq) rms=dsqrt(ochisq)
         par(1:npar)=asave(1:npar)
         quit=.true.
         return
      endif

!     Neutral termination, if number of maximum microiterations reached, quit the LM algorithm
      if (lauf.eq.micit) then
         if (chisq.lt.ochisq) rms=dsqrt(chisq)
         if (chisq.ge.ochisq) rms=dsqrt(ochisq)
         par(1:npar)=asave(1:npar)
         write(6,'(''Warning: iterations exceeded: '',I0)') lauf
         quit=.true.
         return
      endif

!     Neutral-positive termination, if parameter changes are small
      check=0.d0
      do i=1,mfit
         check=check+da(i)**2
      enddo
      
      check=dsqrt(check)        !root mean square of gradient
      if (check.lt.1.d-15) then
         write(6,*) 'change of parameters converged', check, lauf
         rms=dsqrt(chisq)
         par(1:npar)=asave(1:npar)
         quit=.true.
         return
      endif

      !Positive termination, if difference between the previous lowest/optimal chisq (that is ochisq) & the current chisq is small
      dum=max(chisq,ochisq)/min(chisq,ochisq) - 1.d0
      if ((dum.lt.1.d-5).and.(lauf.gt.1)) then !change of chi^2 < 0.01%
         write(6,*) 'change of chi^2 converged', dum, lauf
!         call flush !Fabian 15.03.2022: Not sure, why this is called
!         flush(6)
         rms=dsqrt(chisq)
         par(1:npar)=asave(1:npar)
         quit=.true.
         return
      endif


      end subroutine

      end module marq_mod