!===============================================================================
!
! axion.F90 - operations on axion field
!
!-------------------------------------------------------------------------------
!
! Copyright 201y 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 init_axion(a, start)  ! initialization of axion field (at trajectory 0)

  use module_vol
  use module_switches
  implicit none

  AXION_FIELD :: a
  integer :: start, i, eo

  if (.not. switches%axion) return

  select case (start)
     case (START_HOT)

        do eo = EVEN, ODD
           call ran_vector_volh(a(1, eo), eo)
        enddo

     case (START_COLD)

        do eo = EVEN, ODD
           !$omp parallel do
           do i = 1, volh
              a(i, eo) = ZERO
           enddo
        enddo

     case default
        call die("init_axion(): don't know how to start")
  end select

  call xbound_axion_field(a)

end

!-------------------------------------------------------------------------------
subroutine axion_conf_copy(out, in)

  use module_vol
  use module_switches
  implicit none
  AXION_FIELD, intent(out) :: out
  AXION_FIELD, intent(in)  :: in
  integer :: i, eo
  
  if (.not. switches%axion) return

  do eo = EVEN, ODD
     !$omp parallel do
     do i = 1, volh_tot
        out(i, eo) = in(i, eo)
     enddo
  enddo

end

!-------------------------------------------------------------------------------
REAL function sk_axion(a, energy, average)  ! returns S_axion and average 

  use module_nn
  use module_vol
  use module_switches
  use module_function_decl
  implicit none

  AXION_FIELD, intent(in) :: a
  REAL, intent(out)      :: energy, average
  
  REAL :: sum
  integer :: i, e, o, j1, mu

  if (.not. switches%axion) then
     sk_axion = ZERO
     energy = ZERO
     average = ZERO
     return
  endif

  sum = ZERO

  do mu = 1, DIM
     do e = EVEN, ODD
        o = EVEN + ODD - e
        !$omp parallel do reduction(+: sum) private(j1)
        do i = 1, volh
           j1 = nn(i, e, mu, FWD)
           sum = sum + (a(i, e) - a(j1, o) ) * a(i, e)
        enddo
     enddo
  enddo
              
  sum = global_sum(sum)

  sk_axion = sum
  energy = sk_axion / (FOUR * VOLUME)


  sum = ZERO
  do e = EVEN, ODD
     !$omp parallel do reduction(+: sum)
     do i = 1, volh
        sum = sum + a(i, e)
     enddo
  enddo
  average = global_sum(sum) / dble(VOLUME)

end

!-------------------------------------------------------------------------------
REAL function sp_axion(p)  ! returns action of momenta p_axion

  use module_vol
  use module_function_decl
  use module_switches
  implicit none
  AXION_MOMENTUM_FIELD, intent(in) :: p
  
  integer :: eo
  
  if (.not. switches%axion) then
     sp_axion = ZERO
     return
  endif

  sp_axion = ZERO
  do eo = EVEN, ODD
     sp_axion = sp_axion + dotprod(p(1, eo), p(1, eo), volh)
  enddo
  sp_axion = sp_axion * HALF

end

!-------------------------------------------------------------------------------
subroutine dsk_axion(p, a, step, beta)

  use module_nn
  use module_vol
  use module_switches
  implicit none

  AXION_MOMENTUM_FIELD, intent(inout) :: p
  AXION_FIELD, intent(in)             :: a
  REAL, intent(in)                    :: step, beta
  REAL                                :: s
  integer                             :: e, o, i, j1 ,j2, mu

  if (.not. switches%axion) return
  DEBUG2S("Start: dsk_axion")

  s = step * beta

  do mu = 1, DIM
     do e = EVEN, ODD
        o = EVEN + ODD - e
        !$omp parallel do private(j1, j2)
        do i = 1, volh
           j1 = nn(i, e, mu, FWD)
           j2 = nn(i, e, mu, BWD)
           p(i, e) = p(i, e) - s * ( TWO* a(i, e) - a(j1, o) -a(j2, o))
        enddo
     enddo
  enddo

  DEBUG2S("End: dsk_axion")
end

!-------------------------------------------------------------------------------
subroutine hmc_axion(p, a, step)
  use module_vol
  use module_switches
  implicit none

  AXION_MOMENTUM_FIELD, intent(in)    :: p
  AXION_FIELD, intent(inout) :: a
  REAL, intent(in)                  :: step
  integer                           :: i, eo

  if (.not. switches%axion) return

  DEBUG2S("Start: hmc_axion")

  do eo = EVEN, ODD
     !$omp parallel do
     do i = 1, VOLH
        a(i, eo) = a(i, eo) + p(i, eo) * step
     enddo
  enddo

  call xbound_axion_field(a)

  DEBUG2S("End: hmc_axion")
end

!-------------------------------------------------------------------------------
subroutine hmc_init_p_axion(p, scratch)

  use module_vol
  use module_switches
  implicit none
  AXION_MOMENTUM_FIELD, intent(out) :: p
  COMPLEX, intent(out)            :: scratch(volh)
  integer                         :: i

  if (.not. switches%axion) return

  call ran_gauss_volh(1, scratch, ONE, EVEN)
  !$omp parallel do
  do i = 1, volh
     p(i, EVEN) = real(scratch(i))
     p(i, ODD) = aimag(scratch(i))
  enddo

#ifdef ZERO_MOMENTUM
  p = ZERO
#endif

end
!-------------------------------------------------------------------------------
subroutine axion_force(p, step, b, a, eo)
  use module_vol
  use module_switches
  implicit none

  AXION_MOMENTUM_FIELD, intent(inout) :: p
  SPINCOL_FIELD,   intent(in)    :: a, b
  REAL, intent(in)               :: step
  integer, intent(in)            :: eo
  integer                        :: i

  if (.not. switches%axion) return

  DEBUG2S("Start: axion_force")

  call gamma5(a, volh)

  !$omp parallel do
  do i = 1, VOLH
     p(i, eo) = p(i, eo) + step * (a(1, 1, i) * dconjg(b(1, 1, i)) &
                                  +a(2, 1, i) * dconjg(b(2, 1, i)) &
                                  +a(3, 1, i) * dconjg(b(3, 1, i)) &
                                  +a(4, 1, i) * dconjg(b(4, 1, i)) &
                                  +a(1, 2, i) * dconjg(b(1, 2, i)) &
                                  +a(2, 2, i) * dconjg(b(2, 2, i)) &
                                  +a(3, 2, i) * dconjg(b(3, 2, i)) &
                                  +a(4, 2, i) * dconjg(b(4, 2, i)) &
                                  +a(1, 3, i) * dconjg(b(1, 3, i)) &
                                  +a(2, 3, i) * dconjg(b(2, 3, i)) &
                                  +a(3, 3, i) * dconjg(b(3, 3, i)) &
                                  +a(4, 3, i) * dconjg(b(4, 3, i)) )
  enddo

  call gamma5(a, volh)

  DEBUG2S("End: axion_force")
end

!-------------------------------------------------------------------------------
subroutine clover_dsd_axion(p, eo, step, t)
  use module_vol
  implicit none

  AXION_MOMENTUM_FIELD, intent(inout) :: p
  REAL,            intent(in)    :: step
  COMPLEX, dimension(NDIRAC, NDIRAC, NCOL, NCOL, volh) :: t
  integer, intent(in) :: eo
  integer :: i,c

  !$omp parallel do private(c)
  do i = 1, volh
  do c = 1, NCOL
#ifdef GAMMAC
     p(i, eo) = p(i, eo) + step *(dble(t(1, 1, c, c, i)) + dble(t(2, 2, c, c, i)) - dble(t(3, 3, c, c, i)) - dble(t(4, 4, c, c, i)))
#else
     p(i, eo) = p(i, eo) + step *(dble(t(1, 3, c, c, i)) + dble(t(2, 4, c, c, i)) + dble(t(3, 1, c, c, i)) + dble(t(4, 2, c, c, i)))
#endif
  enddo
  enddo

end

!-------------------------------------------------------------------------------
subroutine axion_sliced_av(traj)
  use module_field
  use module_vol
  use module_lattice
  use module_function_decl
  implicit none
  integer :: ix,iy,iz,it, j(4), coord(4), ii, eo, traj
  REAL :: slice(0:lt-1)
  character(len=50) ::fmt
  character(len=4) ::clt
  integer, external :: e_o, std_xyzt2i

  call pe2coord(my_pe(), coord)

  slice = 0

  !$omp parallel do private(it, iz, iy, ix, j, ii, eo) collapse(2)
  do it = 0, nt-1
  do iz = 0, nz-1
  do iy = 0, ny-1
  do ix = 0, nx-1
     j  = (/ix, iy, iz, it/)
     ii = std_xyzt2i(j)
     eo = e_o(j)
     slice(it + nt*coord(4)) = slice(it + nt*coord(4)) + axion_ptr(ii, eo)
  enddo
  enddo
  enddo
  enddo

  call global_sum_vec(lt, slice)
  slice(:) = slice(:) / dble(lx*ly*lz)

  write(clt, "(i4)")lt
  fmt = "(a, i10, 2x, "//trim(clt)//"(1x,E20.12))"

  if (my_pe()==0) then
     write(UREC,fmt)"  %axsa",traj, slice
  endif
!stop
end

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