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

   ! Public type
   ! -----------
   public :: MW_cg_t

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

   type MW_cg_t

      ! Convergence parameters
      integer :: max_iterations !< maximum number of iterations
      real(wp) :: tol           !< tolerance on residual norm
      logical :: charge_neutrality = .false.
      logical :: preconditioned = .false.
      logical :: prec_jacobi = .false.

      ! 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(:)   !< projection of res on the hyperplan <Q>=cst (only for charge neutrality)
      real(wp), allocatable :: d(:)   !< conjugate direction
      real(wp), allocatable :: Ad(:)  !< A times d
      real(wp), allocatable :: jacobi_precond(:) !< Jacobi preconditioner
      real(wp), allocatable :: precond(:,:) !< General preconditioner
      real(wp), allocatable :: Minvp(:)  !< M^-1 times res

   end type MW_cg_t

contains

   ! ================================================================================
   ! Set maximum number of iterations
   subroutine set_max_iterations(this, max_iterations)
      implicit none
      type(MW_cg_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_cg_t), intent(inout) :: this
      real(wp), intent(in) :: tol
      this%tol = tol
   end subroutine set_tolerance

   ! ================================================================================
   ! Set force neutral
   subroutine set_charge_neutrality(this, charge_neutrality)
      implicit none
      type(MW_cg_t), intent(inout) :: this
      logical, intent(in) :: charge_neutrality
      this%charge_neutrality = charge_neutrality
   end subroutine set_charge_neutrality

   ! ================================================================================
   ! Set force neutral
   subroutine set_preconditioner(this, preconditioned, preconditioner)
      use MW_configuration_line, only: max_word_length
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      implicit none
      type(MW_cg_t),              intent(inout) :: this
      logical,                    intent(in)    :: preconditioned
      character(max_word_length), intent(in)    :: preconditioner
      this%preconditioned = preconditioned
      select case (preconditioner)
      case ("jacobi")
         this%prec_jacobi = .true.
      case default
         call MW_errors_runtime_error("set preconditioner", "cg.f90", &
               "Invalid preconditioner for conjugate gradient method")
      end select
   end subroutine set_preconditioner

   ! ================================================================================
   ! reset statistic counters
   subroutine reset_statistics(this)
      implicit none
      type(MW_cg_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_cg_t), intent(inout) :: this

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

      ! Local
      ! -----
      integer :: i, j
      integer :: ierr

      this%n = n
      allocate(this%b(n), this%res(n), this%p(n), this%d(n), this%Ad(n), this%Minvp(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%d(i) = 0.0_wp
         this%Ad(i) = 0.0_wp
         this%Minvp(i) = 0.0_wp
      end do

      if (this%preconditioned) then
         if (this%prec_jacobi) then
            allocate(this%jacobi_precond(n), stat=ierr)
            do i = 1, n
               this%jacobi_precond(i) = 1.0_wp
            end do
         else
            allocate(this%precond(n, n), stat=ierr)
            do i = 1, n
               this%precond(i,i) = 1.0_wp
               do j = 1, i-1
                  this%precond(i,j) = 0.0_wp
                  this%precond(j,i) = 0.0_wp
               end do
            end do

         end if
         if (ierr /= 0) then
            call MW_errors_allocate_error("define_type", "cg.f90", ierr)
         end if
      end if

   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_cg_t), intent(inout) :: this

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

      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%d, this%Ad, this%Minvp, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "cg.f90", ierr)
         end if
      end if
      if (this%preconditioned) then
         if (allocated(this%jacobi_precond)) then
            deallocate(this%jacobi_precond, stat=ierr)
         else
            deallocate(this%precond, stat=ierr)
         end if
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "cg.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_cg_t), intent(in) :: this
      integer,       intent(in) :: ounit

      write(ounit, '("|cg| maximum number of iterations: ",i12)') this%max_iterations
      write(ounit, '("|cg| Residual error:               ",es12.5)') this%tol
      if (this%charge_neutrality) then
         write(ounit, '("|cg| Force charge neutrality:      ",a3)') "yes"
      else
         write(ounit, '("|cg| Force charge neutrality:      ",a3)') " no"
      end if
      if (this%preconditioned) then
         write(ounit, '("|cg| Preconditioned:      ",a3)') "yes"
      else
         write(ounit, '("|cg| Preconditioned:      ",a3)') " no"
      end if
   end subroutine print_type

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

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

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

   ! ================================================================================
   !! Solve the equation Ax = b, with preconditioner M
   !!
   !! Algorithm
   !! ----------
   !! r_0 = b - Ax_0
   !! z_0 = M^{-1}r_0
   !! p_0 = z_0
   !!
   !! do iter = 1, max_iter
   !!
   !!   alpha_i = (r_i* x z_i) / (p_i* x Ap_i)
   !!   x_{i+1} = x_{i} + alpha_i x p_i
   !!   r_{i+1} = r_{i} - alpha_i x Ap_i
   !!   if (|r_{i+1}| < tol) exit loop
   !!   z_{i+1} = M^{-1}r_{i+1}
   !!   beta_i = (z_{i+1}* x r_{i+1}) / (z_{i}* x r_{i})
   !!   p_{i+1} = z_{i+1} + beta_{i} x p_i
   !!
   !! end do
   !! Adaptation to the constraint of charge neutrality implemented from the paper of Shariff (1995)
   subroutine solve(this, system, apply_A, x)
      use MW_kinds, only: wp
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      use MW_system, only: MW_system_t
      use MW_ion, only: MW_ion_t
      use MW_stdio, only: MW_stderr
      implicit none
      ! Parameters
      ! ----------
      type(MW_cg_t), intent(inout) :: this !> cg algorithm parameters
      type(MW_system_t), intent(inout) :: system !> system parameters

      !> Interface for apply_A routine
      !> Computes y = Ax
      interface
         subroutine apply_A(system, x, y)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            implicit none
            type(MW_system_t), intent(inout) :: system
            real(wp), intent(in) :: x(:)
            real(wp), intent(out) :: y(:)
         end subroutine apply_A
      end interface

      real(wp), intent(inout) :: x(:) !< Initial guess (in) and solution (out)

      ! Locals
      ! ------
      integer :: num, num_ions, num_atoms !< size of the solution vector
      integer :: i, iter
      real(wp) :: gammaold, gammanew, rsnew !< scalar to store residual norm
      real(wp) :: alpha, beta, dAd, dp  !< CG coefficients
      real(wp) :: res_norm, res_tol
      logical :: cg_converged
      real(wp) :: average_res, total_res

      num = size(x, 1)
      num_ions = system%num_ions
      num_atoms = system%num_atoms
      ! Setup tolerance on residual norm
      ! stopping criteria: |r_k| < tol * |b| where tol is a user input parameter
      res_tol = this%tol * sqrt(real(num,wp))

      ! Compute Ax0
      call apply_A(system, x, this%Ad)

      ! Setup initial residual
      do i = 1, num
         this%res(i) = this%b(i) - this%Ad(i)
      end do

      ! p0 = z0

      if (system%electrode_charge_neutrality) then
         ! projection of the gradient on the hypersurface <Q>=cst
         total_res = 0.0_wp
         if (system%pimrun) then
            do i = 3*num_ions+1, num
               total_res = total_res + this%res(i)
            end do
            average_res = total_res / real(num_atoms,wp)
            do i = 1, 3*num_ions
               this%p(i) = this%res(i)
            end do
            do i = 3*num_ions+1, num
               this%p(i) = this%res(i) - average_res
            end do
         else
            do i = 1, num
               total_res = total_res + this%res(i)
            end do
            average_res = total_res / real(num,wp)
            do i = 1, num
               this%p(i) = this%res(i) - average_res
            end do
         end if
      else
         do i = 1, num
            this%p(i) = this%res(i)
         end do
      endif


      do i = 1, num
         this%d(i) = this%p(i)
      end do

      rsnew=0.0_wp
      do i = 1, num
         rsnew = rsnew + this%d(i)*this%d(i)
      end do
      res_norm = sqrt(rsnew)
      gammanew=0.0_wp
      do i = 1, num
         gammanew = gammanew + this%p(i)*this%p(i)
      end do


      ! Begin CG iterations
      cg_converged = .false.
      do iter = 1, this%max_iterations

         call apply_A(system, this%d, this%Ad)

         dAd=0.0_wp
         do i = 1, num
            dAd = dAd + this%d(i)*this%Ad(i)
         end do
         dp=0.0_wp
         do i = 1, num
            dp = dp + this%d(i)*this%p(i)
         end do
         alpha = dp / dAd
         do i = 1, num
            x(i) = x(i) + alpha * this%d(i)
         end do
         gammaold = gammanew

         ! Avoid residual drift by recomputing it as res=b-Ax as suggested in Gingrich master's thesis
         if (mod(iter,50) == 0) then
            call apply_A(system, x, this%Ad)
            do i = 1, num
               this%res(i) = this%b(i) - this%Ad(i)
            end do
         else
            do i = 1, num
               this%res(i) = this%res(i) - alpha * this%Ad(i)
            end do
         endif

         if (system%electrode_charge_neutrality) then
            ! projection of the gradient on the hypersurface <Q>=cst
            total_res = 0.0_wp
            if (system%pimrun) then
               do i = 3*num_ions+1, num
                  total_res = total_res + this%res(i)
               end do
               average_res = total_res / real(num_atoms,wp)
               do i = 1, 3*num_ions
                  this%p(i) = this%res(i)
               end do
               do i = 3*num_ions+1, num
                  this%p(i) = this%res(i) - average_res
               end do
            else
               do i = 1, num
                  total_res = total_res + this%res(i)
               end do
               average_res = total_res / real(num,wp)
               do i = 1, num
                  this%p(i) = this%res(i) - average_res
               end do
            end if
         else
            do i = 1, num
               this%p(i) = this%res(i)
            end do
         endif

         gammanew=0.0_wp
         do i = 1, num
            gammanew = gammanew + this%p(i)*this%p(i)
         end do
         ! Setup for next iteration
         beta = gammanew / gammaold
         do i = 1, num
            this%d(i) = this%p(i) + beta * this%d(i)
         end do

         rsnew=0.0_wp
         do i = 1, num
            rsnew = rsnew + this%d(i)*this%d(i)
         end do

         res_norm = sqrt(rsnew)

         if (res_norm < res_tol) then
            this%last_iteration_count = iter
            this%total_iteration_count = this%total_iteration_count + iter
            this%last_residual = res_norm
            this%last_residual_tol = res_tol
            cg_converged = .true.
            exit
         end if

      end do

      if (.not. cg_converged) then
         this%last_iteration_count = iter
         this%last_residual = res_norm
         this%last_residual_tol = res_tol
         call print_statistics(this, MW_stderr)
         call MW_errors_runtime_error("solve", "cg.f90", "CG failed to converge")
      end if
   end subroutine solve

   subroutine apply_diagMinv(x, diagM, y)
      implicit none
      ! Parameters Input/Output
      ! ----------
      real(wp), intent(in) :: x(:)
      real(wp), intent(in) :: diagM(:)
      real(wp), intent(out) :: y(:)

      integer :: i, n

      n = size(x)
      do i = 1, n
         y(i) = 1.0_wp/diagM(i)*x(i)
      end do

   end subroutine apply_diagMinv

   subroutine apply_Minv(x, Minv, y)
      implicit none
      ! Parameters Input/Output
      ! ----------
      real(wp), intent(in) :: x(:)
      real(wp), intent(in) :: Minv(:,:)
      real(wp), intent(out) :: y(:)

      integer :: i, j, n

      !TO BE PARALLELIZED
      n = size(x)
      do j = 1, n
         do i = 1, n
            y(i) = y(i) + Minv(i, j)*x(i)
         end do
      end do

   end subroutine apply_Minv

   subroutine solve_preconditioned(this, system, apply_A, x)
      use MW_kinds, only: wp
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      use MW_system, only: MW_system_t
      use MW_ion, only: MW_ion_t
      use MW_stdio, only: MW_stderr
      implicit none
      ! Parameters
      ! ----------
      type(MW_cg_t), intent(inout) :: this !> cg algorithm parameters
      type(MW_system_t), intent(inout) :: system !> system parameters

      !> Interface for apply_A routine
      !> Computes y = Ax
      interface
         subroutine apply_A(system, x, y)
            use MW_kinds, only: wp
            use MW_system, only: MW_system_t
            use MW_ion, only: MW_ion_t
            implicit none
            type(MW_system_t), intent(inout) :: system
            real(wp), intent(in) :: x(:)
            real(wp), intent(out) :: y(:)
         end subroutine apply_A
      end interface

      real(wp), intent(inout) :: x(:) !< Initial guess (in) and solution (out)

      ! Locals
      ! ------
      integer :: num, num_ions, num_atoms !< size of the solution vector
      integer :: iter
      real(wp) :: gammaold, gammanew, rsnew !< scalar to store residual norm
      real(wp) :: alpha, beta, dAd  !< CG coefficients
      real(wp) :: res_norm, res_tol
      logical :: cg_converged
      real(wp) :: average_res, total_res

      num = this%n
      num_ions = system%num_ions
      num_atoms = system%num_atoms

      res_tol = this%tol * sqrt(real(num,wp))

      call apply_A(system, x, this%Ad)

      ! Setup initial residual
      this%res(1:num) = this%b(1:num) - this%Ad(1:num)

      if (system%electrode_charge_neutrality) then
         ! projection of the gradient on the hypersurface <Q>=cst
         !$acc kernels
         if (system%pimrun) then
            total_res = sum(this%res(3*num_ions+1:3*num_ions+num_atoms))
            average_res = total_res / real(num_atoms,wp)
            this%p(1:3*num_ions) = this%res(1:3*num_ions)
            this%p(3*num_ions+1:3*num_ions+num_atoms) = this%res(3*num_ions+1:3*num_ions+num_atoms) - average_res
         else
            total_res = sum(this%res(1:num))
            average_res = total_res / real(num,wp)
            this%p(1:num) = this%res(1:num) - average_res
         end if
         !$acc end kernels
      else
         !$acc kernels
         this%p(1:num) = this%res(1:num)
         !$acc end kernels
      endif

      if (this%preconditioned) then
         if (this%prec_jacobi) then
            call apply_diagMinv(this%p, this%jacobi_precond, this%Minvp)
         else
            call apply_Minv(this%p, this%precond, this%Minvp)
         end if
         this%d(1:num) = this%Minvp(1:num)
      else
         this%Minvp(1:num) = this%p(1:num)
         this%d(1:num) = this%Minvp(1:num)
      end if

      rsnew = sum(this%p(1:num)*this%p(1:num))
      res_norm = sqrt(rsnew)

      gammanew = sum(this%p(1:num)*this%Minvp(1:num))

      cg_converged = .false.
      do iter = 1, this%max_iterations

         call apply_A(system, this%d, this%Ad)

         dAd = sum(this%d(1:num)*this%Ad(1:num))
         alpha = gammanew / dAd
         x(1:num) = x(1:num) + alpha*this%d(1:num)

         if (mod(iter,50) == 0) then
            call apply_A(system, x, this%Ad)
            this%res(1:num) = this%b(1:num) - this%Ad(1:num)
         else
            this%res(1:num) = this%res(1:num) - alpha * this%Ad(1:num)
         endif

         if (system%electrode_charge_neutrality) then
            ! projection of the gradient on the hypersurface <Q>=cst
            !$acc kernels
            if (system%pimrun) then
               total_res = sum(this%res(3*num_ions+1:3*num_ions+num_atoms))
               average_res = total_res / real(num_atoms,wp)
               this%p(1:3*num_ions) = this%res(1:3*num_ions)
               this%p(3*num_ions+1:3*num_ions+num_atoms) = this%res(3*num_ions+1:3*num_ions+num_atoms) - average_res
            else
               total_res = sum(this%res(1:num))
               average_res = total_res / real(num,wp)
               this%p(1:num) = this%res(1:num) - average_res
            end if
            !$acc end kernels
         else
            !$acc kernels
            this%p(1:num) = this%res(1:num)
            !$acc end kernels
         endif

         if (this%preconditioned) then
            if (this%prec_jacobi) then
               call apply_diagMinv(this%p, this%jacobi_precond, this%Minvp)
            else
               call apply_Minv(this%p, this%precond, this%Minvp)
            end if
         else
            this%Minvp(1:num) = this%p(1:num)
         end if

         gammaold = gammanew
         gammanew = sum(this%p(1:num)*this%Minvp(1:num))

         beta = gammanew / gammaold
         this%d(1:num) = this%Minvp(1:num) + beta*this%d(1:num)

         rsnew = sum(this%p(1:num)*this%p(1:num))
         res_norm = sqrt(rsnew)

         if (res_norm < res_tol) then

            this%last_iteration_count = iter
            this%total_iteration_count = this%total_iteration_count + iter
            this%last_residual = res_norm
            this%last_residual_tol = res_tol
            cg_converged = .true.
            exit
         end if
      end do

      if (.not. cg_converged) then

         this%last_iteration_count = iter
         this%last_residual = res_norm
         this%last_residual_tol = res_tol
         call print_statistics(this, MW_stderr)
         call MW_errors_runtime_error("solve", "cg.f90", "CG failed to converge")
      end if
   end subroutine solve_preconditioned

end module MW_cg
