!===============================================================================
!
! dsfr_z.F90  --  p(j,x,mu) := p(j,x,mu) - step * D_{x,mu,j} S_{fr}
!                 for rational fractions
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2006-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/>.
!
!-------------------------------------------------------------------------------
! Complex version for PF-RHMC and the Zolotarev approximation
!
! Editor: Taylor Haar
! Last modified: 21/03/2017
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------
subroutine dsfr_z(p, p_qed, p_axion, step, calc_sf, pf)
  use module_field
  use module_action
  use module_input
  use module_switches
  use module_vol
  use module_rhmc
  use module_hmc_forces
  implicit none

  GENERATOR_FIELD, intent(inout) :: p
  U1_GENERATOR_FIELD, intent(inout) :: p_qed
  AXION_MOMENTUM_FIELD, intent(inout) :: p_axion
  integer,         intent(in)    :: calc_sf
  REAL,            intent(in)    :: step
  type(type_pfermi)              :: pf

  COMPLEX, dimension(:,:,:,:,:), allocatable :: w, w1, w2

  if (switches%quenched) return
  DEBUG2S("Start: dsfr")
  call hmc_forces_old(p)

  allocate(w( NCOL, NCOL, volh_tot, EVEN:ODD, DIM)); w=ZERO
  allocate(w1(NCOL, NCOL, volh_tot, EVEN:ODD, DIM)); w1=ZERO
  allocate(w2(NCOL, NCOL, volh_tot, EVEN:ODD, DIM)); w2=ZERO

  if (switches%zolotarev .and.  pf%nth == 2) then
     call dsf_fraction_sep_zolo(pf, step, w, w1, w2, p_axion)
  else
     call dsf_fraction_sep_z(pf, step, w, w1, w2, p_axion)
  endif
  call dsf_w2gen(p, p_qed, w, w1, w2, pf%mid1)

  deallocate(w, w1, w2)

  call hmc_forces_new(p, step, pf%fid)
  DEBUG2S("End: dsfr")

end subroutine dsfr_z

!-------------------------------------------------------------------------------
subroutine dsf_fraction_sep_z(pf, step, w, w1, w2, p_axion)
  use module_action
  use module_vol
  use module_rhmc
  use module_field
  use module_input

  use module_poly
  use module_switches
  use module_function_decl
  implicit none

  type(type_pfermi) :: pf
  GAUGE_FIELD,    intent(inout) :: w, w1, w2
  AXION_MOMENTUM_FIELD, intent(inout) :: p_axion
  REAL,           intent(in)    :: step
  integer, allocatable :: frdo(:), iters(:)
  integer :: mid, rid, type, polyid
  integer :: i, j, k

  REAL :: bof(ratapp(pf%rid)%shift), factor, fe, poly_c
  COMPLEX :: shifts(ratapp(pf%rid)%shift + poly(pf%polyid)%n)
  COMPLEX :: residue, rat_factor
  REAL, allocatable :: tmp_a(:), tmp_b(:), phi(:)
  integer :: m, n, conj_id

  mid = pf%mid1
  rid = pf%rid
  type= pf%type
  polyid = pf%polyid
  call init_quark(mid)

  factor=pf%factor
  fe    =factor**ratapp(rid)%exponents(1) / factor
  bof   =ratapp(rid)%bbar/factor

  allocate( tmp_a(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot))
  allocate( tmp_b(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot))
  allocate( phi(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot))

  do j = 1, pf%npf
     call sc_copy(phi, pf%mpf(j)%sc)

     m = ratapp(pf%rid)%shift
     if (polyid > 0) then
        n = poly(polyid)%n
        poly_c = poly(polyid)%c
     else
        ! No PF, so set to ensure 'no PF' behaviour
        n = 0
        poly_c = 1
     endif

     ! Consolidate shifts into a single vector
     shifts(1:m) = bof(1:m)
     if (polyid > 0) then
        shifts(m+1:m+n) = -poly(polyid)%z_k
     endif

     ! Calculate all the shifted inverses in one go,
     ! saving them at sco(1,i)%sco
     allocate(frdo(m+n), iters(m+n))
     frdo = 0
     call mcg_outer_z(phi, mid, m+n, frdo, shifts, iters, pf%pid, CG_MD)

     do i = 1, m
        ! Determine required residue at the rational pole
        residue = fe * ratapp(rid)%abar(i) / poly_c
        do k = 1, n
           residue = residue / (- bof(i) - poly(polyid)%z_k(k))
        enddo

        ! A = [K + b_j] phi; B = res [K + b_j] phi
        ! Residue should be real by construction, so no conjugation necessary
        call wmul_r8(tmp_b, sco(1,i)%sco(1), mid)
        call sc_cscale(tmp_b, residue)

        call dsf_sum(p_axion, w, w1, w2, step, sco(1,i)%sco(1), tmp_b, mid)
     enddo

     do i = 1, n
        ! Determine required residue at the polynomial pole
        residue = ONE / poly_c
        do k = 1, n
           if (k /= i) then
              residue = residue / (poly(polyid)%z_k(i) - poly(polyid)%z_k(k))
           endif
        enddo
        rat_factor = fe * ratapp(rid)%abar(0) * factor
        do k = 1, m
           rat_factor = rat_factor + (fe * ratapp(rid)%abar(k)) / &
                (poly(polyid)%z_k(i) + bof(k))
        enddo
        residue = residue * rat_factor

        ! A = res [K - z_j] phi
        call sc_copy(tmp_a, sco(1,m+i)%sco(1))
        call sc_cscale(tmp_a, residue)

        ! B = W [K - z_j*]^-1 phi = W [K - z_(n-j+1)]^-1 phi
        conj_id = m + (n - i + 1)
        call wmul_r8(tmp_b, sco(1, conj_id)%sco(1), mid)

        call dsf_sum(p_axion, w, w1, w2, step, tmp_a, tmp_b, mid)
     enddo
     deallocate(frdo, iters)
  enddo

  deallocate(tmp_a, tmp_b, phi)

end subroutine dsf_fraction_sep_z

!-------------------------------------------------------------------------------
subroutine dsf_fraction_sep_zolo(pf, step, w, w1, w2, p_axion)
  ! Version for Zolotarev rationals

  use module_action
  use module_vol
  use module_rhmc
  use module_field
  use module_input

  use module_rhmc_zolo
  use module_truncrat
  use module_poly
  use module_switches
  use module_function_decl
  implicit none

  type(type_pfermi) :: pf
  GAUGE_FIELD,    intent(inout) :: w, w1, w2
  AXION_MOMENTUM_FIELD, intent(inout) :: p_axion
  REAL,           intent(in)    :: step

  integer, allocatable :: frdo(:), iters(:)
  integer :: mid, rid, type, polyid
  integer :: i, j, k

  COMPLEX, allocatable :: shifts(:)
  COMPLEX :: residue
  REAL, allocatable :: tmp_a(:), tmp_b(:), phi(:)
  integer :: conj_id

  ! For polynomial part
  integer :: n
  REAL :: poly_c

  ! For ratapp part, with possible truncations
  integer :: m, i_small, i_large, nshifts
  REAL :: bof(ratapp_zolo(pf%rid)%shift), cof(ratapp_zolo(pf%rid)%shift)
  REAL :: factor, fe

  mid = pf%mid1
  rid = pf%rid
  type= pf%type
  polyid = pf%polyid
  nshifts = ratapp_zolo(rid)%shift
  call init_quark(mid)

  !! Scale approximation to determined eigenvalue spectrum
  !+ see fermi/rhmc/init_rhmc:gene_new_coeffs() for more info
  factor = pf%factor
  fe = factor**ratapp_zolo(rid)%exponent
  bof = ratapp_zolo(rid)%b_k/factor
  cof = ratapp_zolo(rid)%c_k/factor

  allocate(tmp_a(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot))
  allocate(tmp_b(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot))
  allocate(phi(SIZE_COMPLEX*NDIRAC*NCOL*volh_tot))

  do j = 1, pf%npf
     call sc_copy(phi, pf%mpf(j)%sc)

     ! Set PF params to defaults (= no PF)
     n = 0
     poly_c = 1

     if (pf%truncid > 0) then
        ! This ratapp is truncated. Account for this in the shifts.
        ! Remember to reverse indexing as well.
        i_large = nshifts - truncrat(pf%truncid)%n_low + 1
        i_small = modulo(nshifts - truncrat(pf%truncid)%n_high, nshifts) + 1
        m = i_large - i_small + 1
     else
        i_small = 1
        i_large = ratapp_zolo(rid)%shift
        m = i_large
     endif

     ! If a polynomial filter is present, add it in
     if (polyid > 0) then
        n = poly(polyid)%n
        poly_c = poly(polyid)%c
     endif

     ! Consolidate shifts into a single vector
     allocate(shifts(m+n))

     shifts(1:m) = bof(i_small:i_large)
     if (polyid > 0) then
        shifts(m+1:m+n) = -poly(polyid)%z_k
     endif

     ! Calculate all the shifted inverses in one go,
     ! saving them at sco(1,i)%sco
     allocate(frdo(m+n), iters(m+n))
     frdo = 0
     call mcg_outer_z(phi, mid, m+n, frdo, shifts, iters, pf%pid, CG_MD)


     do i = 1, m+n
        ! Determine required residue at the current pole
        !
        ! R(K) / P(K) = deg(m) poly / deg(m+n) poly, so the same code can be
        ! used at each pole.

        ! Only include ratapp normalization for largest shifts
        if (i_large == ratapp_zolo(rid)%shift) then
           residue = fe * ratapp_zolo(rid)%d_n / poly_c
        else
           residue = ONE / poly_c
        endif

        do k = 1, m+n
           if (k /= i) then
              residue = residue / (- shifts(i) + shifts(k))
           endif
           if (k <= m) then
              ! Add in numerator parts as well = prod_k(z + c_k)
              residue = residue * (- shifts(i) + cof(k+i_small-1))
           endif
        enddo

        ! Call to dsf gives F = Bt dW/dU A + h.c., where
        !  A = res [K + shift_j]^-1 phi
        !  B = W [K + shift_j*]^-1 phi
        ! shift_j* is equal to one of the other shifts, j' = conj_id

        call sc_copy(tmp_a, sco(1,i)%sco(1))
        call sc_cscale(tmp_a, residue)

        if (i <= m) then
           ! Rational root: real number => own conjugate
           conj_id = i
        else
           ! Polynomial root: roots come in complex conjugate pairs
           conj_id = m + n - (i - m)  + 1
        endif
        call wmul_r8(tmp_b, sco(1, conj_id)%sco(1), mid)

        call dsf_sum(p_axion, w, w1, w2, step, tmp_a, tmp_b, mid)
     enddo

     deallocate(frdo, iters, shifts)
  enddo

  deallocate(tmp_a, tmp_b, phi)

end subroutine dsf_fraction_sep_zolo

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