!===============================================================================
!
! dsf_poly.F90  --  p(j,x,mu) := p(j,x,mu) - step * D_{x,mu,j} S_{f}
!
!-------------------------------------------------------------------------------
!
! 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/>.
!
!-------------------------------------------------------------------------------
! Author: Taylor Haar
! Last modified: 03/04/2017
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------

subroutine dsf_poly(p, p_qed, p_axion, step, calc_sf, pf)
  ! Change the momentum by the force due to the PF action term
  ! belonging to the pseudofermion 'pf'
  !
  ! Constructs dS_poly/dU from eta_i and chi_i, where
  !    dS_poly/dU = sum(i=1, n) eta_i^t dK/dU chi_i, K = MtM
  ! then uses dsf() to calculate the corresponding change in p
  !
  ! For each i, the action part to calculate is
  !   eta_i^t dK/dU chi_i = [M eta_i]^t dM/dU chi_i + eta_i^t dMt/dU [M chi_i]
  ! Note that the second term is NOT the hermitian conjugate of the first.
  ! Using dS/dU = [dS/dU]^t, it can be shown that this is equal to
  !   1/2 {[M eta_i]^t dM/dU chi_i + [M chi_i]^t dM/dU eta_i + h.c.}
  ! One can take advantage of the symmetry of the polynomial to show
  ! that the two terms' sums over n are equal, so
  !    dS_poly/dU = sum(i=1, n) {[M eta_i]^t dM/dU chi_i +  h.c.}
  ! and only 'n' calls of dsf are required.
  !
  ! A further performance improvement involves splitting the calls to
  ! dsf_w2_gen() and dsf_sum() in dsf(), such that the former only
  ! needs to be called once rather than for each shift.
  
  use module_action
  use module_field
  use module_hmc_forces
  use module_p_interface
  use module_switches
  use module_vol
  use module_poly 
  use module_input
  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

  ! Have to use tmp_poly to distinguish from wdagw's 'tmp'
  P_SPINCOL_FIELD, save          :: tmp_poly
  type(type_chi), pointer        :: chi_i(:)
  type(type_chi), pointer        :: eta_i(:)
  REAL                           :: poly_step
  integer                        :: pid, mid
  integer                        :: i, j

  ! Intermediate fields for the dsf() splitting
  COMPLEX, allocatable, dimension(:,:,:,:,:) :: w, w1, w2
  integer :: ierr

  REAL, external                 :: dotprod

  if (switches%quenched) return

  TIMING_START(timing_bin_dsf_poly)

  call init_quark(pf%mid1)

  ALLOCATE_SC_FIELD(tmp_poly)

  pid = pf%polyid

  if (calc_sf /= 0) then
     ! Going to calculate Sf
     pf%sf = 0
  endif
  
  if (step /= ZERO) then
     ! Going to calculate dSf/dU
     ! Allocate storage for intermediate fields
     allocate(chi_i(poly(pid)%n))
     allocate(eta_i(poly(pid)%n))
     do j = 1, poly(pid)%n
        nullify(chi_i(j)%sc)
        nullify(eta_i(j)%sc)
        ALLOCATE_SC_FIELD(chi_i(j)%sc)
        ALLOCATE_SC_FIELD(eta_i(j)%sc)
     end do

     ! Allocate storage for intermediate force terms (for dsf splitting)
     allocate(w(NCOL, NCOL, volh_tot, EVEN:ODD, DIM))
     allocate(w1(NCOL, NCOL, volh_tot, EVEN:ODD, DIM))
     allocate(w2(NCOL, NCOL, volh_tot, EVEN:ODD, DIM))

     w = ZERO
     w1 = ZERO
     w2 = ZERO
  endif

  ! Correct the step-size's sign based on the type of term
  ! dsf() calculates something like '-dS/dU', so the signs are switched
  if (poly(pid)%type == P_CORRECTION) then
     ! Correction term => -ve change in momentum (like the non-filtered case)
     poly_step = step
  else
     ! Filter term => +ve change in momentum
     poly_step = -step
  end if

  do i = 1, pf%npf 
     ! If this pseudofermion has action term 'Sf = phi^t (PK)^-1 phi'
     if (poly(pid)%type == P_CORRECTION) then
        if (input%solver_cg_poly == 0) then
           ! Method 0 (default):
           ! Expand the inverse polynomial in the force term
           ! into partial fractions, then construct the force term
           ! from the shifted inverses (K - z_i)^-1
           
           ! Calculate required shifts (stored at sco(1,i)%sco(1))
           call gen_shifts(pf%mpf(i)%sc, pf, calc_sf, .true.)
           
           if (calc_sf /= 0) then
              ! Construct sf
              call inv_poly_shifts(tmp_poly, poly(pid)%n, &
                   sco(1,1:poly(pid)%n), pf%mpf(i)%sc, pf)
              pf%sf = pf%sf + dotprod(pf%mpf(i)%sc, tmp_poly, SIZE_SC_FIELD)
           end if
           
        elseif (input%solver_cg_poly == 1) then
           ! Method 1:
           ! Evaluate the inverse [PK]^-1 phi in full, then
           ! construct the action term by expanding the derivative on P(K)K
           
           ! Use the full cg solver
           call cg_poly(tmp_poly, pf%mpf(i)%sc, pid, pf%pid, calc_sf)
           if (calc_sf /= 0) then
              ! Construct sf
              pf%sf = pf%sf + dotprod(pf%mpf(i)%sc, tmp_poly, SIZE_SC_FIELD)
           end if
           
        elseif (input%solver_cg_poly == 2) then
           ! Method 2:
           ! Same as method 0, except that P(K) and K are inverted separately
           
           ! Only generate non-zero shifts
           call gen_shifts(pf%mpf(i)%sc, pf, calc_sf, .false.)
           
           ! Utilize unused sco storage to calculate
           !    tmp_poly = K^-1 P^-1 phi,
           ! which is required later on for dsf
           
           ! sco = P^-1(K) phi
           call inv_poly_shifts(sco(1, poly(pid)%n/2+1)%sco, poly(pid)%n, &
                sco(1,1:poly(pid)%n), pf%mpf(i)%sc, pf)
           
           ! tmp_poly = K^-1 sco
           TIMING_START(timing_bin_cg_mre_mtmp)
           call cg_mre_mtmp(tmp_poly, sco(1, poly(pid)%n/2+1)%sco, pf%mid1, &
                pf%pid, calc_sf)
           TIMING_STOP(timing_bin_cg_mre_mtmp)
           
           if (calc_sf /= 0) then
              pf%sf = pf%sf + dotprod(pf%mpf(i)%sc, tmp_poly, SIZE_SC_FIELD)
           endif
        end if

     ! Otherwise, we have an action term like 'Sf = phi^t P phi'
     else
        if (calc_sf /= 0) then
           ! Construct sf
           call poly_mult(tmp_poly, pf%mpf(i)%sc, pf)
           pf%sf = pf%sf + dotprod(pf%mpf(i)%sc, tmp_poly, SIZE_SC_FIELD)
        end if
     end if
     
     ! Having calculated sf if required, cycle if zero step-size
     if (step == ZERO) cycle
     
     ! Generate required chi_i and eta_i
     TIMING_START(timing_bin_gen_chi_eta)
     if (poly(pid)%type == P_CORRECTION) then
        if (input%solver_cg_poly == 0) then
           call gen_chi_eta(chi_i, eta_i, sco(1,1:poly(pid)%n), pid)
        elseif (input%solver_cg_poly == 1 .or. input%solver_cg_poly == 2) then
           ! Recall tmp_poly = [PK]^-1 phi
           call gen_chi_eta_poly(chi_i, eta_i, tmp_poly, pid, pf%mid1)
        end if
     else
        call gen_chi_eta_poly(chi_i, eta_i, pf%mpf(i)%sc, pid, pf%mid1)
     end if
     TIMING_STOP(timing_bin_gen_chi_eta)
     
     ! Add the corresponding terms to the intermediate force terms (w, w1, w2)
     ! (see comments at the start for an explanation)
     do j = 1, poly(pid)%n
        call wmul_r8(tmp_poly, eta_i(j)%sc, pf%mid1)
        TIMING_START(timing_bin_dsf)
        call dsf_sum(p_axion, w, w1, w2, poly_step, chi_i(j)%sc(1,1,1), &
             tmp_poly(1,1,1), pf%mid1)
        TIMING_STOP(timing_bin_dsf)
     enddo
  enddo

  ! Exit if zero step-size -> no dsf, only (possibly) sf calc
  if (step == ZERO) return

  ! Reset HMC forces
  call hmc_forces_old(p)
  
  ! Having totalled up dSf via 'dsf_sum()', we now call dsf_w2gen()
  ! to advance the momenta (p, p_qed)
  TIMING_START(timing_bin_dsf)
  call dsf_w2gen(p, p_qed, w, w1, w2, pf%mid1)
  TIMING_STOP(timing_bin_dsf)
  
  ! Update HMC forces
  call hmc_forces_new(p, step, pf%fid)
  
  deallocate(w, w1, w2)
  
  ! Deallocate storage for intermediate fields,
  ! making sure to go down to the actual data
  do j = 1, poly(pid)%n
     deallocate(chi_i(j)%sc)
     deallocate(eta_i(j)%sc)
  end do
  deallocate(chi_i)
  deallocate(eta_i)
  
  TIMING_STOP(timing_bin_dsf_poly)

end subroutine dsf_poly

!-------------------------------------------------------------------------------
subroutine gen_chi_eta(chi, eta, inv_shift, pid)
  ! Use the given inv_shift(i) = (K - z_i)^-1 phi to calculate chi and eta
  ! terms for the polynomial correction term.
  !
  ! pid is the corresponding polynomial P(pid).
  !
  ! The derivative of P(K)^-1 can be expressed as
  !     d(P(K)^-1)/dU = - sum(i=1, n) [p(i) (K - z_i)^-1 dK/dU (K - z_i)^-1]
  ! where the complex number p(i) is given by
  !     p(i) = 1/c prod(j=1, n, j != i) [1/(z_i - z_j)].
  !
  ! Hence, we prescribe
  !     chi(i) = p(i) (K - z_i)^-1 phi
  ! and
  !     eta(i) = (K - z_i*)^-1 phi,
  ! such that we can write
  !     dS/dU = - sum(i=1, n) [eta_i^dag dK/dU chi_i]
  ! The negative sign here is accounted for in dsf_poly().

  use module_poly
  use module_cg
  use module_vol
  use typedef
  implicit none

  integer, intent(in) :: pid
  type(type_chi), dimension(poly(pid)%n), intent(inout) :: chi, eta
  type(sco_field), intent(in) :: inv_shift(poly(pid)%n)

  COMPLEX, dimension(poly(pid)%n) :: p_i
  integer :: i, j, n
  
  n = poly(pid)%n
  
  ! First, construct p(i) = 1/c_n prod(j=1, n, j!=i) [1/(z_i - z_j)]
  do i = 1, n
     p_i(i) = ONE / poly(pid)%c
     do j = 1, n
        if (j /= i) then
           p_i(i) = p_i(i) / (poly(pid)%z_k(i) - poly(pid)%z_k(j))
        endif
     enddo
  enddo

  ! chi_i = - p_i /(K - z_i) phi = -p_i * i-th inverse solution
  do i = 1, n
     call sc_copy(chi(i)%sc(1,1,1), inv_shift(i)%sco(1))
     call sc_cscale(chi(i)%sc(1,1,1), p_i(i))
  enddo

  ! Use the fact that z_i = z_(n-i+1)* to calculate eta_i as
  !    eta_i = 1/(K - z_(n-i+1)) phi
  do i = 1, n
     j = n - i + 1
     call sc_copy(eta(i)%sc(1,1,1), inv_shift(j)%sco(1))
  enddo
  
end subroutine gen_chi_eta

!-------------------------------------------------------------------------------
subroutine gen_chi_eta_poly(chi, eta, phi, pid, mid)
  ! Calculate the required chi, eta in the polynomial filter case
  ! and in the correction term case when inverting [P(K)K] in full
  ! (in which case, 'phi' = [P(K)K]^-1 phi)
  !
  ! pid is the corresponding polynomial P(pid).
  !
  ! The derivative of 'S = phi^dag P(K) phi' can be expressed as
  !     dS/dU = sum(i=1, n) [eta_i^dag dK/dU chi_i]
  ! where we prescribe
  !     chi(i) = prod(j=i+1, n) [(K - z_i)] phi
  ! and
  !     eta(i) = c prod(j=1, i-1) [(K - z_i*)] phi.

  use module_poly
  use module_cg
  use module_vol
  use module_field
  use module_p_interface
  use typedef
  implicit none

  integer, intent(in) :: pid, mid
  type(type_chi), dimension(poly(pid)%n), intent(inout) :: chi, eta
  SPINCOL_FIELD, intent(in) :: phi

  P_SPINCOL_FIELD, save :: tmp_poly
  integer :: j, n
  
  n = poly(pid)%n

  ALLOCATE_SC_FIELD(tmp_poly)

  ! Start with chi_n = phi
  call sc_copy(chi(n)%sc(1,1,1), phi)

  ! Construct chi's recursively, using chi_(j-1) = (K - z_j) chi_j
  do j = n, 2, -1
     call wdagw(tmp_poly, chi(j)%sc, mid)
     call sc_cax2(chi(j-1)%sc(1,1,1), tmp_poly, (ONE, ZERO), &
          chi(j)%sc(1,1,1), -poly(pid)%z_k(j))
  end do

  ! For eta's, start with eta_1 = c phi
  call sc_copy(eta(1)%sc(1,1,1), phi)
  call sc_scale(eta(1)%sc(1,1,1), poly(pid)%c)

  ! Construct eta's recursively, using eta_j = (K - z_(j-1)*) eta_(j-1)
  do j = 2, n
     call wdagw(tmp_poly, eta(j-1)%sc, mid)
     call sc_cax2(eta(j)%sc(1,1,1), tmp_poly, (ONE, ZERO), &
          eta(j-1)%sc(1,1,1), -conjg(poly(pid)%z_k(j-1)))
  enddo
  
end subroutine gen_chi_eta_poly
