!===============================================================================
!
! 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"

!-------------------------------------------------------------------------------
subroutine multi_shift_blockcg(matrix, id, b, iters) 
  use multi_lr_solver
  use module_cg
  use module_function_decl
  use module_vol
  implicit none
  integer :: id, iters
  external :: matrix

  REAL,dimension(24*volh_tot, nr) :: b, r, p, aap, w
  REAL,dimension(nr)     :: bb, res
  COMPLEX,dimension(nr, nr, nl) :: sxi0, sxi1, sxi2,  salpha
  integer                :: i,j, niter, n_min, n_con(nl), nblas1, ncon, i1, i2
  COMPLEX,dimension(nr,nr)  :: rr0, rr, alp0, alp, bet, paap, sbeta, tmpp

  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

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

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

  rr0 = ZERO
  do i1 = 1, nr
  do i2 = 1, nr
!if(i1==i2)then
     rr0(i1, i2) = sc_cdotc(r(1, i1),r(1, i2))
!endif
  enddo
  enddo
  call global_sum_vec(2*nr*nr, rr0)

  do i = 1, nr
     bb(i) = rr0(i,i)
  enddo

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

     paap = ZERO
     do i1 = 1, nr
     do i2 = 1, nr
!if(i1==i2)then
        paap(i1,i2) = sc_cdotc(p(1,i1), aap(1,i2)) 
!endif
     enddo
     enddo
     call global_sum_vec(2*nr*nr, paap)

     alp = - rr0
     call matinvmul_local(alp, paap, nr) ! alp = -rr0 / paap


     do i1 = 1, nr
     do i2 = 1, nr
        call sc_caxpy(r(1,i1), aap(1,i2), alp(i2,i1))
     enddo
     enddo

     rr=ZERO
     do i1 = 1, nr
     do i2 = 1, nr
!if(i1==i2)then
        rr(i1, i2) = sc_cdotc(r(1, i1),r(1, i2))
!endif
     enddo
     enddo
     call global_sum_vec(2*nr*nr, rr)

     do j = 1, nl
        if ( n_con(j) == 0) then
           call get_sxi_aa(sxi2(:,:,j), sxi1(:,:,j), sxi0(:,:,j), alp0, alp, bet, lrkv(j)%shift)
           call matmul_local(salpha(:,:,j), alp, sxi2(:,:,j), nr)

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

           ncon = 0
!           call matmul2_local( sxi2(:,:,j),  sxi1(:,:,j), nr)
!           call get_residuals(res, sxi2(1,1,j))
!           write(*,*)"rs", res

           do i = 1, nr
!              if (stop_condition(sxi2(i,i,j)**2*sxi1(i,i,j)**2*rr(i,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
     do i = 1, nr
!        if (stop_condition(sxi2(i,i,j)**2*sxi1(i,i,j)**2*rr(i,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

     bet = rr
     call matinvmul_local(bet, rr0, nr)

     rr0   = rr
     do i1 = 1, nr
	call sc_copy(w(1, i1), r(1, i1))
        do i2 = 1, nr
           call sc_caxpy(w(1, i1), p(1, i2), bet(i2, i1))
	enddo
     enddo
     do i1 = 1, nr
        call sc_copy(p(1, i1), w(1, i1))
     enddo

     do j = 1, nl
        if ( n_con(j) == 0) then
!           call get_matinv_local(tmpp, sxi1(:,:,j), nr)
!           call matmul_local(sbeta, salpha(:,:,j), sxi2(:,:,j), nr)
!           call matmul2_local(sbeta, paap, nr)
!           call matmul2_local(sbeta, sxi2(:,:,j), nr)
!           call matmul2_local(sbeta, bet, nr)

           call matmul_local(sbeta, bet, sxi2(:,:,j), nr)
           call matmul2_local(sbeta, sxi2(:,:,j), nr)

           call matmul2_local( sxi2(:,:,j),  sxi1(:,:,j), nr)
           do i1 = 1, nr
              call sc_zero(w(1, i1))
              do i2 = 1, nr
                 call sc_caxpy(w(1, i1), pp(j)%v(1,i2), sbeta(i2,i1))
                 call sc_caxpy(w(1, i1), r(1,i2), sxi2(i2,i1,j))
              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

     alp0=alp
  enddo
!----------------------------------------

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

!!  if (cg_para%check /= 0) then
     call check_multi_lr_solver(matrix, id, b)
!!  endif
!!  call update_cg_stat(cgm_stat, niter, nblas1)

end

!===============================================================================
!-------------------------------------------------------------------------------
!subroutine dmatinvmul_local(out, in, n)
!  implicit none
!  REAL, dimension(n, n) :: out, in
!  integer :: ipiv(n), info, n!
!
!  call DGESV( n, n, in, n, ipiv, out, n, info )!
!
!end
!
!-------------------------------------------------------------------------------
subroutine get_sxi_aa(sxi2, sxi1, sxi0, alp0, alp, bet, shift)
  use multi_lr_solver
  implicit none
  COMPLEX, dimension(nr, nr) :: sxi2, sxi1, sxi0
  COMPLEX, dimension(nr, nr) :: alp0, alp, bet, unit, tmp0, tmp1, tmp2, tmp3, alpbet
  REAL :: shift
  integer :: i, j
  
  unit = ZERO
  do i = 1, nr
     unit(i, i) = ONE
  enddo
#ifdef LAPACK
  call get_matinv_local(tmp0, sxi0, nr)
  call zgemm('N','N', nr, nr, nr, zone, sxi1, nr, tmp0, nr, zzero, tmp1, nr)
  tmp1 = unit - tmp1

  call get_matinv_local(tmp0, alp0, nr)
  call zgemm('N','N', nr, nr, nr, zone, tmp1, nr, tmp0, nr, zzero, tmp2, nr) ! (1-sx1 sxi0^-1) alp0^-1
  call zgemm('N','N', nr, nr, nr, zone, tmp2, nr, bet,  nr, zzero, tmp1, nr) ! (1-sx1 sxi0^-1) alp0^-1 bet


  tmp0 = unit - shift*alp
  call zgemm('N','N', nr, nr, nr, zone, tmp1, nr, alp,  nr, zone, tmp0, nr) !
  call get_matinv_local(sxi2, tmp0, nr) ! []^-1
#else
  call die("multi_shift_blockcg needs LAPACK")
#endif
!!  call matinvmul_local(sxi2, tmp2, nr)

end

!-------------------------------------------------------------------------------
subroutine matmul_local(out, in1, in2, n)
  implicit none
  COMPLEX, dimension(n, n) :: out, in1, in2
  integer :: i, j, k, n

  do j = 1, n
     do i = 1, n
        out(i,j) = ZERO
        do k = 1, n
           out(i,j) = out(i,j) + in1(i,k) * in2(k,j)
        enddo
     enddo
  enddo

end

!-------------------------------------------------------------------------------
subroutine matmul2_local(out, in, n)
  implicit none
  COMPLEX, dimension(n, n) :: out, in, tmp
  integer :: i, j, k, n

  do j = 1, n
     do i = 1, n
        tmp(i,j) = ZERO
        do k = 1, n
           tmp(i,j) = tmp(i,j) + out(i,k) * in(k,j)
        enddo
     enddo
  enddo
  out = tmp

end
