!===============================================================================
!
! multi_shift_cr.F90 - conjugate-residual multi shift solver
!                      ( M + a(i) ) x(i) = b
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2006-2010 Yoshifumi Nakamura
!
! This file is part of BQCD -- Berlin Quantum ChromoDynamics program
!
! BQCD is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! BQCD is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with BQCD.  If not, see <http://www.gnu.org/licenses/>.
!
!-------------------------------------------------------------------------------
! Added this file to accomodate init of Zolotarev ratapp action terms
!
! Editor: Taylor Haar
! Last modified: 29/09/2016
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------
! Multi-shift Conjugate Residual method for complex shifts
! -- only works if matrix() is hermitian
!-------------------------------------------------------------------------------
subroutine multi_shift_crz(matrix, id, n, shift_switch, betabar, b, iterations, niter_shift, res) 
  use module_field
  use module_cg
  use module_function_decl
  use module_p_interface
  use module_vol
  implicit none
  SPINCOL_OVERINDEXED,           intent(in)  :: b
  integer,                       intent(out) :: iterations, niter_shift(n)
  integer,                       intent(in)  :: shift_switch(n), n, id
  COMPLEX,                       intent(in)  :: betabar(n)
  REAL,                          intent(in)  :: res(n)
  external :: matrix

  REAL, allocatable :: r(:),p(:),ap(:),ar(:)
  REAL              :: alpha, alpha_old, beta, rtr, rtar, rtarold, paap, bb
  COMPLEX           :: sbeta, salpha
  COMPLEX, dimension(n) :: sxi_old, sxi, sxi_new
  integer           :: i,j, niter, n_min, n_con(n), niter_sub, ierr, ncon
  character(72)     :: msg


  TIMING_START(timing_bin_cg_mcg)

  allocate(r(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot), STAT = ierr)
  if (ierr /= 0) then
     call stderr2_int("allocation failed", 1, ierr)
     call die("cgm: allocation failed r")
  endif
  allocate(p(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot), STAT = ierr)
  if (ierr /= 0) then
     call stderr2_int("allocation failed", 1, ierr)
     call die("cgm: allocation failed p")
  endif
  allocate(ap(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot), STAT = ierr)
  if (ierr /= 0) then
     call stderr2_int("allocation failed", 1, ierr)
     call die("cgm: allocation failed ap")
  endif
  allocate(ar(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot), STAT = ierr)
  if (ierr /= 0) then
     call stderr2_int("allocation failed", 1, ierr)
     call die("cgm: allocation failed ar")
  endif

  n_min = n
  do i = 1, n
     n_con(i) = shift_switch(i)
     if (shift_switch(i) == 0) n_min = min(n_min,i)
  enddo

  alpha_old=ONE
  beta=ZERO
  sxi_old  =ONE
  sxi      =ONE

  call sc_copy(r,b)
  call matrix(ar, r, id)
  rtarold = sc_dot(r, ar); rtarold = global_sum(rtarold)

  rtr = sc_dot(r, r); rtr = global_sum(rtr)
  bb = rtr

  call sc_copy(p,r)
  call sc_copy(ap,ar)
  do i=1,n
     call sc_copy(sco(2,i)%sco(1),b)
     call sc_zero(sco(1,i)%sco(1))
  enddo

niter_sub=0
!----------------------------------------
  do niter = 1, cg_para%maxiter
     paap = sc_dot(ap, ap) ; paap = global_sum(paap)
     alpha= - rtarold / paap

     call sc_axpy(r, ap, alpha)
     call matrix(ar, r, id)
     rtar = sc_dot(r, ar); rtar = global_sum(rtar)
     rtr = sc_dot(r, r); rtr = global_sum(rtr)

     do j = 1,n
        if ( n_con(j) == 0) then
           sxi_new(j)=sxi(j)*sxi_old(j)*alpha_old/(alpha*beta*(sxi_old(j) &
                     -sxi(j))+sxi_old(j)*alpha_old*(ONE-betabar(j)*alpha))
           salpha=alpha*sxi_new(j)/sxi(j)
           call sc_caxpy(sco(1,j)%sco(1), sco(2,j)%sco(1), -salpha)
           niter_sub = niter_sub + 1

           if (stop_condition((Re(sxi_new(j))**2 + Im(sxi_new(j))**2)*rtr, &
                res(j), bb)) then
              niter_shift(j)= niter
              n_con(j)=1
           endif
        endif
     enddo

#ifdef DEBUG2
if(my_pe()==0 .and. mod(niter,100)==0  )write(0,*)niter,"mcr:rtr",sxi_new(n_min)*sxi_new(n_min)*rtr
#endif
     if (stop_condition((Re(sxi_new(n_min))**2 + Im(sxi_new(n_min))**2)*rtr, &
          res(n_min), bb)) then
        ncon=0
        do j = n_min, n
           ncon=ncon + n_con(j)
        enddo
        if (ncon == n - n_min +1) goto 9999
     endif

     beta = rtar / rtarold
     rtarold   = rtar
     call sc_xpby(p, r, beta)  ! p = r + beta * p
     call sc_xpby(ap, ar, beta)  ! ap = ar + beta * ap

     do j = 1,n
        if ( n_con(j) == 0) then
           sbeta = beta * (sxi_new(j)/sxi(j))**2
           call sc_caxpby(sco(2, j)%sco(1), r, sbeta, sxi_new(j))
           sxi_old(j)  =sxi(j)
           sxi(j)      =sxi_new(j)
        endif
     enddo

     alpha_old=alpha
  enddo
!----------------------------------------
  niter = niter - 1

  if (cg_para%log /= 2) then
     if(my_pe()==0)write(UREC, *) "cgr(): no convergence; rtr = ", rtr
     if(my_pe()==0)write(UREC, *) "b    = ",b(1)
     if(my_pe()==0)write(UREC, *) "r    = ",r(1)
     if(my_pe()==0)write(UREC, *) "aa   = ",sco(1,1)%sco(1)
     if(my_pe()==0)write(UREC, *) "pm   = ",sco(2,1)%sco(1)
     if(my_pe()==0)write(UREC, *) "alpha= ",alpha
     if(my_pe()==0)write(UREC, *) "beta = ",beta
     write(msg, *) "multi_shift_cgr(): no convergence; rtr = ", rtr
     call die(msg)
  endif

9999 iterations = niter
  if (cg_para%check /= 0) then
     call check_solution_zshifts(matrix, id, n, shift_switch, betabar, niter_shift, b) !?
  endif
  call update_cg_stat(cgm_stat, niter, niter_sub)

  deallocate(r, STAT = ierr)
  if (ierr /= 0) then
     call stderr2_int("deallocation failed", 1, ierr)
     call die("cgm: deallocation failed r")
  endif
  deallocate(p, STAT = ierr)
  if (ierr /= 0) then
     call stderr2_int("deallocation failed", 1, ierr)
     call die("cgm: deallocation failed p")
  endif
  deallocate(ap, STAT = ierr)
  if (ierr /= 0) then
     call stderr2_int("deallocation failed", 1, ierr)
     call die("cgm: deallocation failed ap")
  endif
  deallocate(ar, STAT = ierr)
  if (ierr /= 0) then
     call stderr2_int("deallocation failed", 1, ierr)
     call die("cgm: deallocation failed ar")
  endif

  TIMING_STOP(timing_bin_cg_mcg)

end subroutine multi_shift_crz
