!===============================================================================
!
! cg_kernel_mbcg.F90 - multi_shift_blockcg_qr kernel of BQCDi
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2012 Yoshifumi Nakamura
!               2016 Hinnerk Stueben
!
! 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/>.
!
!-------------------------------------------------------------------------------
# define LANGUAGE_FORTRAN
# include "defs.h"
# include "cg_defs.h"
# include "alignment.h"

!-------------------------------------------------------------------------------
# define TRACE_NAME T(cg_kernel_mbcg)
!-------------------------------------------------------------------------------
subroutine T(cg_kernel_mbcg)(matrix, id, b, iters) 
  use multi_lr_solver
  use module_cg
  use module_cg_procs, exclude => T(cg_kernel_mbcg)  ! avoid name clash
  use module_function_decl
  use module_vol
  use T(module_vol), volh_rrs => volh, volh_tot_rrs => volh_tot  ! two defs of volh and volh_tot are needed
  implicit none

  external :: matrix  ! not used here
  integer  :: id
  COMPLEX  :: b(12*volh_tot, nr)
  integer  :: iters


  COMPLEX,dimension(:,:), allocatable :: r, w 
  ATTRIBUTE_ALIGN2(r, w)
  COMPLEX,dimension(:,:), allocatable :: p, aap  ! only these arrays are used with SIMD (ie with routines cg_procs%...)
  ATTRIBUTE_ALIGN2(p, aap)
  REAL8,dimension(nr) :: bb, res, restmp
  COMPLEX8,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, del
  COMPLEX :: ztmp, sig(nr, nr), tmp0(nr, nr), tmp1(nr, nr), tmp2(nr,nr)


  if (nr == 2) then
     call T(cg_kernel_mbcg2)(matrix, id, b, iters)
     return
  endif

  TRACE_ENTRY
#ifdef LAPACK

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

  do i = 1, nr
     call sc_zero(aap(1, i))
     call sc_zero_halo(aap(1, i))
     call sc_zero(p(1, i))
     call sc_zero_halo(p(1, i))
  enddo

  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

  TIMING_START(timing_bin_cg_global_sum)
  call global_sum_vec(nr, bb)
  TIMING_STOP(timing_bin_cg_global_sum)

  call gram_schmidt2(r, rho0, 12*volh_tot, 12*volh, nr, .true.)

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

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

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

     do i1 = 1, nr
     do i2 = 1, nr
        call cg_procs%cdotc(paap(i2,i1), p(1,i2), aap(1,i1), volh_tot_rrs, volh_rrs)
     enddo
     enddo

     TIMING_START(timing_bin_cg_global_sum)
     call global_sum_vec(2*nr*nr, paap)
     TIMING_STOP(timing_bin_cg_global_sum)

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

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


     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))
           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 = conjg(rho(i1, i2))
           call cg_procs%copy_caxpy(w(1, i1), p(1, i2), ztmp)
	enddo
     enddo
     do i1 = 1, nr
        call cg_procs%copyin_spincol(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

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

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

!!  call update_cg_stat(cgm_stat, niter, nblas1)

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

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

  TRACE_RETURN
end
