!===============================================================================
!
! integrator.F90
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2008-2013 Yoshifumi Nakamura
!                    2013 Hinnerk Stueben
!
! 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/>.
!
!-------------------------------------------------------------------------------
! Modified to allow for the generalized multi-scale integration scheme
!
! Editor: Taylor Haar
! Last modified: 16/03/2017
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------
module module_integrator

  integer, save :: n_time_scale, tot_pq_loop
  integer, dimension(:), pointer, save  :: scale, m_scale
  REAL, dimension(:), pointer, save  :: step_p
  REAL, dimension(:,:), pointer, save  :: step_q
  character(len=6), pointer, save :: integrator(:)

  ! Storage arrays for the generalized integrator scheme
  REAL, dimension(:,:), pointer, save :: step_ps

  integer,save :: itraj
  integer,save :: iscale
  REAL, save :: len

  logical, save :: initialized = .false.

end module module_integrator

!-------------------------------------------------------------------------------
subroutine integrate(p, p_qed, p_axion, para, conf, sf1, sf2, sf3)
  use module_action
  use typedef_hmc
  use module_integrator
  use module_vol
  implicit none

  GENERATOR_FIELD,    intent(inout) :: p
  U1_GENERATOR_FIELD, intent(inout) :: p_qed
  AXION_MOMENTUM_FIELD, intent(inout) :: p_axion
  type(hmc_para),     intent(in)    :: para
  type(hmc_conf),     intent(inout) :: conf
  REAL,               intent(out)   :: sf1, sf2, sf3
  integer :: i, j
  logical, external :: zero_step_q

  DEBUG2S("Start MD step")
  TIMING_START(timing_bin_integrator)

  action%pfermi(:)%sf = ZERO

  DEBUG2S("Start init_integrator")
  call init_integrator(para%traj_length, para%tau)
  DEBUG2S("End init_integrator")

  ! Alternate between 'Q' (Q -> Q + ieP) and 'P' (P -> P - eF) updates
  do i = 0, tot_pq_loop - 1
     call hmc_integrator_q(p, p_qed, p_axion, para, conf, step_p(i))
     call hmc_integrator_p(p, p_qed, p_axion, para, conf, i, 0)
  enddo

  call hmc_integrator_q(p, p_qed, p_axion, para, conf, step_p(tot_pq_loop))

  ! Handle last 'P' update differently to ensure
  ! we calculate the action terms 'action%pfermi(i)%sf'
  if (zero_step_q(tot_pq_loop)) then
     ! But, if there are no true 'P' updates,
     ! we need to get 'sf' via other means
     call calc_fermi_action(.false.)
  else
     call hmc_integrator_p(p, p_qed, p_axion, para, conf, tot_pq_loop, 1)
  endif

  call calc_fermi_action_rhmc()

  sf1 = action%pfermi(1)%sf
  sf2 = action%pfermi(2)%sf

  sf3 = 0
  do i = 3, n_pfermi
     sf3 = sf3 + action%pfermi(i)%sf
  enddo

  TIMING_STOP(timing_bin_integrator)
  DEBUG2S("End MD step")

end subroutine integrate

!-------------------------------------------------------------------------------
subroutine init_integrator(traj_length, tau)
  use module_integrator
  use module_switches
  implicit none

  REAL, intent(in) :: traj_length, tau
!!  integer,save :: count = 0
  integer :: i, max_pq_scale(n_time_scale)

  if (initialized .and. tau > 0 ) return

  len = traj_length
  if (tau < 0) len = - len

!!  tot_pq_loop = scale(n_time_scale) * 2**(n_time_scale)
!!  tot_pq_loop = scale(n_time_scale) * 2**1 * 3**2
!!  tot_pq_loop = scale(n_time_scale) * 1**(n_time_scale)

  if (switches%genint) then
     ! Use a generalized multi-scale integrator scheme
     !
     ! As the step-counts are arbitrary, we have to check each scale
     ! for the maximum number of P-Q 'loops',
     ! such that we can construct the integrators
     do i = 1, n_time_scale
        ! Calculate the maximum possible # of P-Q loops (minus 1)
        ! for this scale
        select case (integrator(i))
        case ("LPFSTS"); max_pq_scale(i) = scale(i)
        case ("LPFTST"); max_pq_scale(i) = scale(i) * 2
        case ("2MNSTS"); max_pq_scale(i) = scale(i) * 2
        case ("2MNTST"); max_pq_scale(i) = scale(i) * 3
        case ("4MN4FP"); max_pq_scale(i) = scale(i) * 5
        case ("4MN5FV"); max_pq_scale(i) = scale(i) * 5
        case default   ; call die("subintegrator: unknown generalized integrator" &
          // integrator(iscale)  )
        end select
     enddo
     ! Then use the maximum for array allocation
     tot_pq_loop = maxval(max_pq_scale)

  else
     ! Default BQCD nested integrator scheme
     tot_pq_loop = scale(n_time_scale)
     do i = 1, n_time_scale
        select case (integrator(i))
        case ("LPFSTS"); tot_pq_loop = tot_pq_loop
        case ("LPFTST"); tot_pq_loop = tot_pq_loop * 2
        case ("2MNSTS"); tot_pq_loop = tot_pq_loop * 2
        case ("2MNTST"); tot_pq_loop = tot_pq_loop * 3
        case ("4MN4FP"); tot_pq_loop = tot_pq_loop * 5
        case ("4MN5FV"); tot_pq_loop = tot_pq_loop * 5
        case default   ; call die("subintegrator: unknown integrator" &
             // integrator(i)  )
        end select
     enddo
  endif

  if (switches%genint) then
     ! The generalized scheme requires storage for each composite integrator
     allocate(step_ps(0:tot_pq_loop,n_time_scale))
     allocate(step_q(0:tot_pq_loop,n_time_scale))
  else
     allocate(step_p(0:tot_pq_loop))
     allocate(step_q(0:tot_pq_loop,n_time_scale))
  endif

  if (switches%genint) then
     step_ps = ZERO
     step_q = ZERO
     ! Construct the composite integrators for each time-scale
     do i = 1, n_time_scale
        select case (integrator(i))
        case ("LPFSTS"); call hmc_lpfsts_genintegrator(i)
        case ("LPFTST"); call hmc_lpftst_genintegrator(i)
        case ("2MNSTS"); call hmc_2mnsts_genintegrator(i)
        case ("2MNTST"); call hmc_2mntst_genintegrator(i)
        case ("4MN4FP"); call hmc_4mn4fp_genintegrator(i)
        case ("4MN5FV"); call hmc_4mn5fv_genintegrator(i)
        case default   ; call die("subintegrator: unknown integrator" &
             // integrator(i)  )
        end select
     enddo
     ! Then meld them together into a single generalized scheme
     call uber_gen_meld()

  else
     step_p = ZERO
     step_q = ZERO
     itraj = 0
     iscale = 0

     call subintegrator(ONE)
  endif

  initialized = .true.

end subroutine init_integrator

!-------------------------------------------------------------------------------
subroutine delete_integrator()
  use module_integrator
  implicit none

  deallocate(scale, m_scale)
  if (associated(step_p)) deallocate(step_p)
  if (associated(step_ps)) deallocate(step_ps)
  deallocate(step_q)
  deallocate(integrator)
  initialized = .false.
end

!-------------------------------------------------------------------------------
subroutine subintegrator(fac)
  use module_integrator
  implicit none
  REAL, intent(in) :: fac

  select case (integrator(iscale+1))
     case ("LPFSTS"); call hmc_lpfsts_integrator(fac)
     case ("LPFTST"); call hmc_lpftst_integrator(fac)
     case ("2MNSTS"); call hmc_2mnsts_integrator(fac)
     case ("2MNTST"); call hmc_2mntst_integrator(fac)
     case ("4MN4FP"); call hmc_4mn4fp_integrator(fac)
     case ("4MN5FV"); call hmc_4mn5fv_integrator(fac)
     case default   ; call die("subintegrator: unknown integrator" &
                      // integrator(iscale)  )
  end select

end subroutine subintegrator

!-------------------------------------------------------------------------------
logical function zero_step_q(i)
  use module_integrator
  implicit none
  integer, intent(in) :: i
  integer :: j
  REAL :: max_step

  max_step = ZERO

  do j = 1, n_time_scale
     max_step = max(max_step, step_q(i,j)**2)
  enddo

  zero_step_q = (max_step == ZERO)

end function zero_step_q

!-------------------------------------------------------------------------------
subroutine uber_gen_meld()
  ! Meld composite integrators into ONE timeline in tau = step_p

  use module_integrator
  implicit none

  REAL, allocatable :: new_qi(:,:), new_p(:)
  integer :: pi_pos(n_time_scale)
  REAL :: tau, next_step, pi_step(n_time_scale)

  REAL :: tol
  integer :: i, j, n_step

  ! If only one composite integrator, we're already done
  if (n_time_scale == 1) then
     return
  endif

  ! Allocate storage arrays for new integrators
  allocate(new_qi(0:n_time_scale*tot_pq_loop, n_time_scale))
  allocate(new_p(0:n_time_scale*tot_pq_loop))

  new_qi = ZERO
  new_p = ZERO
  tau = ZERO
  n_step = ZERO
  tol = 1e-10

  ! Start with the first U update steps
  pi_pos = ZERO
  do i = 1, n_time_scale
     pi_step(i) = step_ps(pi_pos(i), i)
  enddo

  do while (abs(tau - len) > tol)
     ! Find smallest next step
     next_step = minval(pi_step)
     do i = 1, n_time_scale
        if (abs(pi_step(i) - next_step) > tol) then
           ! This integrator's step > next step:
           ! don't insert force update, but decrement step-size
           new_qi(n_step, i) = 0
           pi_step(i) = pi_step(i) - next_step
        else
           ! This integrator's step = next step:
           ! insert force update, and move to next step
           new_qi(n_step, i) = step_q(pi_pos(i), i)
           if (pi_pos(i) >= tot_pq_loop) then
              pi_step(i) = 0
           else
              pi_step(i) = step_ps(pi_pos(i) + 1, i)
           endif
           pi_pos(i) = pi_pos(i) + 1
        endif
     enddo
     tau = tau + next_step
     new_p(n_step) = next_step
     n_step = n_step + 1
  enddo

  n_step = n_step - 1

  ! Having constructed melded versions of steps, update storage arrays
  deallocate(step_q)
  allocate(step_p(0:n_step), step_q(0:n_step, n_time_scale))
  do i = 0, n_step
     step_p(i) = new_p(i)
     do j = 1, n_time_scale
        step_q(i,j) = new_qi(i,j)
     enddo
  enddo
  deallocate(new_qi, new_p)

  ! Revise number of P/Q loops required
  tot_pq_loop = n_step

end subroutine uber_gen_meld

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