!===============================================================================
!
! init_rhmc.F90 - program for RHMC
!
!-------------------------------------------------------------------------------
!
! Copyright (C) 2006-2011 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/>.
!
!-------------------------------------------------------------------------------
! Modified for Zolotarev approximations and ratapp I/O
!
! Editor: Taylor Haar
! Last modified: 23/03/2017
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------
subroutine init_rhmc(para)
  use module_action
  use module_field
  use module_p_interface
  use typedef_hmc
  use module_input
  use module_rhmc
  use module_switches
  use module_function_decl
  implicit none
  type(hmc_para),  intent(in)    :: para
  integer ::i, j, nth, max_frac, iostat, n_rid, rid_tmp(10)
  FILENAME:: list
  REAL :: dmax
  character (len=100) :: key
  logical, save :: initialized = .false.

  if ( .not. switches%rhmc ) return
  if (initialized) return
  DEBUG2S("Start: init_rhmc")

  ratapp%max_rhmc =1 ! set 0 range for all ratapp, first
  ratapp%min_rhmc =1 ! set 0 range for all ratapp, first
  call init_rhmc_2_e5_5( ratapp(1))
  call init_rhmc_4_e8_16(ratapp(2))
  call init_rhmc_8_e8_5( ratapp(3))
  n_ratapp = 3

  max_frac = 0
  do i = 1, n_ratapp
     max_frac = max(max_frac, ratapp(i)%shift)
     max_frac = max(max_frac, ratapp(i)%shift_exact)
  enddo

  if (.not. associated(sco)) then
     allocate(sco(2, max_frac))
     do i = 1, 2
        do j = 1, max_frac
           nullify(sco(i, j)%sco)
           call allocate_sc_overindexed(sco(i, j)%sco)
        enddo
     enddo
  endif

  if (input%tuning_approx_range /=0) then
     list = input%tuning_approx_range_list
     if (list /= " ") then
        open(ULIST, file = trim(list), action = "read", status = "old")
        iostat = 0
        n_rid  = 0
        do while (iostat == 0)
           read(ULIST, *, iostat = iostat) rid_tmp(n_rid + 1)
           n_rid = n_rid + 1
        enddo
        n_rid = n_rid - 1

        rewind(ULIST)

        iostat = 0
        do j = 1, n_rid
           read(ULIST, *) i, ratapp(rid_tmp(j))%shift, ratapp(rid_tmp(j))%shift_exact, ratapp(rid_tmp(j))%margin_factor
!!           if (i/=j) call die("list is not correct format, skipped rid")
        enddo
        close(ULIST)
        if (my_pe()==0)write(0,*)"tuning_approx_range",j-1,"lines are read"
        do j = 1, n_rid
           ratapp(rid_tmp(j))%max_rhmc =1
           ratapp(rid_tmp(j))%min_rhmc =1
           if (my_pe()==0)write(0,*)"tuning_approx_range reset minmax",rid_tmp(j)
        enddo

     else
        call die("init_rhmc(): tuning_approx_range_list is missing")
     endif
  endif

  list = input%tuning_fraction_tolerance
  if (list /= " ") then
     open(ULIST, file = trim(list), action = "read", status = "old")
     iostat = 0
     i = 0
     do while (iostat == 0)
        read(ULIST, *, iostat = iostat) key
        i = i + 1
     enddo
     i=i-1

     if (my_pe()==0)write(0,*)"tuning_fraction_tolerance",i,"lines are read"

     rewind(ULIST)

     iostat = 0
     allocate(relax(i))
     do j = 1, i
        read(ULIST, *) relax(j)
     enddo

     dmax=ZERO
     do j = 1, i
        dmax=max(relax(j),dmax)
     enddo
     if (dmax==ZERO) call die("init_rhmc(): max of tuning_fraction_tolerance is zero")

     relax=dmax/relax
     if (my_pe()==0) then
        write(0,*)"multi shift solver tolerances are relaxed"
        write(0,*)"shift factor"
        do j = 1, i
           write(0,*)j, relax(j)
        enddo
     endif

!!     deallocate(relax_tmp)
     close(ULIST)
  endif
  initialized = .true.
  DEBUG2S("End: init_rhmc")

end

!-------------------------------------------------------------------------------
subroutine gene_new_coeff()
  use module_action
  use module_rhmc
  use module_rhmc_zolo
  use module_poly
  use module_input
  use module_switches
  use module_function_decl
  implicit none
  REAL,allocatable :: dr(:,:),dp(:,:)
  REAL :: min_new, max_new
  REAL :: min_ev, max_ev, factor, rat_range
  integer :: fid, mid, rid, ftype, min_i, max_i, frid
  integer :: deg(2), pm, num, check_nth, polyid
  logical :: id_done(n_mtilde,8)

  ! Variables for ratapp saving/loading
  logical :: calc_coeffs, fixed_zolo

  ! Variables for Zolotarev
  integer :: n_zolo
  REAL, dimension(:), allocatable :: a_k, b_k
  REAL :: d_n, delta

  DEBUG2S("Start: gene_new_coeff")

  id_done = .false.
  min_new =  99999
  min_ev  =  99999
  max_new = -99999
  max_ev  = -99999

  rid   = 0

!! check when tuning_approx_range,
!! "non zero nth" of all of ftype must be same
if (input%tuning_approx_range==0) then

  check_nth=0
  do fid = 1, n_pfermi
     ftype = action%pfermi(fid)%type
     mid   = action%pfermi(fid)%mid1
     if ( action%pfermi(fid)%nth /= 1 ) then
        if (check_nth==0) then
           check_nth=action%pfermi(fid)%nth
        else
           if (check_nth/=action%pfermi(fid)%nth) then
              write(0,*) "ERROR:with tuning_approx_range=0, different nth", &
              check_nth,action%pfermi(fid)%nth
              call die("with tuning_approx_range=0: different nth")
           endif
        endif
     endif
  enddo
endif

  do fid = 1, n_pfermi
     ftype = action%pfermi(fid)%type
     mid   = action%pfermi(fid)%mid1
     rid   = action%pfermi(fid)%rid
     if ( action%pfermi(fid)%nth /= 1 .and. rid > 0) then
        if (.not. id_done(mid,ftype)) then
           call minmax_id(ftype, mid, min_ev, max_ev, min_i, max_i)
           id_done(mid,ftype) = .true.
!!           if (input%tuning_approx_range/=0)rid = rid + 1
        endif

        ! Choose correct ratapp for the ratapp geometric range
        if (switches%zolotarev .and. action%pfermi(fid)%nth == 2) then
           if (.not. associated(ratapp_zolo(rid)%c_k)) then
              ! Zolotarev approx not initialized yet: force creation
              rat_range = 0
           else
              rat_range = ratapp_zolo(rid)%max_rhmc/ratapp_zolo(rid)%min_rhmc
           endif
           ! Determine if this ratapp is fixed
           fixed_zolo = (input%hmc_zolo_min > 0 .or. &
                .not. switches%rescale_rhmc)
        else
           rat_range = ratapp(rid)%max_rhmc/ratapp(rid)%min_rhmc
           fixed_zolo = .false.
        endif

        ! If eigenvalue range too large
        ! and the ratapp is not a fixed, loaded Zolotarev ratapp,
        ! we need a new ratapp
        if (max_ev/min_ev > rat_range .and. &
             .not. (fixed_zolo .and. ratapp_zolo(rid)%shift > 0)) then

           calc_coeffs = .true.

           if (calc_coeffs) then
              if (switches%zolotarev .and. action%pfermi(fid)%nth == 2) then
                 ! Zolotarev -- generate approximation with correct range
                 if (my_pe() == 0) then
                    write(0, *) "Calculating ratapp with " // &
                         "Zolotarev approximation..."
                 endif

                 if (input%hmc_zolo_min > 0) then
                    ! Approximation range provided
                    min_new = input%hmc_zolo_min
                    max_new = input%hmc_zolo_max
                 else
                    min_new = 0.8 * min_ev
                    max_new = 1.2 * max_ev
                 endif

                 if (input%hmc_zolo_n > 0) then
                    n_zolo = input%hmc_zolo_n
                 else
                    ! Estimate number of shifts required
                    call estZolotarevOrder(n_zolo, input%hmc_zolo_delta, &
                         min_new, max_new)
                 endif

                 ! Knowing n_zolo, allocate tmp storage for coefficients
                 allocate(a_k(n_zolo), b_k(n_zolo))
                 ! Then generate the required coefficients
                 call getZolotarevCoeffs(n_zolo, min_new, max_new, &
                      a_k, b_k, d_n, delta)

                 if (input%hmc_zolo_delta < delta ) then
                    call warn("Produced Zolo delta < input%hmc_zolo_delta")
                 endif

                 if (n_zolo /= ratapp_zolo(rid)%shift) then
                    ! Order of ratapp changed: need to resize storage arrays
                    if (associated(ratapp_zolo(rid)%c_k)) then
                       deallocate(ratapp_zolo(rid)%c_k, ratapp_zolo(rid)%b_k)
                    endif
                    allocate(ratapp_zolo(rid)%c_k(n_zolo))
                    allocate(ratapp_zolo(rid)%b_k(n_zolo))

                    ! Also may need to expand sco storage
                    polyid = action%pfermi(fid)%polyid
                    if (polyid >= 0) then
                       ! Include any polynomial filters
                       call expand_sco(n_zolo + poly(polyid)%n)
                    else
                       call expand_sco(n_zolo)
                    endif
                 endif

                 ! Update the rational approximation accordingly
                 ratapp_zolo(rid)%min_rhmc = min_new
                 ratapp_zolo(rid)%max_rhmc = max_new
                 ratapp_zolo(rid)%shift = n_zolo
                 ratapp_zolo(rid)%exponent = HALF
                 ratapp_zolo(rid)%d_n = d_n
                 ratapp_zolo(rid)%c_k(1:n_zolo) = a_k(1:n_zolo)
                 ratapp_zolo(rid)%b_k(1:n_zolo) = b_k(1:n_zolo)
                 deallocate(a_k, b_k)

              else
                 ! Perform Remez instead
                 if (my_pe() == 0) then
                    write(0, *) "Calculating ratapp with remez algorithm..."
                 endif
                 ratapp(rid)%max_rhmc      = max_ev * ratapp(rid)%margin_factor
                 ratapp(rid)%min_rhmc      = min_ev
                 deg(1) = ratapp(rid)%shift
                 deg(2) = ratapp(rid)%shift_exact

                 num=action%pfermi(fid)%nth
                 do pm = 1, 2
                    allocate(dr(0:deg(pm), pm))
                    allocate(dp(1:deg(pm), pm))
#ifdef FMLIB
                    call remez(ratapp(rid)%min_rhmc,ratapp(rid)%max_rhmc, &
                         1, num*pm, deg(pm),100, dr, dp, pm)
#else
                    call die("remez method needs FMLIB")
#endif
                    call storeremez(rid, deg(pm), pm, num*pm, dr, dp)
                    deallocate(dr)
                    deallocate(dp)
                 enddo
              endif

           endif
        endif

        action%pfermi(fid)%rid = rid
        frid = action%pfermi(fid)%rid

        !! Rescale the rational approximation such that
        !+ the determined eigenvalue range [min_ev, max_ev]
        !+ is geometrically centred in the ratapp range
        !+ [min_rhmc/k, max_rhmc/k], where k = factor.
        !! i.e.
        !+    min_ev/(min_rhmc*1/k) = (max_rhmc*1/k)/max_ev

        if (switches%rescale_rhmc) then
           if (switches%zolotarev) then
              factor= sqrt(ratapp_zolo(frid)%max_rhmc &
                   * ratapp_zolo(frid)%min_rhmc/max_ev/min_ev)
           else
              factor= sqrt(ratapp(frid)%max_rhmc &
                   * ratapp(frid)%min_rhmc/max_ev/min_ev)
           endif
        else
           factor = 1
        endif
        action%pfermi(fid)%factor = factor
     endif
  enddo

  DEBUG2S("End: gene_new_coeff")

end subroutine gene_new_coeff

!-------------------------------------------------------------------------------
subroutine storeremez(rid, deg, pm, num, dr, dp)
  use module_rhmc
  implicit none
  integer :: rid,deg,pm,num
  real(8) :: dr(0:deg,pm),dp(deg,pm)

  if (pm == 1) then
     ratapp(rid)%shift = deg
     ratapp(rid)%exponents(1) = ONE/dble(num)

     if (associated(ratapp(rid)%time_switch)) deallocate(ratapp(rid)%time_switch)
     allocate(ratapp(rid)%time_switch(1:deg))

     if (associated(ratapp(rid)%abar)) deallocate(ratapp(rid)%abar)
     if (associated(ratapp(rid)%bbar)) deallocate(ratapp(rid)%bbar)
     allocate(ratapp(rid)%abar(0:deg))
     allocate(ratapp(rid)%bbar(1:deg))

     ratapp(rid)%abar(0:deg) = dr(0:deg,1)
     ratapp(rid)%bbar(1:deg) = dp(1:deg,1)

  else
     ratapp(rid)%shift_exact = deg
     ratapp(rid)%exponents(2) = ONE/dble(num)

     if (associated(ratapp(rid)%abar_exact_p)) deallocate(ratapp(rid)%abar_exact_p)
     if (associated(ratapp(rid)%bbar_exact_p)) deallocate(ratapp(rid)%bbar_exact_p)
     allocate(ratapp(rid)%abar_exact_p(0:deg))
     allocate(ratapp(rid)%bbar_exact_p(1:deg))
     if (associated(ratapp(rid)%abar_exact_m)) deallocate(ratapp(rid)%abar_exact_m)
     if (associated(ratapp(rid)%bbar_exact_m)) deallocate(ratapp(rid)%bbar_exact_m)
     allocate(ratapp(rid)%abar_exact_m(0:deg))
     allocate(ratapp(rid)%bbar_exact_m(1:deg))

     ratapp(rid)%abar_exact_m(0:deg) = dr(0:deg,1)
     ratapp(rid)%bbar_exact_m(1:deg) = dp(1:deg,1)
     ratapp(rid)%abar_exact_p(0:deg) = dr(0:deg,2)
     ratapp(rid)%bbar_exact_p(1:deg) = dp(1:deg,2)

  endif

end

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