!===============================================================================
!
! wilson_flow.F90 - measurement of the wilson flow
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2012, 2016 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"


#ifndef INCLUDE_MPIF_H
#define INCLUDE_MPIF_H include 'mpif.h'
#endif

!-------------------------------------------------------------------------------
module module_wilson_flow
  implicit none

  integer, save :: init = 0
  integer, save :: traj_qcd, traj_qed

  REAL, dimension(:, :), pointer, save :: q_qcd ! topological density of QCD
  REAL, dimension(:, :), pointer, save :: q_qed ! topological density of QED
  REAL, dimension(:, :, :, :), pointer, save :: qq_qcd ! topological density of QCD
  REAL, dimension(:, :, :, :), pointer, save :: qq_qed ! topological density of QED

end module module_wilson_flow

!-------------------------------------------------------------------------------
subroutine init_module_wilson_flow()
  use module_wilson_flow
  use module_lattice
  use module_vol
  implicit none

  if (init /= 0) return

  allocate(q_qcd(volh,0:1))
  allocate(q_qed(volh,0:1))

  allocate(qq_qcd(0:nx-1, 0:ny-1, 0:nz-1, 0:nt-1))
  allocate(qq_qed(0:nx-1, 0:ny-1, 0:nz-1, 0:nt-1))

  init = 1

end

!-------------------------------------------------------------------------------
subroutine topological_density_dump(char, istep)
  use module_wilson_flow
  use module_lattice
  use module_counter
  use module_decomp
  use module_function_decl
  implicit none

  character(len = 3), intent(in) :: char
  integer, intent(in)            :: istep
  character(len = 30)            :: filename, filename_t

  integer            :: x, y, z, t, z_global, t_global, size
  integer            :: pe_x, pe_y, pe_z, pe_t
  integer            :: pe, pe_io
  integer            :: count, block_length, stride, buf_type
  logical            :: io_pe
  integer, external  :: i_global


  integer :: ix, iy, iz, it, i, e, j(4), ierror
  integer, external :: e_o, std_xyzt2i
  REAL :: buffer(0:LX-1, 0:LY-1)
  REAL :: field(0:nx-1, 0:ny-1, 0:nz-1, 0:nt-1)
  integer            :: i_pe(DIM)
  INCLUDE_MPIF_H

  if (counter%run <= 999) then
     write(filename,233)char, counter%run, counter%traj, istep
  else
     write(filename,266)char, counter%run, counter%traj, istep
  endif
  233 format ("q_",a3,".",i3.3,".",i6.6,".",i4.4)
  266 format ("q_",a3,".",i6.6,".",i6.6,".",i4.4)
  open(200, file = filename, action = "write")

  if (char == "QCD") then
     !$omp parallel do private(ix, iy, iz, it, j, i, e)
     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/)
        i  = std_xyzt2i(j)
        e  = e_o(j)
        qq_qcd(ix,iy,iz,it) = q_qcd(i,e)
     enddo
     enddo
     enddo
     enddo
   
     do it = 0, nt-1
     do iz = 0, nz-1
     do iy = 0, ny-1
     do ix = 0, nx-1
        write(200,'(4i5, e20.10)')ix,iy,iz,it,qq_qcd(ix,iy,iz,it)
     enddo
     enddo
     enddo
     enddo

     field=qq_qcd

  else
     !$omp parallel do private(ix, iy, iz, it, j, i, e)
     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/)
        i  = std_xyzt2i(j)
        e  = e_o(j)
        qq_qed(ix,iy,iz,it) = q_qed(i,e)
     enddo
     enddo
     enddo
     enddo
   
     do it = 0, nt-1
     do iz = 0, nz-1
     do iy = 0, ny-1
     do ix = 0, nx-1
        write(200,'(4i5, e20.10)')ix,iy,iz,it,qq_qed(ix,iy,iz,it)
     enddo
     enddo
     enddo
     enddo

     field=qq_qed
  endif

  close(200)


  count        = NY
  block_length = NX
  stride       = block_length * NPE(1)
  size         = block_length * NY               ! words in send/recv

  i_pe         = decomp%std%i_pe                 ! pe coordinates
  io_pe        = (i_pe(1) == 0 .and. i_pe(2) == 0 .and. i_pe(3) == 0) ! am I doing IO?
  call mpi_type_vector(count, block_length, stride, BQCD_REAL, buf_type, ierror)
  call mpi_type_commit(buf_type, ierror)

  pe_t = i_pe(4)
  do t = 0, NT - 1
     t_global = i_global(t, NT, i_pe(4))

     write(filename_t,'(a,".",i2.2)')trim(filename), t_global
     open(210, file = filename_t, action = "write")

     do pe_z = 0, NPE(3) - 1
     do z = 0, NZ - 1
        z_global = i_global(z, NZ, pe_z)

        do pe_y = 0, NPE(2) - 1
        do pe_x = 0, NPE(1) - 1
           y = count * pe_y
           x = block_length * pe_x
           call field_io_pes(pe, pe_io, (/pe_x, pe_y, pe_z, pe_t/))
           if (io_pe) then
              CHECK(my_pe() == pe_io)
           endif      

           if (my_pe() == pe .and. my_pe() == pe_io) then
              call field_io_seq("write", count, block_length, stride, field(0,0,z,t), buffer(x,y))
           else
              if (my_pe() == pe_io) then
                 call mpi_recv(buffer(x,y), 1, buf_type, pe, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE, ierror)
              endif
              if (my_pe() == pe) then
                 call mpi_ssend(field(0,0,z,t), size, BQCD_REAL, pe_io, 0, MPI_COMM_WORLD, ierror)
              endif
           endif

        enddo
        enddo

        if (io_pe) then
        do y = 0, ly-1
        do x = 0, lx-1
            write(210,'(4i5, e20.10)')x,y, z_global, t_global, buffer(x,y)
        enddo
        enddo
        endif
     enddo
     enddo

     close(210)
  enddo

end

!-------------------------------------------------------------------------------
!
! Wilson flow for QCD filed
!
!-------------------------------------------------------------------------------
subroutine wilson_flow(u_in, traj, i_ensemble1, i_ensemble2)
  use module_wilson_flow
  use module_input
  use module_function_decl
  use module_vol
  implicit none
  integer, intent(in)     :: traj, i_ensemble1, i_ensemble2
  GAUGE_FIELD, intent(in) :: u_in
  GAUGE_FIELD             :: u
  integer                 :: i, nstep
  REAL                    :: clov, t, epsilon, plaq, qq, clovtt_old
  logical, save :: fileopen = .false.
  character(len = 14) :: filename

  call init_module_wilson_flow()
  traj_qcd = traj
  nstep = input%measure_wilson_flow_step
  read(input%measure_wilson_flow_eps,*)epsilon

  u = u_in

  if (.not. fileopen) then
     if (len(trim(input%measure_wilson_flow_file)) > 0) then
        if (my_pe()==0) open(40, file = trim(input%measure_wilson_flow_file), action = "write")
     else
        write(filename, "(a, i6.6)" ) "wilflow.", traj
        if (my_pe()==0) open(40, file = trim(filename), action = "write")
     endif
     fileopen = .true.
  endif

  clovtt_old = ZERO
  call clover_energy(clov, u, plaq, qq)
  if (my_pe()==0)write(40,'(I9,1x,7E20.12)')traj, t, 1-plaq, clov, ZERO, ZERO, qq, ZERO
  do i = 1, nstep
     call runge_kutta4(u, epsilon)
     call clover_energy(clov, u, plaq, qq)
     t = dble(i) * epsilon

     if (my_pe()==0)write(40,'(I9,1x,7E20.12)')traj, t, 1-plaq, clov, clov*t**2, (clov*t**2-clovtt_old)*dble(i), qq, plaq*t**2
     clovtt_old = clov*t**2
  enddo

  call topological_density_dump("QCD",i)

end

!-------------------------------------------------------------------------------
subroutine runge_kutta4(u, epsilon)
  use module_vol
  use module_action
  implicit none
  GAUGE_FIELD, intent(inout) :: u
  GENERATOR_FIELD :: p
  REAL :: s, ss, epsilon

  p = ZERO

!!! last HALF mult comes from im_tr_j(p, u, s)  ! p(j) := p(j) + s * Im Tr(lambda_j U)
  s =  17.0_8/36.0_8/ action%plaq%beta * 3.0_8 * HALF 
  ss=  9.0_8/17.0_8*epsilon
  call dsg(p, u, action%plaq, s)
  call hmc_u(p, u, ss)

  s =  -8.0_8/9.0_8/ action%plaq%beta * 3.0_8 * HALF
  ss=  -epsilon
  call dsg(p, u, action%plaq, s)
  call hmc_u(p, u, ss)

  s =  3.0_8/4.0_8/action%plaq%beta * 3.0_8 * HALF
  ss=  epsilon
  call dsg(p, u, action%plaq, s)
  call hmc_u(p, u, ss)

end

!-------------------------------------------------------------------------------
! mostly copied from cooling.F90
!
subroutine clover_energy(clov_energy, u, plaq_energy, qq)
  use module_wilson_flow
  use module_vol
  use module_nn
  implicit none

  REAL, intent(out)       :: clov_energy, plaq_energy, qq
  GAUGE_FIELD, intent(in) :: u

  integer :: e, o, mu, nu, i, j1, j2, j3, j4, j5, j6, j7, c1, c2
  SU3     :: uuu, left, right
  COMPLEX, dimension(NCOL, NCOL, DIM - 1, DIM) :: ut  ! U~(x,mu,nu) - h.c.
  COMPLEX :: trut
  REAL    :: clov, plaq, inv
  REAL, external :: global_sum, Re_Tr_uu

  clov = ZERO
  plaq = ZERO

  do e = EVEN, ODD
     o = EVEN + ODD - e
     !$omp parallel do reduction(+: clov, plaq) private(uuu, left, right, ut, &
     !$omp mu, nu, i, j1, j2, j3, j4, j5, j6, j7, c1, c2)  
     do i = 1, VOLH
        do mu = 1, DIM - 1
           do nu = mu + 1, DIM
              j1 = nn(i,  e, mu, FWD)
              j2 = nn(i,  e, nu, FWD)
              j3 = nn(j2, o, mu, BWD)
              j4 = nn(j3, e, nu, BWD)
              j5 = nn(j4, o, nu, BWD)
              j6 = nn(j5, e, mu, FWD)
              j7 = nn(j6, o, mu, FWD)
              uuu = ZERO

              call uuu_fwd(uuu, u(1, 1, j1, o, nu), &
                                u(1, 1, j2, o, mu), &
                                u(1, 1, i,  e, nu))

              plaq = plaq + Re_Tr_uu(uuu, u(1, 1, i, e, mu))

              call uuu_bwd_m(uuu, u(1, 1, j7, e, nu), &
                                  u(1, 1, j6, o, mu), &
                                  u(1, 1, j6, o, nu))

              call uu(right, u(1, 1, i, e, mu), uuu)

              uuu = ZERO

              call uuu_fwd(uuu, u(1, 1, i,  e, nu), &
                                u(1, 1, j3, e, mu), &
                                u(1, 1, j4, o, nu))

              call uuu_bwd_m(uuu, u(1, 1, j6, o, nu), &
                                  u(1, 1, j5, e, mu), &
                                  u(1, 1, j5, e, nu))

              call uu(left, uuu, u(1, 1, j4, o, mu))

              do c2 = 1, NCOL
                 do c1 = 1, NCOL
                    ut(c1, c2, mu, nu) = right(c1, c2) - conjg(right(c2, c1)) &
                                       +  left(c1, c2) - conjg(left(c2, c1))
                 enddo
              enddo

              trut = ZERO
              do c1 = 1, NCOL
                    trut = trut + ut(c1, c1, mu, nu)
              enddo
              trut = trut / THREE
              do c1 = 1, NCOL
                 ut(c1, c1, mu, nu) = ut(c1, c1, mu, nu) - trut
              enddo

           enddo ! nu
        enddo ! mu

        q_qcd(i,e) = Re_Tr_uu(ut(1, 1, 1, 2), ut(1, 1, 3, 4)) &
                   - Re_Tr_uu(ut(1, 1, 1, 3), ut(1, 1, 2, 4)) &
                   + Re_Tr_uu(ut(1, 1, 1, 4), ut(1, 1, 2, 3))
        clov = clov - Re_Tr_uu(ut(1, 1, 1, 2), ut(1, 1, 1, 2)) &
                    - Re_Tr_uu(ut(1, 1, 1, 3), ut(1, 1, 1, 3)) &
                    - Re_Tr_uu(ut(1, 1, 1, 4), ut(1, 1, 1, 4)) &
                    - Re_Tr_uu(ut(1, 1, 2, 3), ut(1, 1, 2, 3)) &
                    - Re_Tr_uu(ut(1, 1, 2, 4), ut(1, 1, 2, 4)) &
                    - Re_Tr_uu(ut(1, 1, 3, 4), ut(1, 1, 3, 4))

     enddo ! i
  enddo ! e/o
  
  clov = global_sum(clov)
  plaq = global_sum(plaq)
  plaq_energy = ONE - plaq / dble(THREE * SIX * volume)
  clov_energy =       clov / dble(64 * volume)

  inv = - ONE / dble(256 * PI**2)
  qq = ZERO
  do e = EVEN, ODD
     !$omp parallel do reduction(+: qq) private(i)
     do i = 1, volh
        qq = qq + q_qcd(i,e)*inv
     enddo
  enddo
  qq = global_sum(qq)

end

!-------------------------------------------------------------------------------
!
! Wilson flow for QED field
!
!-------------------------------------------------------------------------------
subroutine wilson_flow_qed(a_in, traj, i_ensemble1, i_ensemble2)
  use module_wilson_flow
  use module_input
  use module_function_decl
  use module_vol
  implicit none
  integer, intent(in)     :: traj, i_ensemble1, i_ensemble2
  U1_GENERATOR_FIELD, intent(in) :: a_in
  U1_GENERATOR_FIELD             :: a
  integer                 :: i, nstep
  REAL                    :: clov, t, epsilon, plaq, qq, clovtt_old
  logical, save :: fileopen = .false.
  character(len = 14) :: filename

  call init_module_wilson_flow()
  traj_qed = traj

  nstep = input%measure_wilson_flow_qed_step
  read(input%measure_wilson_flow_qed_eps,*)epsilon

  a = a_in

  if (.not. fileopen) then
     if (len(trim(input%measure_wilson_flow_qed_file)) > 0) then
        if (my_pe()==0) open(40, file = trim(input%measure_wilson_flow_qed_file), action = "write")
     else
        write(filename, "(a, i6.6)" ) "wilflow_qed.", traj
        if (my_pe()==0) open(40, file = trim(filename), action = "write")
     endif
     fileopen = .true.
  endif

  clovtt_old = ZERO
  call clover_energy_qed(clov, a, plaq, qq)
  if (my_pe()==0)write(40,'(I9,1x,7E20.12)')traj, t, plaq, clov, ZERO, ZERO, qq, ZERO
  do i = 1, nstep
     call runge_kutta4_qed(a, epsilon)
     call clover_energy_qed(clov, a, plaq, qq)
     t = dble(i) * epsilon

     if (my_pe()==0)write(40,'(I9,1x,7E20.12)')traj, t, plaq, clov, clov*t**2, (clov*t**2-clovtt_old)*dble(i), qq, plaq*t**2
     clovtt_old = clov*t**2
  enddo

  call topological_density_dump("QED",i)

end

!-------------------------------------------------------------------------------
subroutine runge_kutta4_qed(a, epsilon)
  use module_vol
  use module_action
  implicit none
  U1_GENERATOR_FIELD, intent(inout) :: a
  U1_GENERATOR_FIELD :: p
  REAL :: s, ss, epsilon

  p = ZERO

  s =  17.0_8/36.0_8
  ss=  9.0_8/17.0_8*epsilon
  call dsg_qed(p, a, s, ONE)
  call hmc_a_qed(p, a, ss)

  s =  -8.0_8/9.0_8
  ss=  -epsilon
  call dsg_qed(p, a, s, ONE)
  call hmc_a_qed(p, a, ss)

  s =  3.0_8/4.0_8
  ss=  epsilon
  call dsg_qed(p, a, s, ONE)
  call hmc_a_qed(p, a, ss)

end

!-------------------------------------------------------------------------------
subroutine clover_energy_qed(clov_energy, a, plaq_energy, qq)
  use module_wilson_flow
  use module_vol
  use module_nn
  implicit none

  REAL, intent(out)       :: clov_energy, plaq_energy, qq
  U1_GENERATOR_FIELD, intent(in) :: a

  integer :: e, o, mu, nu, i, j1, j2, j3, j4, j5, j6, j7
  REAL, dimension(DIM - 1, DIM) :: ut
  REAL    :: clov, plaq, inv
  REAL, external :: global_sum

  clov = ZERO
  plaq = ZERO

  do e = EVEN, ODD
     o = EVEN + ODD - e
     !$omp parallel do reduction(+: clov, plaq) private(ut, &
     !$omp mu, nu, i, j1, j2, j3, j4, j5, j6, j7)  
     do i = 1, VOLH
        do mu = 1, DIM - 1
           do nu = mu + 1, DIM
              j1 = nn(i,  e, mu, FWD)
              j2 = nn(i,  e, nu, FWD)
              j3 = nn(j2, o, mu, BWD)
              j4 = nn(j3, e, nu, BWD)
              j5 = nn(j4, o, nu, BWD)
              j6 = nn(j5, e, mu, FWD)
              j7 = nn(j6, o, mu, FWD)

              plaq = plaq + (a(i, e, mu) + a(j1, o, nu) - a(j2, o, mu) - a(i, e, nu))**2

              ut(mu, nu) = a(i, e, mu) + a(j1, o, nu) - a(j2, o, mu) - a(i, e, nu) &
                         + a(i, e, mu) - a(j7, e, nu) - a(j6, o, mu) + a(j6, o, nu) &
                         - a(j4, o, mu) + a(j4, o, nu) + a(j3, e, mu) - a(i, e, nu) &
                         - a(j4, o, mu) - a(j5, e, nu) + a(j5, e, mu) + a(j6, o, nu)
           enddo ! nu
        enddo ! mu

        q_qed(i,e) = ut(1, 2) * ut(3, 4) &
                   - ut(1, 3) * ut(2, 4) &
                   + ut(1, 4) * ut(2, 3)
        clov = clov + ut(1, 2) * ut(1, 2) &
                    + ut(1, 3) * ut(1, 3) &
                    + ut(1, 4) * ut(1, 4) &
                    + ut(2, 3) * ut(2, 3) &
                    + ut(2, 4) * ut(2, 4) &
                    + ut(3, 4) * ut(3, 4)

     enddo ! i
  enddo ! e/o
  
  clov = global_sum(clov)
  plaq = global_sum(plaq)
  plaq_energy = plaq / dble(6   * volume)
  clov_energy = clov / dble(6*16* volume)

  inv = - ONE / dble(64 * PI**2)
  qq = ZERO
  do e = EVEN, ODD
     !$omp parallel do reduction(+: qq) private(i)
     do i = 1, volh
        qq = qq + q_qed(i,e)*inv
     enddo
  enddo
  qq = global_sum(qq)

end

!-------------------------------------------------------------------------------
subroutine topological_density_qed(a)
  use module_vol
  use module_nn
  use module_lattice
  implicit none

  U1_GENERATOR_FIELD, intent(in) :: a

  integer :: e, o, mu, nu, i, j1, j2, j3, j4, j5, j6, j7, ix,iy,iz,it
  REAL, dimension(DIM - 1, DIM)   :: ut
  REAL, dimension(volh, 0:1)      :: q
  REAL, dimension(0:nx-1, 0:ny-1, 0:nz-1, 0:nt-1) :: qq
  integer, external :: e_o, std_xyzt2i
  integer :: j(4)
!!  REAL :: tq

  !$omp parallel do private(ix, iy, iz, it, j, i, e)
  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/)
     i  = std_xyzt2i(j)
     e  = e_o(j)
     qq(ix,iy,iz,it) = -q(i,e) / (64 * PI**2)
  enddo
  enddo
  enddo
  enddo

  !! tq = ZERO
  do it = 0, nt-1
  do iz = 0, nz-1
  do iy = 0, ny-1
  do ix = 0, nx-1
  !! tq= tq+ qq(ix,iy,iz,it)
     write(*,'(4i5, e20.10)')ix,iy,iz,it,qq(ix,iy,iz,it)
  enddo
  enddo
  enddo
  enddo

end

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