!===============================================================================
!
! stout.F90  -  make stout link   Ref. Phys. Rev. D69, 054501 (2004)
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2006-2008 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"

!-------------------------------------------------------------------------------
module typedef_stout
  type :: type_stout
     P_GAUGE_FIELD :: u
     P_GAUGE_FIELD :: p
     P_GAUGE_FIELD :: q1
     P_GAUGE_FIELD :: q2
     P_GAUGE_FIELD :: b1
     P_GAUGE_FIELD :: b2
     P_GAUGE_FIELD :: fb
  end type type_stout
end

!-------------------------------------------------------------------------------
module module_stout_field

  use typedef_stout

  type  stout_field
     logical :: init
     integer :: n
     REAL    :: alpha
     type(type_stout), dimension(:), pointer :: iter
  end type stout_field

  type(stout_field), dimension(:),pointer,save  :: stout

end

!-------------------------------------------------------------------------------
!
! stout_allocate.F90  -  allocate u-fields for stout link smearing
!
!-------------------------------------------------------------------------------
subroutine stout_allocate(stout)

  use typedef_stout
  use module_p_interface
  implicit none

  type(type_stout), intent(out) :: stout

  nullify(stout%u)
  nullify(stout%p)
  nullify(stout%q1)
  nullify(stout%q2)
  nullify(stout%b1)
  nullify(stout%b2)
  nullify(stout%fb)

  ALLOCATE_G_FIELD(stout%u)
  ALLOCATE_G_FIELD(stout%p)
  ALLOCATE_G_FIELD(stout%q1)
  ALLOCATE_G_FIELD(stout%q2)
  ALLOCATE_G_FIELD(stout%b1)
  ALLOCATE_G_FIELD(stout%b2)
  ALLOCATE_G_FIELD(stout%fb)

end

!-------------------------------------------------------------------------------
module module_stout_link_smearing
  interface

     subroutine stout_link_smearing_up(u, alpha, stout, dobs)
        use typedef_stout
        P_GAUGE_FIELD :: u
        REAL :: alpha
        type(type_stout) :: stout
        logical :: dobs
     end subroutine stout_link_smearing_up

     subroutine stout_link_smearing_down(ff, u, stout, alpha)
        use typedef_stout
        use module_vol
        P_GAUGE_FIELD :: u
        GAUGE_FIELD :: ff
        REAL :: alpha
        type(type_stout) :: stout
     end subroutine stout_link_smearing_down

  end interface
end

!-------------------------------------------------------------------------------
subroutine stout_link_smearing(id)
  use module_field
  use module_stout_field
  use module_stout_link_smearing
  implicit none
  integer,intent(in) :: id
  logical :: dobs
  integer :: i

  if (id == 0) return
  DEBUG2S("Start: stout_link_smearing")
  if (.not. stout(id)%init ) return
  if ( stout(id)%n == 0 .or. stout(id)%alpha == ZERO) return

  if (.not. associated(stout(id)%iter)) then
     allocate(stout(id)%iter( stout(id)%n ))
     do i = 1, stout(id)%n
        call stout_allocate(stout(id)%iter(i))
     enddo
  endif

  dobs = .true.

  call stout_link_smearing_up(gauge(1)%u, &
                              stout(id)%alpha, &
                              stout(id)%iter(1), dobs)
  if ( stout(id)%n > 1 ) then
     do i=2, stout(id)%n
        call stout_link_smearing_up(stout(id)%iter(i-1)%u, &
                                    stout(id)%alpha, &
                                    stout(id)%iter(i), dobs)
     enddo
  endif

  stout(id)%init = .false.

  DEBUG2S("End: stout_link_smearing")
end

!-------------------------------------------------------------------------------
subroutine dstout(uuu, id)
  use module_field
  use module_stout_field
  use module_vol
  use module_stout_link_smearing
  implicit none

  GAUGE_FIELD,     intent(inout) :: uuu
  integer,         intent(in)    :: id
  integer :: i

  if (id == 0) return
  if ( stout(id)%n == 0 .or. stout(id)%alpha == ZERO) return

  if ( stout(id)%n > 1) then
     do i = stout(id)%n, 2, -1
        call stout_link_smearing_down(uuu, stout(id)%iter(i-1)%u, &
                                           stout(id)%iter(i), &
                                           stout(id)%alpha)
     enddo
  endif
  call stout_link_smearing_down(uuu, gauge(1)%u, &
                                     stout(id)%iter(1), &
                                     stout(id)%alpha)

end

!-------------------------------------------------------------------------------
subroutine check_stout()
  use module_field
  use module_stout_field
  use module_function_decl
  use module_vol
  implicit none

  REAL, external :: sg
  REAL           :: stout0, stout1, stout5

  stout0 = sg(gauge(1)%u)
  stout1 = sg(stout(1)%iter(1)%u)
!  stout5 = sg(stout(1)%iter(5)%u)
  stout0 = ONE - stout0 / (SIX * volume)
  stout1 = ONE - stout1 / (SIX * volume)
!  stout5 = ONE - stout5 / (SIX * volume)

  if (my_pe() == 0) write(UREC,*)" %stout_plaq ", stout0,stout1,stout5

end

!-------------------------------------------------------------------------------
subroutine stout_link_smearing_up(u, alpha, stout, dobs)

  use typedef_stout
  use module_nn
  use module_vol
  use module_staple
  use module_switches
  implicit none

  type(type_stout), intent(out) :: stout
  P_GAUGE_FIELD :: u
!!  GAUGE_FIELD :: u
  REAL,        intent(in)  :: alpha
  logical,     intent(in)  :: dobs

  SU3                            :: uuu, w
  integer                        :: mu, eo, i
  REAL :: alpha_half

  DEBUG2S("Start: stout_link_smearing_up")
  TIMING_START(timing_bin_stout_smear)

  alpha_half = alpha * HALF

  do mu = 1, DIM
     do eo = EVEN, ODD
        !$omp parallel do private(uuu, w)
        do i = 1, volh
           call staple(uuu, u, i, eo, mu)
           call uu(w, u(1, 1, i, eo, mu), uuu)
           call u_scale(w, alpha_half)
           call u_i_times(w)
           call traceless_hermit(stout%q1(1, 1, i, eo, mu), w)
        enddo
     enddo
  enddo

  do mu = 1, DIM
     do eo = EVEN, ODD
        !$omp parallel do
        do i = 1, volh
           !
           ! calculate fb(3,3), which is f(3), 
           ! if dobs .true., b1(3) and b2(3) as well
           !
           ! calculate p(3,3),
           ! if dobs .true., B1(3) and B2(3) as well
           !
           call projection_su3(stout%q1(1, 1, i, eo, mu), &
                               stout%p(1, 1, i, eo, mu), &
                               stout%q2(1, 1, i, eo, mu), &
                               stout%b1(1, 1, i, eo, mu), &
                               stout%b2(1, 1, i, eo, mu), &
                               stout%fb(1, 1, i, eo, mu), dobs)
           call uu(stout%u(1, 1, i, eo, mu), &
                   stout%p(1, 1, i, eo, mu), &
                   u(1, 1, i, eo, mu))
        enddo
     enddo
  enddo

  TIMING_START(timing_bin_hmc_xbound_g)
  call xbound_g_field(stout%u(1,1,1,0,1))
  TIMING_STOP(timing_bin_hmc_xbound_g)

!!  if (switches%boundary_sf) call schr_boundary_zero4(stout%u)
  if (switches%boundary_sf) call schr_boundary_gauge(stout%u)

  TIMING_STOP(timing_bin_stout_smear)
  DEBUG2S("End: stout_link_smearing_up")
end

!-------------------------------------------------------------------------------
subroutine projection_su3(q1, p, q2, b1, b2, fb, dobs)
  implicit none

  SU3, intent(in) :: q1
  SU3, intent(out) :: p, q2, b1, b2, fb
  logical, intent(in) :: dobs
  REAL :: c0, c1, c0max
  REAL, external :: Re_Tr_uu

  call uu(q2, q1, q1)
  c0 = Re_Tr_uu(q1,q2) * 0.3333333333333333333333333333333333333333_8
  c1 = ( Re(q2(1,1)) + Re(q2(2,2)) + Re(q2(3,3)) ) * HALF
  c0max= TWO * sqrt( (c1 / THREE )**3 )


  !
  ! a couple of coefficients taken from Chroma
  !
  if (c1 < 0.00000001_8) then
     fb(1,1)=cmplx( &
            1.0_8-c0*c0/720.0_8, &
            -(c0/6.0_8)*(1.0_8 - (c1/20.0_8)*(1.0_8 - (c1/42.0_8))), &
            kind = RKIND)
     fb(2,1)=cmplx( &
            c0/24.0_8*(1.0_8-c1/15.0_8*(1.0_8-3.0_8*c1/112.0_8)), &
            1.0_8-c1/6.0_8*(1.0_8-c1/20.0_8*(1.0_8-c1/42.0_8))-c0*c0/5040.0_8, &
            kind = RKIND)
     fb(3,1)=cmplx( &
            0.5_8*(-1.0_8+c1/12.0_8*(1.0_8-c1/30.0_8*(1.0_8-c1/56.0_8))+c0*c0/20160.0_8), &
            0.5_8*(c0/60.0_8*(1.0_8-c1/21.0_8*(1.0_8-c1/48.0_8))), &
            kind = RKIND)
     if (dobs)then
        fb(1,2)=cmplx( &
            0.0_8, &
            (c0/120.0_8)*(1.0_8-c1/21.0_8), &
            kind = RKIND)
        fb(2,2)=cmplx( &
            -c0/360.0_8*(1.0_8 - 3.0_8*c1/56.0_8 ), &
            -1.0_8/6.0_8*(1.0_8-c1/10.0_8*(1.0_8-c1/28.0_8)), &
            kind = RKIND)
        fb(3,2)=cmplx( &
            0.5_8*( 1.0_8/12.0_8*(1.0_8-(2.0_8*c1/30.0_8)*(1.0_8-3.0_8*c1/112.0_8)) ), &
            0.5_8*( -c0/1260.0_8*(1.0_8-c1/24.0_8)), &
            kind = RKIND)
        fb(1,3)=cmplx( &
            -c0/360.0_8, &
            -(1.0_8/6.0_8)*(1.0_8-(c1/20.0_8)*(1.0_8-c1/42.0_8)), &
            kind = RKIND)
        fb(2,3)=cmplx( &
            (1.0_8/24.0_8)*(1.0_8-c1/15.0_8*(1.0_8-3.0_8*c1/112.0_8)), &
            -c0/2520.0_8, &
            kind = RKIND)
        fb(3,3)=cmplx( &
            0.5_8*c0/10080.0_8, &
            0.5_8*( 1.0_8/60.0_8*(1.0_8-c1/21.0_8*(1-c1/48.0_8)) ), &
            kind = RKIND)
     endif
  else
     call coefficient_f_b(fb,c0,c1,c0max, dobs)
  endif

  if ( c0 < ZERO) then
     fb(1,1) = conjg(fb(1,1))
     fb(2,1) =-conjg(fb(2,1))
     fb(3,1) = conjg(fb(3,1))

     fb(1,2) = conjg(fb(1,2))
     fb(2,2) =-conjg(fb(2,2))
     fb(3,2) = conjg(fb(3,2))

     fb(1,3) =-conjg(fb(1,3))
     fb(2,3) = conjg(fb(2,3))
     fb(3,3) =-conjg(fb(3,3))
  endif

  !
  ! calculate p
  !
  call a0_a1xu1_a2xu2(p, q1, q2, fb(:,1))

  if (dobs)then
     !
     ! calculate B1
     !
     call a0_a1xu1_a2xu2(b1, q1, q2, fb(:,2))
     !
     ! calculate B2
     !
     call a0_a1xu1_a2xu2(b2, q1, q2, fb(:,3))
  endif

end

!-------------------------------------------------------------------------------
subroutine coefficient_f_b(fb,c0,c1,c0max, dobs)

  implicit none
  COMPLEX, dimension(3,3),intent(out)::fb
  REAL,intent(in)::c0,c1,c0max
  logical, intent(in) :: dobs

  COMPLEX :: e2u, emu, h(3), r1(3), r2(3)
  REAL    :: c0abs, theta3, u, w, xi, u2, w2, xi13, w3
  REAL    :: cw, cwt2, ut2, ut3, ut2tcwt2, u2t3, u2t9, u2t9mw2, u2t3pw2, u2t3mw2, xit3
  REAL    :: pb

  COMPLEX :: i_times, c
  i_times(c) = cmplx(-aimag(c), real(c), kind = RKIND)
  integer :: i

!--------------------------op 21
  c0abs =  abs(c0)
  pb =  (c0max - c0abs)/c0max

  if (pb < 0) then
     theta3 = ZERO
  elseif(pb < 0.0001_8) then
     theta3 = sqrt(TWO)*sqrt(pb)*( 1.0_8 + ( 1.0_8 /12.0_8  &
     + ( 3.0_8 /160.0_8 + ( 5.0_8 /896.0_8  + ( 35.0_8 /18432.0_8  &
     + 63.0_8 /90112.0_8 *pb ) *pb) *pb) *pb) *pb)
     theta3 = theta3 / THREE
  else
     theta3 = acos(c0abs / c0max) / THREE
  endif

  u = cos( theta3 ) * sqrt(c1 / THREE )
  w = sin( theta3 ) * sqrt(c1)

  e2u = cmplx( cos(TWO*u),  sin(TWO*u), kind = RKIND)
  emu = cmplx( cos(    u), -sin(    u), kind = RKIND)
  u2 = u**2
  w2 = w**2
  w3 = w2*w

!--------------------------op 8
  if (abs(w) >= 0.00000001_8 ) then
     xi  = sin(w) / w
     xi13= THREE * ( cos(w) / w2 - sin(w) / w3 )
  else
     xi13 = - ONE + w2/10.0_8 * (ONE - w2/28.0_8 * (ONE - w2/54.0_8))
     xi   =   ONE - w2/SIX    * (ONE - w2/20.0_8 * (ONE - w2/42.0_8))
  endif

!--------------------------op 11
  cw  = cos(w)
  cwt2    = cw*TWO
  ut2     = u *TWO
  ut3     = u *THREE
  ut2tcwt2= ut2*cwt2
  u2t3    = u2  *THREE
  u2t9    = u2t3*THREE
  u2t9mw2 = u2t9 - w2
  u2t3pw2 = u2t3+w2
  u2t3mw2 = u2t3-w2
  xit3    = xi * THREE

!--------------------------op 18 + 6 = 24

  h(1) = (u2-w2)*e2u + emu*cmplx(EIGHT*u2*cw, ut2* c1*xi, kind = RKIND)
  h(2) =   ut2  *e2u - emu*cmplx(     ut2*cw,-u2t3mw2*xi, kind = RKIND)
  h(3) =         e2u - emu*cmplx(         cw, ut3    *xi, kind = RKIND)

  do i=1,3
     fb(i,1) = h(i) / u2t9mw2
  enddo

!--------------------------op 30 + 21 + 12 + 18 + 11 + 7 = 99
  if (dobs) then
     r1(1)=TWO*e2u*cmplx(u, u2-w2, kind = RKIND) &
          +TWO*emu*cmplx(EIGHT*u*cw+u*c1*xi, (9.0_8*u2+w2)*xi-FOUR*u2*cw, kind = RKIND)

     r1(2)= e2u*cmplx(TWO, FOUR*u, kind = RKIND) &
          + emu*cmplx(u2t3mw2*xi-TWO*cw, TWO*u*(cw+xit3), kind = RKIND)

     r1(3)= i_times(TWO*e2u+emu*cmplx(cw-xit3, u*xit3, kind = RKIND))

     r2(1)=-TWO*e2u + TWO*u*i_times(emu*cmplx(cw+xi+u2*xi13, FOUR*u*xi  ,kind = RKIND))
     r2(2)=         -       i_times(emu*cmplx(cw+xi-u2*xi13,  TWO*u*xi  ,kind = RKIND))
     r2(3)=                         emu*cmplx(   xi        , -    u*xi13,kind = RKIND)

!--------------------------op 3 * (12 + 10 + 12 ) = 102
     do i=1,3
        fb(i,2) = TWO*u*r1(i) +       u2t3mw2*r2(i) - TWO*(15.0_8*u2 + w2)*fb(i,1)
        fb(i,3) =       r1(i) -  THREE*u     *r2(i) -      24.0_8*u       *fb(i,1)
     enddo

     do i=1,3
        fb(i,2) = HALF * fb(i,2)/u2t9mw2**2
        fb(i,3) = HALF * fb(i,3)/u2t9mw2**2
     enddo
  endif

end

!-------------------------------------------------------------------------------
subroutine a0_a1xu1_a2xu2(u,u1,u2,a)

  implicit none

  SU3, intent(out)   :: u
  SU3, intent(in)    :: u1, u2
  COMPLEX,intent(in) :: a(3)

  integer :: i, j

  do j = 1, NCOL
     do i = 1, NCOL
        u(i, j) = a(2) * u1(i, j) + a(3) * u2(i, j)
     enddo
  enddo

  do i = 1, NCOL
     u(i, i) = u(i, i) + a(1)
  enddo

end

!-------------------------------------------------------------------------------
subroutine stout_link_smearing_down(ff, u, stout, alpha)

  use typedef_stout
  use module_vol
  use module_p_interface
  use module_staple
  use module_switches
  implicit none

  GAUGE_FIELD,       intent(inout) :: ff
  P_GAUGE_FIELD :: u
  type(type_stout),  intent(in)    :: stout
  REAL,              intent(in)    :: alpha

  P_GAUGE_FIELD,save :: lam
  SU3 :: uf, uuu, w ,ulam
  REAL :: t1, t2, alpha_half
  REAL, external :: Re_Tr_uu
  integer :: mu, i ,eo

  DEBUG2S("Start: stout_link_smearing_down")
  alpha_half = alpha * HALF

  TIMING_START(timing_bin_stout_diffe)
  ALLOCATE_G_FIELD(lam)

  do mu = 1, DIM !!- 1
     do eo = EVEN, ODD
        !$omp parallel do private(uf, t1, t2)
        do i = 1, volh
           call uu(uf,u(1,1,i,eo,mu),ff(1,1,i,eo,mu))
           t1 = Re_Tr_uu(stout%b1(1,1,i,eo,mu), uf)
           t2 = Re_Tr_uu(stout%b2(1,1,i,eo,mu), uf)
           call lambda( lam(1,1,i,eo,mu), t1, t2, &
                stout%fb(2,1,i,eo,mu), &
                stout%fb(3,1,i,eo,mu), &
                stout%q1(1,1,i,eo,mu), &
                stout%q2(1,1,i,eo,mu), uf)
        enddo
     enddo
  enddo

  TIMING_START(timing_bin_hmc_xbound_g)
  call xbound_g_field(lam(1,1,1,0,1))
  TIMING_STOP(timing_bin_hmc_xbound_g)

  if (switches%boundary_sf) call schr_boundary_zero4(lam)

!---------------------------- new force matrix
  do mu = 1, DIM
     do eo = EVEN, ODD
        !$omp parallel do private(uuu, w, ulam)
        do i = 1, volh
           call u_update2( ff(1,1,i,eo,mu), stout%p(1,1,i,eo,mu) )
           call staple(uuu, u, i, eo, mu)
           call uu(w, uuu, lam(1,1,i,eo,mu))
           call ulam_space(ulam, u(1,1,1,0,1), lam(1,1,1,0,1), i, eo, mu)
           call stout_force_update(ff(1,1,i,eo,mu), w, ulam, alpha_half)
        enddo
     enddo
  enddo
!---------------------------------------------



  if (switches%boundary_sf) call schr_boundary_zero4(ff)
  if (switches%boundary_sf) call schr_boundary_zero3(ff)

  TIMING_STOP(timing_bin_stout_diffe)
  DEBUG2S("End: stout_link_smearing_down")

end

!-------------------------------------------------------------------------------
subroutine lambda(lam,t1,t2,f2,f3,q1,q2,u)

  implicit none

  SU3,     intent(out) :: lam
  SU3,     intent( in) :: q1, q2, u
  COMPLEX, intent( in) :: f2, f3
  REAL,    intent( in) :: t1, t2

  integer :: i, j
  SU3 :: qu, uq, gamma

  call uu(qu, q1,  u)
  call uu(uq,  u, q1)

  do j = 1, NCOL
     do i = 1, NCOL
        gamma(i,j) = f2*u(i,j) + f3*(uq(i,j) + qu(i,j)) + t1*q1(i,j) + t2*q2(i,j)
     enddo
  enddo

  call traceless_hermit(lam, gamma)

end

!-------------------------------------------------------------------------------
subroutine ulam_space(ulam, u, lam, i, e, mu)

  use module_nn
  use module_vol
  implicit none

  SU3, intent(out) :: ulam
  GAUGE_FIELD, intent(in) :: u, lam
  integer, intent(in) :: i, e, mu
  integer :: o, nu, j1, j2, j3, j4
  SU3 :: clp, tmp12, tmp34, tmp_lam33sub, tmp33

  o = EVEN + ODD - e
  ulam = ZERO

  do nu = 1, DIM !!- 1
     if (nu /= mu) then

        !  (j2,o) --<--   x        nu
        !    |            |
        !    v            ^         ^
        !    |            |         |
        !  (i,e)  -->-- (j1,o)      x-->  mu
        !    |            |
        !    ^            v
        !    |            |
        !  (j3,o) --<-- (j4,e)

        j1 = nn(i, e, mu, FWD)
        j2 = nn(i, e, nu, FWD)
        j3 = nn(i, e, nu, BWD)
        j4 = nn(j3,o, mu, FWD)

        if (j4 /= nn(j1, o, nu, BWD)) call die('staple(): j4 inconsistent')
        if (nn(j1, o, nu, FWD) /= nn(j2, o, mu, FWD)) &
           call die('staple(): j12 inconsistent')

        call uud(tmp12, u(1, 1, j1, o, nu), u(1, 1, j2, o, mu))
        call udu(tmp33, u(1, 1, j3, o, mu), u(1, 1, j3, o, nu))
        call uud(clp, tmp12, u(1, 1, i, e, nu))
        call uu(tmp34, u(1,1,j3,o,mu), u(1,1,j4,e,nu))
        call u_sub2(tmp_lam33sub, lam(1,1,j3,o,mu), lam(1,1,j3,o,nu))

        call  duup(ulam,  tmp34,tmp_lam33sub,  u(1,1,j3,o,nu))
        call   uup(ulam,  clp,lam(1,1,i ,e,nu))
        call   uum(ulam,lam(1,1,j1,o,nu),  clp)
        call  duup(ulam,  u(1,1,j4,e,nu),lam(1,1,j4,e,nu), tmp33)
        call  uudp(ulam,  tmp12,lam(1,1,j2,o,mu),  u(1,1,i ,e,nu))

     endif
  enddo

end

!-------------------------------------------------------------------------------
subroutine stout_force_update(f, w, ulam, alpha_half)

  SU3, intent(inout) :: f
  SU3, intent(in)    :: w, ulam
  REAL,intent(in)    :: alpha_half

  SU3 :: tmp
!  integer :: i, j
!  COMPLEX :: i_times, c, w_minus_ulam
!  i_times(c) = cmplx(-aimag(c), real(c), kind = RKIND)

  tmp(1:NCOL,1:NCOL) = w(1:NCOL,1:NCOL) - ulam(1:NCOL,1:NCOL)
  call u_scale(tmp,alpha_half)
  call u_i_times(tmp)
  f(1:NCOL,1:NCOL) = f(1:NCOL,1:NCOL) + tmp(1:NCOL,1:NCOL)

!  do j = 1, NCOL
!     do i = 1, NCOL
!        w_minus_ulam = w(i,j) - ulam(i,j)
!        w_minus_ulam = w_minus_ulam * alpha_half
!        f(i, j)= f(i, j) +  i_times(w_minus_ulam)
!     enddo
!  enddo

end

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