!===============================================================================
!
! multi_shift_blockcg.F90 - block cg with multi shift solver  ( A + a(i) ) x(i,j) = b(j)
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2012 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/>.
!
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------
# define TRACE_NAME multi_shift_blockcg_qr
!-------------------------------------------------------------------------------
subroutine multi_shift_blockcg_qr(matrix, id, b, iters) 
  use multi_lr_solver
  use module_cg
  use module_function_decl
  use module_input
  use module_vol
  implicit none
  integer :: id, iters
  external :: matrix

  COMPLEX,dimension(12*volh_tot, nr) :: b !!!, r, p, aap, w, x
  COMPLEX,dimension(:,:), allocatable :: r, p, aap, w, x
  REAL,dimension(nr)     :: bb, res, restmp
  COMPLEX,dimension(nr, nr, nl) :: sxi0, sxi1, sxi2, sxinv, asxi
  integer                :: i,j, niter, n_min, n_con(nl), nblas1, ncon, i1, i2
  COMPLEX,dimension(nr, nr)  :: alp0, alp, paap, rho0, rho, del0, del
  COMPLEX :: ztmp, sig(nr, nr), tmp0(nr, nr), tmp1(nr, nr), tmp2(nr,nr)

  TRACE_ENTRY
  TIMING_START(timing_bin_cg_mbcg)

#ifdef LAPACK

  if (input%tuning_cg_version == 2) then

#ifdef PRECISION_R4
     call die("cg(): multi_shift_blockcg_qr version 2 is not available for PRECISION_R4")
#elif PRECISION_R16
     call die("cg(): multi_shift_blockcg_qr version 2 is not available for PRECISION_R16")
#else
     call cg_wrapper_mbcg(matrix, id, b, iters)
#endif

  else  ! (input%cg_version /= 2)

  allocate(r(12*volh_tot, nr))
  allocate(p(12*volh_tot, nr))
  allocate(aap(12*volh_tot, nr))
  allocate(w(12*volh_tot, nr))
  allocate(x(12*volh_tot, nr))

  n_min = nl
  do i = 1, nl
     n_con(i) = lrkv(i)%frdo
     if (lrkv(i)%frdo == 0) n_min = min(n_min, i)
  enddo

  sxi0 = ZERO
  alp0 = ZERO
  do i = 1, nr
     alp0(i,i) = ONE
     do j = 1, nl
        sxi0(i,i,j)=ONE
     enddo
  enddo

  do i = 1, nr
     call sc_copy(r(1, i), b(1, i))
     bb(i) = sc_norm2(b(1, i))
  enddo
  call global_sum_vec(nr, bb)
  call gram_schmidt2(r, rho0, 12*volh_tot, 12*volh, nr, .true.)

  do i = 1, nr
     call sc_copy(p(1, i), r(1, i))
     call sc_zero(x(1, i))
     do j = 1, nl
        call sc_copy( pp(j)%v(1, i), p(1, i))
        call sc_zero( xx(j)%v(1, i))
     enddo
  enddo

  del0 = rho0
  do j = 1, nl
     sxi1(:,:,j) = rho0
  enddo

  nblas1 = 0
!----------------------------------------
  do niter = 1, cg_para%maxiter
     do i = 1, nr
        call matrix(aap(1,i), p(1,i), id)
     enddo

     do i1 = 1, nr
     do i2 = 1, nr
        paap(i2,i1) = sc_cdotc(p(1,i2), aap(1,i1)) 
     enddo
     enddo
     call global_sum_vec(2*nr*nr, paap)

     call get_matinv_local(alp, paap, nr) ! alp = paap^-1

     do i1 = 1, nr
     do i2 = 1, nr
        ztmp = - alp(i2,i1)
        call sc_caxpy(r(1,i1), aap(1,i2), ztmp)
     enddo
     enddo
     call gram_schmidt2(r, rho, 12*volh_tot, 12*volh, nr, .true.)

!!     !------------------------------------------------------------------------------for checking
!!     call zgemm('N','N', nr, nr, nr, zone, rho, nr, del0, nr, zzero, del, nr)
!!     call zgemm('N','N', nr, nr, nr, zone, alp, nr, del0, nr, zzero, sig, nr)
!!     do i1 = 1, nr
!!     do i2 = 1, nr
!!        call sc_caxpy(x(1,i1), p(1,i2), sig(i2,i1))
!!     enddo
!!     enddo
!!     do i = 1, nr
!!        call matrix(w(1,i), x(1,i), id)
!!        w(:,i) = b(:,i) - w(:,i)
!!        restmp(i)=sc_norm2(w(1,i))
!!     enddo
!!     !!write(200,*)"rr", restmp
!!     !------------------------------------------------------------------------------for checking

     do j = 1, nl
        if ( n_con(j) == 0) then
           call get_sxi(sxi2(:,:,j), sxinv(:,:,j), sxi1(:,:,j), sxi0(:,:,j), alp0, alp, rho0, rho, lrkv(j)%shift)
           call zgemm('N','N', nr, nr, nr, zone, alp,         nr, sxinv(:,:,j), nr, zzero, asxi(:,:,j), nr)
           call zgemm('N','N', nr, nr, nr, zone, asxi(:,:,j), nr, sxi1(:,:,j),  nr, zzero, sig,      nr)

           do i1 = 1, nr
           do i2 = 1, nr
              call sc_caxpy(xx(j)%v(1,i1), pp(j)%v(1,i2), sig(i2,i1))
           enddo
           enddo
           nblas1 = nblas1 + nr

           ncon = 0
           call get_residuals(res, sxi2(1,1,j))
!!           if (j==30) then
!!              lrkv(j)%iter= niter
!!              write(*,'(2x,a7 ,i2,a1,i6,1x,3(es16.6))')"%rresft(",j,")", niter, (res(1)+res(2))/2,res(1),res(2)
!!           endif
           do i = 1, nr
              if (stop_condition(res(i), lrkv(j)%tol, bb(i))) ncon = ncon + 1
           enddo
           if (ncon == nr) then
               lrkv(j)%iter= niter
               n_con(j)=1
           endif
        endif
     enddo

     ncon = 0
     call get_residuals(res, sxi2(1,1,n_min))
     do i = 1, nr
        if (stop_condition(res(i), lrkv(n_min)%tol, bb(i))) ncon = ncon + 1
     enddo

     if (ncon == nr) then
        ncon=0
        do j = n_min, nl
           ncon=ncon + n_con(j)
        enddo
        if (ncon == nl - n_min +1) goto 9999
     endif

     do i1 = 1, nr
	call sc_copy(w(1, i1), r(1, i1))
        do i2 = 1, nr
           ztmp = dconjg(rho(i1, i2))
           call sc_caxpy(w(1, i1), p(1, i2), ztmp)
	enddo
     enddo
     do i1 = 1, nr
        call sc_copy(p(1, i1), w(1, i1))
     enddo

     call zgemm('N','C', nr, nr, nr, zone, paap, nr, rho, nr, zzero, tmp1, nr)  ! tmp1 = a^-1 rho^H
     do j = 1, nl
        if ( n_con(j) == 0) then
           call zgemm('N','N', nr, nr, nr, zone, asxi(:,:,j), nr, tmp1, nr, zzero, sig, nr)
           do i1 = 1, nr
              call sc_copy(w(1, i1), r(1, i1))
              do i2 = 1, nr
                 call sc_caxpy(w(1, i1), pp(j)%v(1,i2), sig(i2,i1))
              enddo
           enddo
           do i1 = 1, nr
              call sc_copy(pp(j)%v(1,i1), w(1, i1))
           enddo
           sxi0(:,:,j) = sxi1(:,:,j)
           sxi1(:,:,j) = sxi2(:,:,j)
        endif
     enddo

     del0 = del
     rho0 = rho
     alp0 = alp
  enddo
!----------------------------------------

  niter = niter - 1
  if (cg_para%log /= 2) call die("multi_shift_blockcg_qr(): no convergence")
  9999 iters = niter

!!  call update_cg_stat(cgm_stat, niter, nblas1)

  deallocate(r)
  deallocate(p)
  deallocate(aap)
  deallocate(w)
  deallocate(x)

  endif  ! (input%cg_version /= 2)

#else
  call die("multi_shift_blockcg_qr needs LAPACK")
#endif

  if (cg_para%check /= 0) then
     call check_multi_lr_solver(matrix, id, b)
  endif

  TIMING_STOP(timing_bin_cg_mbcg)
  TRACE_RETURN
end

!-------------------------------------------------------------------------------
subroutine get_sxi(sxi2, sxinv, sxi1, sxi0, alp0, alp, rho0, rho, shift)
  use multi_lr_solver
  implicit none
  COMPLEX, dimension(nr, nr) :: sxi2, sxinv, sxi1, sxi0, alp0, alp, rho0, rho
  COMPLEX, dimension(nr, nr) :: unit, tmp0, tmp1, tmp2, tmp3
  REAL :: shift
  integer :: i

  unit = ZERO
  do i = 1, nr
     unit(i, i) = ONE
  enddo
#ifdef LAPACK
  call zgemm('N','C', nr, nr, nr, zone, unit,  nr, rho0, nr, zzero, tmp2, nr) ! I rho^H
  call matinvmul_local(tmp2, alp0, nr)                                        ! tmp2 = alp0^-1 rho^H 

  tmp1 = unit
  call matinvmul_local(tmp1, sxi0, nr)                                        ! sxi0^-1
  call zgemm('N','N', nr, nr, nr, zone, sxi1,  nr, tmp1, nr, zzero, tmp0, nr) ! tmp0 = sxi1 sxi0^-1

  tmp0 = rho0 - tmp0                                                          ! = (rho0 - sxi1 sxi0^-1)=A
  call zgemm('N','N', nr, nr, nr, zone, tmp0,  nr, tmp2, nr, zzero, tmp1, nr) ! tmp1 = A * alp0^-1 rho^H
  sxinv = unit + shift * alp

  call zgemm('N','N', nr, nr, nr, zone, tmp1,  nr, alp, nr, zone,  sxinv, nr)
  call matinv_local(sxinv, nr)

  call zgemm('N','N', nr, nr, nr, zone, rho,  nr, sxinv, nr, zzero, tmp0, nr)
  call zgemm('N','N', nr, nr, nr, zone, tmp0, nr, sxi1, nr, zzero, sxi2, nr)
#else
  call die("multi_lr_solver_qr needs LAPACK")
#endif
end

!===============================================================================
!-------------------------------------------------------------------------------
!-------------------------------------------------------------------------------
subroutine get_sxi2(sxi2, sxinv, sxi1, sxi0, alp0, alp, rho0, rho, shift)
  use multi_lr_solver
  implicit none
  COMPLEX, dimension(nr, nr) :: sxi2, sxinv, sxi1, sxi0, alp0, alp, rho0, rho
  COMPLEX, dimension(nr, nr) :: unit, tmp0, tmp1, tmp2, tmp3, alp0inv, isa
  REAL :: shift
  integer :: i

  unit = ZERO
  do i = 1, nr
     unit(i, i) = ONE
  enddo

  alp0inv = alp0
  call matinv_local(alp0inv, nr)
  isa = unit + shift * alp

#ifdef LAPACK
  call zgemm('C','N', nr, nr, nr, zone, rho0,  nr, alp, nr, zzero, tmp2, nr)  ! tmp2=rho^H alp

  tmp1 = unit
  call matinvmul_local(tmp1, sxi0, nr)                                        ! sxi0^-1
  call zgemm('N','N', nr, nr, nr, zone, sxi1,  nr, tmp1, nr, zzero, tmp0, nr) ! tmp0 = sxi1 sxi0^-1

  tmp0 = rho0 - tmp0                                                          ! = (rho0 - sxi1 sxi0^-1)=A
  call zgemm('N','N', nr, nr, nr, zone, tmp0,  nr, tmp2, nr, zzero, tmp1, nr) ! tmp1 = A * rho^H alp


  call zgemm('N','N', nr, nr, nr, zone, alp0inv,  nr, tmp1, nr, zone,  isa, nr)
  sxinv=isa
  call matinv_local(sxinv, nr)

  call zgemm('N','N', nr, nr, nr, zone, rho,  nr, sxinv, nr, zzero, tmp0, nr)
  call zgemm('N','N', nr, nr, nr, zone, tmp0, nr, sxi1, nr, zzero, sxi2, nr)
#else
  call die("multi_lr_solver_qr needs LAPACK")
#endif
end

!===============================================================================
