! Module defines CG type to store CG algorithm parameters
module MW_cgaim
   use MW_kinds, only: wp, sp
   implicit none
   private

   ! Public type
   ! -----------
   public :: MW_cgaim_t

   ! Public subroutines
   ! ------------------
   public :: set_max_iterations
   public :: set_tolerance
   public :: reset_statistics
!  public :: allocate_arrays
   public :: void_type
   public :: print_type
   public :: print_statistics
   public :: solve

   type MW_cgaim_t

      ! Convergence parameters
      integer :: max_iterations !< maximum number of iterations
      real(wp) :: tol           !< tolerance on residual norm
      ! Statistics
      integer  :: last_iteration_count    !< Number of iteration for the last CG
      real(wp) :: last_residual
      real(wp) :: last_residual_tol
      real(sp) :: total_iteration_count   !< Sum of all iteration count

      ! work arrays
!     integer :: n                    !< dimension of the work arrays
!     real(wp), allocatable :: b(:)   !< right hand side
!     real(wp), allocatable :: res(:) !< residual
!     real(wp), allocatable :: p(:)   !< conjugate direction
!     real(wp), allocatable :: Ap(:)  !< A times p
      real(wp), allocatable :: xi(:)  ! gradient

   end type MW_cgaim_t

contains

   ! ================================================================================
   ! Set maximum number of iterations
   subroutine set_max_iterations(this, max_iterations)
      implicit none
      type(MW_cgaim_t), intent(inout) :: this
      integer, intent(in) :: max_iterations
      this%max_iterations = max_iterations
   end subroutine set_max_iterations

   ! ================================================================================
   ! Set tolerance parameter
   subroutine set_tolerance(this, tol)
      implicit none
      type(MW_cgaim_t), intent(inout) :: this
      real(wp), intent(in) :: tol
      this%tol = tol
   end subroutine set_tolerance

   ! ================================================================================
   ! reset statistic counters
   subroutine reset_statistics(this)
      implicit none
      type(MW_cgaim_t), intent(inout) :: this
      this%last_iteration_count = 0
      this%last_residual = 0.0_wp
      this%last_residual_tol = 0.0_wp
      this%total_iteration_count = 0
   end subroutine reset_statistics

   !================================================================================
!  ! Allocate data arrays
!  subroutine allocate_arrays(this, n)
!     use MW_errors, only: MW_errors_allocate_error => allocate_error
!     implicit none

!     ! Parameters inout
!     ! ----------------
!     type(MW_cgaim_t), intent(inout) :: this

!     ! Parameters in
!     ! -------------
!     integer, intent(in) :: n !< size of the data arrays

!     ! Local
!     ! -----
!     integer :: i
!     integer :: ierr

!     this%n = n
!     allocate(this%b(n), this%res(n), this%p(n), this%Ap(n), stat=ierr)
!     if (ierr /= 0) then
!        call MW_errors_allocate_error("define_type", "cg.f90", ierr)
!     end if

!     do i = 1, n
!        this%b(i) = 0.0_wp
!        this%res(i) = 0.0_wp
!        this%p(i) = 0.0_wp
!        this%Ap(i) = 0.0_wp
!     end do

!  end subroutine allocate_arrays

   !================================================================================
   ! Void the data structure
   subroutine void_type(this)
      use MW_errors, only: MW_errors_deallocate_error => deallocate_error
      implicit none

      ! Parameters inout
      ! ----------------
      type(MW_cgaim_t), intent(inout) :: this

      ! Local
      ! -----
      this%max_iterations = 0
      this%tol = 0.0_wp
      this%last_iteration_count = 0
      this%total_iteration_count = 0.0_sp
      this%last_residual_tol = 0.0_wp
!     this%n = 0

!     if (allocated(this%b)) then
!        deallocate(this%b, this%res, this%p, this%Ap, stat=ierr)
!        if (ierr /= 0) then
!           call MW_errors_deallocate_error("void_type", "cgaim.f90", ierr)
!        end if
!     end if
   end subroutine void_type

   !================================================================================
   ! Print the data structure parameters
   subroutine print_type(this, ounit)
      implicit none

      ! Parameters inout
      ! ----------------
      type(MW_cgaim_t), intent(in) :: this
      integer,       intent(in) :: ounit

      write(ounit, '("|cgaim| maximum number of iterations: ",i12)') this%max_iterations
      write(ounit, '("|cgaim| Residual error:               ",es12.5)') this%tol
   end subroutine print_type

   !================================================================================
   ! Print the data structure parameters
   subroutine print_statistics(this, ounit)
      implicit none

      ! Parameters inout
      ! ----------------
      type(MW_cgaim_t), intent(in) :: this
      integer,       intent(in) :: ounit

      write(ounit, '("|cgaim| number of iterations: ",i12)') this%last_iteration_count
      write(ounit, '("|cgaim| Residual difference:        ",es12.5)') this%last_residual
   end subroutine print_statistics

!subroutine solve(system,p,num,itmax,tolaim,ftolaim)
subroutine solve(this,system,p,num)
   use MW_system, only: MW_system_t
   use MW_daim, only: &
         MW_daim_repulsionenergy => repulsionenergy, &
         MW_daim_diffrepulsionenergy => diffrepulsionenergy, &
         MW_daim_selfenergy => selfenergy, &
         MW_daim_diffselfenergy => diffselfenergy
   use MW_kinds, only: wp
   use MW_errors, only: MW_errors_runtime_error => runtime_error
   use MW_stdio, only: MW_stderr

   implicit none

   ! Parameters in
   ! -------------
   type(MW_system_t), intent(inout) :: system
   type(MW_cgaim_t), intent(inout) :: this
   integer :: num        ! number of ions

   ! Parameters out 
   ! -------------
   real(wp),  intent(inout) :: p(:)

   ! LOCAL 
   ! -------------
   integer :: its,iter,xyz_now
   integer :: itmax     ! max number of iterations
   real(wp) :: tolaim,ftolaim    ! tolerances
   real(wp), dimension(4*num) :: g,hhh,pairxi,selfxi,xi
   real(wp) :: fp,gg,dgg,gam,fret
   real(wp) :: pairenergy,selfenergy
   real(wp), dimension(num) :: locdeltas
   real(wp), dimension(num,3) :: locepsilons
   logical :: cg_converged

   pairenergy=0.d0
   selfenergy=0.d0
   itmax=this%max_iterations
   tolaim=this%tol
   ftolaim=this%tol
   cg_converged=.false.

   xyz_now=system%xyz_ions_step(1)

   locdeltas(1:num)=p(1:num)
   locepsilons(1:num,1)=p(num+1:2*num)
   locepsilons(1:num,2)=p(2*num+1:3*num)
   locepsilons(1:num,3)=p(3*num+1:4*num)

   pairxi=0.d0
   selfxi=0.d0

! I need to implement the derivatives there
   call MW_daim_diffrepulsionenergy(system%num_PBC, system%localwork, system%daim, system%box, &
               system%ions, system%xyz_ions(:,:,xyz_now), &
               system%electrodes, system%xyz_atoms,locdeltas, &
               locepsilons,pairenergy,pairxi)

   call MW_daim_diffselfenergy(system%num_ion_types,system%ions,locdeltas,locepsilons,selfenergy,selfxi)

   fp=pairenergy+selfenergy
   xi=pairxi+selfxi

   g=-xi
   hhh=g
   xi=hhh

   iter=0
   do its=1,itmax
      iter=its
      call linminaim(system,num,tolaim,p,xi,fret)

!     if(2.*abs(fret-fp).le.ftolaim*(abs(fret)+abs(fp)+EPS))then
      if(abs(fret-fp).le.ftolaim)then
         this%last_residual=abs(fret-fp)
         this%last_iteration_count=iter
         cg_converged=.true.
         return
      endif

      locdeltas(1:num)=p(1:num)
      locepsilons(1:num,1)=p(num+1:2*num)
      locepsilons(1:num,2)=p(2*num+1:3*num)
      locepsilons(1:num,3)=p(3*num+1:4*num)

      pairenergy=0.0_wp
      selfenergy=0.0_wp
      pairxi=0.0_wp
      selfxi=0.0_wp
      call MW_daim_diffrepulsionenergy(system%num_PBC, system%localwork, system%daim, &
               system%box, &
               system%ions, system%xyz_ions(:,:,xyz_now), &
               system%electrodes, system%xyz_atoms,locdeltas, &
               locepsilons,pairenergy,pairxi)

      call MW_daim_diffselfenergy(system%num_ion_types,system%ions, &
                locdeltas,locepsilons,selfenergy,selfxi)

      fp=pairenergy+selfenergy
      xi=pairxi+selfxi
      gg=dot_product(g,g)
      dgg=dot_product(xi,g+xi)
      if(gg.eq.0)return
      gam=dgg/gg
      g=-xi
      hhh=g+gam*hhh
      xi=hhh
   enddo
   if(.not.cg_converged)then
      this%last_residual=abs(fret-fp)
      this%last_iteration_count=iter
      call print_statistics(this, MW_stderr)
      call MW_errors_runtime_error("solve","cgaim.f90", "CGaim failed to converge")
   endif
return

end subroutine solve

subroutine linminaim(system,num,tolaim,p,xi,fret)
   use MW_system, only: MW_system_t

   implicit none

   ! Parameters in
   ! -------------
   type(MW_system_t), intent(inout) :: system
   integer :: num        ! number of ions
   real(wp) :: tolaim    ! tolerances

   ! Parameters out 
   ! -------------
   real(wp),  intent(inout) :: p(:)
   real(wp),  intent(inout) :: xi(:)
   real(wp), intent(out) :: fret

   ! LOCAL 
   ! -------------
   real(wp) :: axx,bxx,xx,xmin !dbrent!,F1DIMAIM
   real(wp), dimension(2) :: resultdbrent  
   real(wp), dimension(4*num) :: pcomaim,xicomaim 
!EXTERNAL F1DIMAIM,DF1DIMAIM

   pcomaim=p
   xicomaim=xi

   axx=0.000_wp
   xx=0.001_wp
   bxx=0.002_wp
   
   call mnbrakaim(system,num,axx,xx,bxx,pcomaim,xicomaim) 
   resultdbrent=dbrent(system,num,axx,xx,bxx,tolaim,pcomaim,xicomaim)
   fret=resultdbrent(1)
   xmin=resultdbrent(2)
   xi=xmin*xi
   p=p+xi

return
end subroutine linminaim

subroutine mnbrakaim(system,num,axx,bxx,cx,pcomaim,xicomaim)
   use MW_system, only: MW_system_t

   implicit none

   ! Parameters in
   ! -------------
   type(MW_system_t), intent(inout) :: system
   integer, intent(in) :: num
   real(wp), intent(in) :: pcomaim(:)
   real(wp), intent(in) :: xicomaim(:) 
   ! Parameters out 
   ! -------------
   real(wp), intent(inout) :: axx,bxx,cx !F1DIMAIM

   ! LOCAL 
   ! -------------
   real(wp) :: dum,r,qq,u,ulim,fu
   real(wp) :: fa,fb,fc 
   real(wp), parameter :: GOLD=1.618034, GLIMIT=100.0_wp, TINY=1.E-20

   fa=f1dimaim(system,num,pcomaim,xicomaim,axx) 
   fb=f1dimaim(system,num,pcomaim,xicomaim,bxx) 


   if(fb.gt.fa)then
      dum=axx
      axx=bxx
      bxx=dum
      dum=fb
      fb=fa
      fa=dum
   endif
   cx=bxx+GOLD*(bxx-axx)
   fc=f1dimaim(system,num,pcomaim,xicomaim,cx) 
1  if(fb.ge.fc)then
      r=(bxx-axx)*(fb-fc)
      qq=(bxx-cx)*(fb-fa)
      u=bxx-((bxx-cx)*qq-(bxx-axx)*r)/(2.0_wp*sign(max(abs(qq-r),TINY),QQ-R))
      ulim=bxx+GLIMIT*(cx-bxx)
      if((bxx-u)*(u-cx).gt.0.0_wp)then
         fu=f1dimaim(system,num,pcomaim,xicomaim,u)
         if(fu.lt.fc)then
            axx=bxx
            fa=fb
            bxx=u
            fb=fu
            go to 1
         else if(fu.gt.fb)then
            cx=u
            fc=fu
            go to 1
         endif
         u=cx+GOLD*(cx-bxx)
         fu=f1dimaim(system,num,pcomaim,xicomaim,u)
      else if((cx-u)*(u-ulim).gt.0.0_wp)then
         fu=f1dimaim(system,num,pcomaim,xicomaim,u)
         if(fu.lt.fc)then
            bxx=cx
            cx=u
            u=cx+GOLD*(cx-bxx)
            fb=fc
            fc=fu
            fu=f1dimaim(system,num,pcomaim,xicomaim,u)
         endif
      else if((u-ulim)*(ulim-cx).ge.0.0_wp)then
         u=ulim
         fu=f1dimaim(system,num,pcomaim,xicomaim,u)
      else
        u=cx+GOLD*(cx-bxx)
        fu=f1dimaim(system,num,pcomaim,xicomaim,u)
      endif
      axx=bxx
      bxx=cx
      cx=u
      fa=fb
      fb=fc
      fc=fu
      go to 1
   end if
return
end subroutine mnbrakaim

function dbrent(system,num,ax,bx,cx,tol,pcomaim,xicomaim) result(resultdbrent)
   use MW_system, only: MW_system_t

   implicit none

   ! Parameters in
   ! -------------
   type(MW_system_t), intent(inout) :: system
   integer, intent(in) :: num
   real(wp), intent(in) :: tol
   real(wp), intent(in) :: ax,bx,cx
   real(wp), dimension(4*num), intent(in) :: pcomaim,xicomaim 

   ! Result
   ! -------------
   real(wp),dimension(2) ::resultdbrent 

   ! Local
   ! -------------
   integer :: iter
   real(wp) :: xminloc,tol1,tol2,olde,a,b,d,e,u,v,w,x,d1,d2, &
                       du,dv,dw,dx,fu,fv,fw,fx,u1,u2,xm
   real(wp),dimension(2) :: energyanddiff
!DOUBLE PRECISION :: F,DF
   integer, parameter :: ITMAX=100
   real(wp), parameter :: ZEPS=1.0d-10
   logical ok1,ok2

   a=min(ax,cx)
   b=max(ax,cx)
   v=bx
   w=v
   x=v
   d=0.0_wp
   e=0.0_wp
!  x=f1dimaim(system,num,pcomaim,xicomaim,x)  ! probably merged with the calculation of dx
   energyanddiff=df1dimaim(system,num,pcomaim,xicomaim,x) ! probably get everything together at once
   fx=energyanddiff(1)
   dx=energyanddiff(2)

   fv=fx
   fw=fx
   dv=dx
   dw=dx   

   do iter=1,ITMAX
      xm=0.5_wp*(a+b)
      tol1=tol*abs(x)+ZEPS
      tol2=2.0_wp*tol1
      if(abs(x-xm).le.(tol2-0.5_wp*(b-a))) exit
      if(abs(e).gt.tol1)then
         d1=2.0_wp*(b-a)
         d2=d1
         if(dw.ne.dx) d1=(w-x)*dx/(dx-dw)
         if(dv.ne.dx) d2=(v-x)*dx/(dx-dv)
         u1=x+d1
         u2=x+d2
         ok1=((a-u1)*(u1-b).gt.0.0_wp).and.(dx*d1.le.0.0_wp)
         ok2=((a-u2)*(u2-b).gt.0.0_wp).and.(dx*d2.le.0.0_wp)
         olde=e
         e=d
         if(.not.(ok1.or.ok2))then
            go to 1
         else if(ok1.and.ok2)then
            if(abs(d1).lt.abs(d2))then
               d=d1
            else
               d=d2
            endif
         else if(ok1)then
            d=d1
         else
            d=d2
         endif
         if(abs(d).gt.abs(0.5_wp*olde)) go to 1
         u=x+d
         if((u-a.lt.tol2).or.(b-u.lt.tol2)) d=sign(tol1,xm-x)   
         go to 2
      endif
1     if(dx.ge.0.0_wp)then
         e=a-x
      else
         e=b-x
      endif
      d=0.5_wp*e
2     if(abs(d).ge.tol1)then
         u=x+d
!        fu=f1dimaim(system,num,pcomaim,xicomaim,u)  
      else
         u=x+sign(tol1,d)
!        fu=f1dimaim(system,num,pcomaim,xicomaim,u)  
!        if(fu.gt.fx)exit
      endif
!     fu=f1dimaim(system,num,pcomaim,xicomaim,u)  
!     du=df1dimaim(system,num,pcomaim,xicomaim,u) ! probably get everything together at once
      energyanddiff=df1dimaim(system,num,pcomaim,xicomaim,u) ! probably get everything together at once
      fu=energyanddiff(1)
      du=energyanddiff(2)
      if(abs(d).lt.tol1)then
         if(fu.gt.fx)exit
      endif
      if(fu.le.fx)then
         if(u.ge.x)then
            a=x
         else
            b=x
         endif
         v=w
         fv=fw
         dv=dw
         w=x
         fw=fx
         dw=dx
         x=u
         fx=fu
         dx=du
      else
         if(u.lt.x)then
            a=u
         else
            b=u
         endif
         if((fu.le.fw).or.(w.eq.x))then
            v=w
            fv=fw
            dv=dw
            w=u
            fw=fu
            dw=du
         else if((fu.le.fv).or.(v.eq.x).or.(v.eq.w))then
            v=u
            fv=fu
            dv=du
         endif
      endif
   enddo
   xminloc=x
   resultdbrent(1)=fx
   resultdbrent(2)=xminloc
   return
end function dbrent

function f1dimaim(system,num,pcomaim,xicomaim,xxx) result(energytot)
   use MW_system, only: MW_system_t
   use MW_daim, only: &
         MW_daim_repulsionenergy => repulsionenergy, &
         MW_daim_selfenergy => selfenergy

   implicit none  

   ! Parameters in
   ! -------------
   type(MW_system_t), intent(inout) :: system
   integer :: num        ! number of ions
   real(wp), intent(in) :: pcomaim(:)
   real(wp), intent(in) :: xicomaim(:) 
   real(wp), intent(in) :: xxx

   ! Result 
   ! -------------
   real(wp) :: energytot

   ! LOCAL 
   ! ------------- 
   integer :: i,itype,iion,offset
   integer :: xyz_now
   real(wp) :: pairenergy,selfenergy
   real(wp), dimension(4*num) :: xt
   real(wp), dimension(num) :: locdeltas
   real(wp), dimension(num,3) :: locepsilons

   xyz_now=system%xyz_ions_step(1)

   xt=0.0_wp

   do itype=1,system%num_ion_types
      if(system%ions(itype)%selfdaimD>0.0_wp)then
         offset=system%ions(itype)%offset
         do i=1,system%ions(itype)%count
            iion=offset+i
            xt(iion)=pcomaim(iion)+xxx*xicomaim(iion) 
            xt(num+iion)=pcomaim(num+iion)+xxx*xicomaim(num+iion) 
            xt(2*num+iion)=pcomaim(2*num+iion)+xxx*xicomaim(2*num+iion) 
            xt(3*num+iion)=pcomaim(3*num+iion)+xxx*xicomaim(3*num+iion) 
         enddo
      endif
   enddo
  
   locdeltas(1:num)=xt(1:num)
   locepsilons(1:num,1)=xt(num+1:2*num)
   locepsilons(1:num,2)=xt(2*num+1:3*num)
   locepsilons(1:num,3)=xt(3*num+1:4*num)

   pairenergy=0.0_wp
   selfenergy=0.0_wp
   call MW_daim_repulsionenergy(system%num_PBC, system%localwork, system%daim, system%box, &
               system%ions, system%xyz_ions(:,:,xyz_now), &
               system%electrodes, system%xyz_atoms,locdeltas, &
               locepsilons,pairenergy)

   call MW_daim_selfenergy(system%num_ion_types,system%ions,locdeltas,locepsilons,selfenergy)

   energytot=pairenergy+selfenergy

return
end function f1dimaim

function df1dimaim(system,num,pcomaim,xicomaim,xxx) result(energyanddiff)
   use MW_system, only: MW_system_t
   use MW_daim, only: &
         MW_daim_diffrepulsionenergy => diffrepulsionenergy, &
         MW_daim_diffselfenergy => diffselfenergy

   implicit none  

   ! Parameters in
   ! -------------
   type(MW_system_t), intent(inout) :: system
   integer :: num        ! number of ions
   real(wp), intent(in) :: pcomaim(:)
   real(wp), intent(in) :: xicomaim(:) 
   real(wp), intent(in) :: xxx

   ! Result 
   ! -------------
   real(wp),dimension(2) :: energyanddiff

   ! LOCAL 
   ! ------------- 
   integer :: i,itype,iion,offset
   integer :: xyz_now
   real(wp) :: pairenergy,selfenergy,energytot
   real(wp), dimension(4*num) :: xt,locxi,pairxi,selfxi
   real(wp), dimension(num) :: locdeltas
   real(wp), dimension(num,3) :: locepsilons

   xyz_now=system%xyz_ions_step(1)

   xt=0.0_wp

   do itype=1,system%num_ion_types
      if(system%ions(itype)%selfdaimD>0.0_wp)then
         offset=system%ions(itype)%offset
         do i=1,system%ions(itype)%count
            iion=offset+i
            xt(iion)=pcomaim(iion)+xxx*xicomaim(iion) 
            xt(num+iion)=pcomaim(num+iion)+xxx*xicomaim(num+iion) 
            xt(2*num+iion)=pcomaim(2*num+iion)+xxx*xicomaim(2*num+iion) 
            xt(3*num+iion)=pcomaim(3*num+iion)+xxx*xicomaim(3*num+iion) 
         enddo
      endif
   enddo
   locdeltas(1:num)=xt(1:num)
   locepsilons(1:num,1)=xt(num+1:2*num)
   locepsilons(1:num,2)=xt(2*num+1:3*num)
   locepsilons(1:num,3)=xt(3*num+1:4*num)

   pairenergy=0.0_wp
   selfenergy=0.0_wp
   pairxi=0.0_wp
   selfxi=0.0_wp
   call MW_daim_diffrepulsionenergy(system%num_PBC, system%localwork, system%daim, system%box, &
               system%ions, system%xyz_ions(:,:,xyz_now), &
               system%electrodes, system%xyz_atoms,locdeltas, &
               locepsilons,pairenergy,pairxi)

   call MW_daim_diffselfenergy(system%num_ion_types,system%ions,locdeltas,locepsilons,selfenergy,selfxi)

   energytot=pairenergy+selfenergy
   energyanddiff(1)=energytot
   locxi=pairxi+selfxi
   energyanddiff(2)=dot_product(locxi,xicomaim)

return
end function df1dimaim
end module MW_cgaim
