!===============================================================================
!
! cg_kernel_mcg.F90 - multi_shift_cg kernel of BQCDi
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2006-2010 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"

!-------------------------------------------------------------------------------
# define TRACE_NAME T(cg_kernel_mcg)
!-------------------------------------------------------------------------------
subroutine T(cg_kernel_mcg)(matrix, id, n, shift_switch, betabar, b, iterations, niter_shift, res) 

  use module_field  ! contains sco
  use module_input
  use module_cg
  use module_cg_arrays
  use module_cg_procs, exclude => T(cg_kernel_mcg)  ! avoid name clash
  use T(module_vol)
  use typedef_simd
  implicit none

  external                   :: matrix  ! not used here
  integer,       intent(in)  :: id
  integer,       intent(in)  :: n
  integer,       intent(in)  :: shift_switch(n)
  REAL,          intent(in)  :: betabar(n)
  spincol_FIELD, intent(in)  :: b
  integer,       intent(out) :: iterations
  integer,       intent(out) :: niter_shift(n)
  DOUBLE,        intent(in)  :: res(n)

  DOUBLE        :: alpha, alpha_old, beta, rtr, rtrold, paap, bb
  DOUBLE        :: sbeta(n), salpha(n)
  REAL          :: sxi_old(n), sxi(n), sxi_new(n)
  integer       :: i, j, niter, n_min, n_con(n), niter_sub, ncon
  integer       :: i1, i2, block_len
  character(72) :: msg
  logical       :: converged, update_1(n), update_2(n)

  TRACE_ENTRY

  if (input%tuning_cg_mcg_block_length > 0) then
     CHECK(mod(volh, input%tuning_cg_mcg_block_length) == 0)
  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

  ! sco is initialised in cg_wrapper_mcg()

  call sc_zero(cg_arrays%r)
  call sc_zero_halo(cg_arrays%r)

  call cg_procs%loop1(rtrold, cg_arrays%r, cg_arrays%p, b, volh_tot, volh)

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

  alpha_old  = ONE
  beta       = ZERO
  sxi_old(:) = ONE
  sxi(:)     = ONE
  bb         = rtrold
  niter_sub  = 0
  converged  = .false.
  !----------------------------------------
  do niter = 1, cg_para%maxiter
     call cg_procs%ww(cg_arrays%aap, cg_arrays%p, id)
     call cg_procs%loop2(paap, cg_arrays%p, cg_arrays%aap, volh_tot, volh)

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

     alpha = -rtrold / paap

     call T(cg_mcg_loop3)(rtr, cg_arrays%r, cg_arrays%aap, -alpha, volh_tot, volh)

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

     do j = 1, n
        if (n_con(j) == 0) then
           update_1(j) = .true.
           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(j) = alpha * sxi_new(j) / sxi(j)
           niter_sub = niter_sub + 1

           if (stop_condition(rtr * sxi_new(j)**2, res(j), bb)) then
              niter_shift(j) = niter
              n_con(j) = 1
           endif
        else
           update_1(j) = .false.
        endif
     enddo

     if (stop_condition(rtr * sxi_new(n_min)**2, res(n_min), bb)) then
        ncon = 0
        do j = n_min, n
           ncon = ncon + n_con(j)
        enddo
        if (ncon == n - n_min + 1) converged = .true.
     endif

     if (converged) then
        update_2(:) = .false.
     else
        beta = rtr / rtrold
        rtrold = rtr

        call cg_procs%xpby(cg_arrays%p, cg_arrays%r, beta, volh_tot, volh)  ! p = r + beta * p

        do j = 1, n
           if (n_con(j) == 0) then
              update_2(j) = .true.
              sbeta(j) = beta * (sxi_new(j) / sxi(j))**2
              sxi_old(j) = sxi(j)
              sxi(j) = sxi_new(j)
           else
              update_2(j) = .false.
           endif
        enddo

        alpha_old = alpha
     endif

     block_len = input%tuning_cg_mcg_block_length

     if (block_len > 0) then
        !$omp parallel do schedule(static, 1) &
        !$omp default(none) &
        !$omp shared(volh, volh_tot, block_len, n, update_1, update_2, sco, cg_arrays, salpha, sbeta, sxi_new) &
        !$omp private(i1, i2, j)
        do i1 = 1, volh, block_len
           i2 = i1 + block_len - 1
           do j = 1, n
              if (update_1(j) .and. update_2(j)) then
                 call T(cg_mcg_update_sco_block)(sco(1,j)%sco(1), sco(2,j)%sco(1), cg_arrays%r, salpha(j), sbeta(j), sxi_new(j), volh_tot, i1, i2)
              else
                 if (update_1(j)) then
                    call T(cg_mcg_update_sco1_block)(sco(1,j)%sco(1), sco(2,j)%sco(1), salpha(j), volh_tot, i1, i2)
                 endif
                 if (update_2(j)) then
                    call T(cg_mcg_update_sco2_block)(sco(2,j)%sco(1), cg_arrays%r, sbeta(j), sxi_new(j), volh_tot, i1, i2)
                 endif
              endif
           enddo
        enddo

     else

        do j = 1, n
           if (update_1(j) .and. update_2(j)) then
              call T(cg_mcg_update_sco)(sco(1,j)%sco(1), sco(2,j)%sco(1), cg_arrays%r, salpha(j), sbeta(j), sxi_new(j), volh_tot, volh)
           else
              if (update_1(j)) then
                 call T(cg_mcg_update_sco1)(sco(1,j)%sco(1), sco(2,j)%sco(1), salpha(j), volh_tot, volh)
              endif
              if (update_2(j)) then
                 call T(cg_mcg_update_sco2)(sco(2,j)%sco(1), cg_arrays%r, sbeta(j), sxi_new(j), volh_tot, volh)
              endif
           endif
        enddo

     endif

     if (converged) goto 9999
  enddo
  !----------------------------------------
  niter = niter - 1

  if (cg_para%log /= 2) then
     write(msg, *) "cg_kernel_mcg(): no convergence; rtr = ", rtr 
     call die(msg)
  endif

9999 iterations = niter

  call update_cg_stat(cgm_stat, niter, niter_sub)

  TRACE_RETURN

end

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