!===============================================================================
!
! zolotarev.F90
!
! Subroutines for calculating Zolotarev approximations R(x) to x^-1/2
! See e.g. hep-lat/0206007 for an explanation
!
!-------------------------------------------------------------------------------
! Author: Taylor Haar
! Last modified: 23/03/2017
!-------------------------------------------------------------------------------
# include "defs.h"

!-------------------------------------------------------------------------------
subroutine estZolotarevOrder(n, delta, lambda_min, lambda_max)
  ! Given desired error delta 'delta'
  ! and estimation range [lambda_min, lambda_max],
  ! return an estimate of the required Zolotarev approximation order 'n'.
  implicit none

  REAL :: delta, lambda_min, lambda_max
  integer :: n

  REAL, parameter :: cc = 0.982d0 * 9.17d0, ce = -0.774d0
  REAL, parameter :: ac = 0.465d0, ae = 0.596d0
  REAL :: kappa_c, c_k, a_k

  if (n <= 0 .or. delta <= 0 .or. lambda_min <= 0 .or. lambda_max <= 0) then
     call die("estZolotarevOrder(): negative parameter given.")
  endif

  kappa_c = (lambda_max/lambda_min)**2
  c_k = cc * log(kappa_c) ** ce
  a_k = ac * log(kappa_c) ** ae

  n = ceiling(-log(HALF * delta / a_k)/c_k)

end subroutine estZolotarevOrder

!-------------------------------------------------------------------------------
subroutine getZolotarevCoeffs(n, lambda_min, lambda_max, &
     a_k, b_k, d_n, delta)
  ! Given an estimation range [lambda_min, lambda_max] and order 'n',
  ! calculate the n-th order Zolotarev approximation R(K) to K^-1/2.
  !
  ! Returns the approximation in the format:
  !   R(K) = d_n * prod(k=1, n) (K + a_k)/(K + b_k)
  ! It also returns the corresponding error delta 'delta'.

  implicit none

  integer, intent(in) :: n
  REAL, intent(in) :: lambda_min, lambda_max
  REAL, dimension(n), intent(out) :: a_k, b_k
  REAL, intent(out) :: d_n, delta

  REAL, dimension(n) :: b_m, bp_m, a_m, ap_m
  REAL :: kappa, kappap, Kp, Kpp, sn, cn, dn, xi, z, lambda, M
  integer :: i, j

  interface
     subroutine sncndnK(z, k, sn, cn, dn, K_k)
       REAL, intent(in) :: z, k
       REAL, intent(out) :: sn, cn, dn, K_k
     end subroutine sncndnK
  end interface

  ! Set the modulus parameter to the geometric eigenvalue range
  kappa = abs(lambda_min/lambda_max)
  ! kappap = complementary modulus parameter
  kappap = sqrt(ONE - kappa**2)

  ! First call: determine Kp = 'K-prime' = K(kappap)
  call sncndnK(ZERO, kappap, sn, cn, dn, Kp)
  ! Second call: determine elliptic functions at z=Kp/(2*n + 1), k=kappap
  call sncndnK(Kp/(2*n+1), kappap, sn, cn, dn, Kpp)

  xi = ONE / dn

  do j = 1,n
     ! Construct even terms
     z = (2*j*Kp) / (2*n+1)
     call sncndnK(z, kappap, sn, cn, dn, Kpp)
     b_m(j) = -(cn/sn)**2
     a_m(j) = (sn/cn)**2

     ! Construct odd terms
     z = ((2*j-1)*Kp) / (2*n+1)
     call sncndnK(z, kappap, sn, cn, dn, Kpp)
     bp_m(j) = -(cn/sn)**2
     ap_m(j) = (sn/cn)**2
  end do

  ! Construct the M parameter from Zolotarev's approximation
  M = ONE
  do j = 1,n
     M = M  * (ONE - b_m(j)) / (ONE - bp_m(j))
  enddo

  ! Construct lambda parameter from Zolotarev's approximation
  lambda = M/xi
  do j = 1,n
     lambda = lambda * (ONE - bp_m(j)*xi**2) / (ONE - b_m(j)*xi**2)
  enddo

  ! Calculate the error delta = maximum error in the approximation
  delta = (ONE - lambda)/(ONE + lambda)

  ! Determine the normalization coefficient d_n
  d_n = TWO / (abs(lambda_min) * (ONE + ONE/lambda))
  do j = 1,n
     d_n = d_n * b_m(j)*(ONE - bp_m(j))/(bp_m(j) * (ONE - b_m(j)))
  enddo

  ! a_k = numerator shifts, b_k = denominator shifts in ratapp
  a_k = a_m * (lambda_min ** 2)
  b_k = ap_m * (lambda_min ** 2)

  ! Rescale such that the approximation is <= 1.0
  d_n = d_n/(ONE + delta)
  delta = TWO * delta

end subroutine getZolotarevCoeffs


!-------------------------------------------------------------------------------
subroutine sncndnK(z, k, sn, cn, dn, K_k)
  ! Given z and k = kappa,
  ! calculate the Jacobi elliptic functions sn(z;k), cn(z;k) and dn(z;k),
  ! as well as the quarter period K(k).
  implicit none

  REAL :: z, k
  REAL :: sn, cn, dn, K_k

  REAL :: agm, pb
  integer :: qdrt, sgn

  interface
     recursive subroutine arithGeom(z, a, g, agm, sn)
       REAL, intent(in) :: z, a, g
       REAL, intent(out) :: agm, sn
     end subroutine arithGeom
  end interface

  pb = -ONE
  ! Calculate sn(z;k) and the arithmetic geometric mean of 1 and sqrt(1-k^2)
  call arithGeom(z, ONE, sqrt(ONE - k*k), agm, sn)

  ASSERT(sn <= 1 .and. sn >= 0)
 
  ! Quarter period is easily constructed from the agm
  K_k = PI/(TWO * agm)

  ! sgn = the correct sign of the cosine for the given z and k
  ! If in the first or fourth quadrant, sgn = +1
  ! If in the second or third quadrant, sgn = -1
  qdrt = mod(int(abs(z)/K_k), 4)
  if (qdrt == 1 .or. qdrt == 2) then
     sgn = -1
  else
     sgn = 1
  endif

  cn = sgn * sqrt(ONE - sn**2)
  dn = sqrt(ONE - (k*sn)**2)

end subroutine sncndnK

!-------------------------------------------------------------------------------
subroutine arithGeom(z, a, g, agm, sn)
  ! Calculate the arithmetic geometric mean 'agm' of 'a' and 'g',
  ! and the elliptic function sn
  implicit none

  REAL, intent(in) :: z, a, g
  REAL, intent(out) :: agm, sn
  REAL :: pg

  pg = -1
  call arithGeom_r(z, a, g, pg, agm, sn)

end subroutine arithGeom

!-------------------------------------------------------------------------------
recursive subroutine arithGeom_r(z, a, g, pg, agm, sn)
  ! Recursively calculate the arithmetic geometric mean of 'a' and 'g',
  ! and the elliptic function sn
  implicit none

  REAL, intent(in) :: z, a, g
  REAL, intent(out) :: agm, sn
  REAL, intent(inout) :: pg
  REAL :: xi

  if (g <= pg) then
     ! Previous g < g => have reached arithmetic-geometric mean (agm)
     ! Can start construction of elliptic function sn(z; k)
     pg = -ONE
     agm = a
     sn = sin(z*a)
  else
     ! Goto next step in AGM calculation
     ! a_{i+1} = arithmetic mean of a_i and g_i = 0.5*(a_i + g_i)
     ! g_{i+1} = geometric mean of a_i and g_i = sqrt(a_i * g_i)
     pg = g
     call arithGeom_r(z, HALF*(a+g), sqrt(a*g), pg, agm, xi)
     ! Build sn(z; k) on top
     sn = (TWO * a * xi) / ((a+g) + (a-g)*xi*xi)
  endif
end subroutine arithGeom_r

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