!===============================================================================
!
! cg_kernel_mbcg2.F90 - multi_shift_blockcg_qr kernel for nr = 2
!
!-------------------------------------------------------------------------------
!
! 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_mbcg2)
!-------------------------------------------------------------------------------
subroutine T(cg_kernel_mbcg2)(matrix, id, b, iters) 
  use multi_lr_solver, nr_multi_lr_solver => nr  ! provides: scratch space pp(:)%v and output xx(:)%v
  use module_cg
  use module_cg_procs
  use module_function_decl
  use module_vol
  implicit none

  integer, parameter :: nr = 2

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

  COMPLEX, dimension(:,:), allocatable :: r, p, aap
  ATTRIBUTE_ALIGN3(r, p, aap)
  REAL8   :: bb(nr), res
  COMPLEX8, dimension(nr, nr, nl) :: sxi0, sxi1, sxi2, sxinv, asxi
  COMPLEX,  dimension(nr, nr, nl) :: sig_ak, sig_bk
  COMPLEX,  dimension(nr, nr) :: alp0, alp, paap, rho0, rho, del, tmp
  COMPLEX :: w1, w2
  COMPLEX :: r12
  REAL    :: r11, r22, sum(3)
  integer :: i,j, niter, n_min, n_con(nl), ncon
  logical :: converged, update_xx(nl), update_pp(nl)

  TRACE_ENTRY

  allocate(r(12*volh_tot, nr))
  allocate(p(12*volh_tot, nr))
  allocate(aap(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
     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

  r22 = ONE
  rho(:, :) = ZERO

  converged = .false.
  !----------------------------------------
  do niter = 1, cg_para%maxiter

     call T(cg_mbcg2_update_p)(p(1,1), p(1,2), r(1,1), r(1,2), r22, rho)

     call cg_procs%ww(aap(1,1), p(1,1), id)
     call cg_procs%ww(aap(1,2), p(1,2), id)

     call T(cg_mbcg2_paap)(paap, p(1,1), p(1,2), aap(1,1), aap(1,2))

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

     call STRCAT(matinv2_c, RKIND)(alp, paap)

     call T(cg_mbcg2_update_r)(r(1,1), r(1,2), aap(1,1), aap(1,2), r22, alp, sum)

     TIMING_START(timing_bin_cg_global_sum)
     call global_sum_vec(3, sum)
     TIMING_STOP(timing_bin_cg_global_sum)

     r11 = sum(1)
     r12 = cmplx(sum(2), sum(3), RKIND)

     rho(1,1) = sqrt(r11); r11 = ONE/rho(1,1)
     rho(2,1) = ZERO
     rho(1,2) = r12 * r11; r12 = -rho(1,2)

     r22 = ZERO
     !$omp parallel do default(none) &
     !$omp shared(volh, r, r11, r12) &
     !$omp reduction(+:r22)
     do i = 1, 12 * volh
        r(i, 1) = r(i, 1) * r11
        r(i, 2) = r(i, 2) + r(i, 1) * r12
        r22 = r22 + Re(r(i, 2))**2 + Im(r(i, 2))**2
     enddo

     TIMING_START(timing_bin_cg_global_sum)
     call global_sum_vec(1, r22)
     TIMING_STOP(timing_bin_cg_global_sum)

     rho(2,2) = sqrt(r22); r22 = ONE/rho(2,2)
     !!call sc_scale(r(1,2), r22) -> handled elsewhere to save this loop

     do j = 1, nl
        if ( n_con(j) == 0) then
           update_xx(j) = .true.
           call cg_mbcg2_get_sxi(sxi2(1,1,j), sxinv(1,1,j), sxi1(1,1,j), sxi0(1,1,j), alp0, alp, rho0, rho, lrkv(j)%shift)

           call STRCAT(matmul2_nn_c, RKIND)(asxi(1,1,j), alp, sxinv(1,1,j))
           call STRCAT(matmul2_nn_c, RKIND)(sig_ak(1,1,j), asxi(1,1,j), sxi1(1,1,j))

           call STRCAT(matmul2_cn_c, RKIND)(tmp, sxi2(1,1,j), sxi2(1,1,j)) 

           ncon = 0
           do i = 1, nr
              res = tmp(i,i)
              if (stop_condition(res, lrkv(j)%tol, bb(i))) ncon = ncon + 1
           enddo
           if (ncon == nr) then
               lrkv(j)%iter= niter
               n_con(j)=1
           endif
        else
           update_xx(j) = .false.
        endif
     enddo

     call STRCAT(matmul2_cn_c, RKIND)(tmp, sxi2(1,1,n_min), sxi2(1,1,n_min)) 

     ncon = 0
     do i = 1, nr
        res = tmp(i,i)
        if (stop_condition(res, 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) converged = .true.
     endif

     if (converged) then
        update_pp(:) = .false.
     else
        call STRCAT(matmul2_nc_c, RKIND)(tmp, paap, rho)
        do j = 1, nl
           if ( n_con(j) == 0) then
              update_pp(j) = .true.
              call STRCAT(matmul2_nn_c, RKIND)(sig_bk(1,1,j), asxi(1,1,j), tmp)
              sxi0(:,:,j) = sxi1(:,:,j)
              sxi1(:,:,j) = sxi2(:,:,j)
           else
              update_pp(j) = .false.
           endif
        enddo
     endif

     do j = 1, nl
        if (update_xx(j) .and. update_pp(j)) then

           !$omp parallel do default(none) &
           !$omp shared(volh, j, r, r22, xx, pp, sig_ak, sig_bk) &
           !$omp private(w1, w2)
           do i = 1, 12 * volh
              w1 = r(i, 1)
              w2 = r(i, 2) * r22

              xx(j)%v(i,1) = xx(j)%v(i,1) + pp(j)%v(i,1) * sig_ak(1,1,j)
              xx(j)%v(i,2) = xx(j)%v(i,2) + pp(j)%v(i,1) * sig_ak(1,2,j)

              w1 = w1 + pp(j)%v(i,1) * sig_bk(1,1,j)
              w2 = w2 + pp(j)%v(i,1) * sig_bk(1,2,j)

              xx(j)%v(i,1) = xx(j)%v(i,1) + pp(j)%v(i,2) * sig_ak(2,1,j)
              xx(j)%v(i,2) = xx(j)%v(i,2) + pp(j)%v(i,2) * sig_ak(2,2,j)

              w1 = w1 + pp(j)%v(i,2) * sig_bk(2,1,j)
              w2 = w2 + pp(j)%v(i,2) * sig_bk(2,2,j)

              pp(j)%v(i,1) = w1
              pp(j)%v(i,2) = w2
           enddo

        elseif (update_xx(j) .and. (.not. update_pp(j))) then

           !$omp parallel do default(none) &
           !$omp shared(volh, j, xx, pp, sig_ak)
           do i = 1, 12 * volh
              xx(j)%v(i,1) = xx(j)%v(i,1) + pp(j)%v(i,1) * sig_ak(1,1,j)
              xx(j)%v(i,2) = xx(j)%v(i,2) + pp(j)%v(i,1) * sig_ak(1,2,j)
              xx(j)%v(i,1) = xx(j)%v(i,1) + pp(j)%v(i,2) * sig_ak(2,1,j)
              xx(j)%v(i,2) = xx(j)%v(i,2) + pp(j)%v(i,2) * sig_ak(2,2,j)
           enddo

        elseif ((.not. update_xx(j)) .and. update_pp(j)) then
           call die("cg_kernel_mbcg(): should not happen")
        endif
     enddo

     rho0(:,:) = rho(:,:)
     alp0(:,:) = alp(:,:)

     if (converged) goto 9999
  enddo
  !----------------------------------------

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

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

  TRACE_RETURN
end
