! Copyright (c) 2017 - 2021, Philip Herd
! This file is distributed under the BSD 2-Clause License

module analytic  ! ASCENDING ORDER!
    implicit none
    ! A collection of general purpose analytic formulas
    ! for polynomials of degree 0 through 4

    ! Polynomials with degrees higher than four do not have analytic solutions

    ! this file accepts the coefficients in ascending order: cof(n) /= 0
    !    (cof(1) * x^0 + cof(2) * x^1 + cof(3) * x^2 + ...)
    ! for coefficients in descending order, see 'analytic_dsc.f90'

    integer, parameter  :: dp = selected_real_kind(8)  ! default precision
    real(dp), parameter :: pi = 4.0_dp * atan(1.0_dp)

    ! second and third cubic unity roots (first is just 1)
    real(dp), parameter :: d60 = sin(pi / 3.0_dp), nd60 = -sin(pi / 3.0_dp)
    complex(dp), parameter :: cu2 = (-0.5_dp, d60), cu3 = (-0.5_dp, nd60)

    private
    public :: root_0, root_1, root_2, root_3, root_4

    contains
        ! cof - complex array of coefficients,
        !       in ascending order; cof(n) /= 0
        ! rts - complex array of found roots
        subroutine root_0(cof, rts)  ! constant,  a = 0
            complex(dp), intent(in)  :: cof(:)
            complex(dp), intent(out) :: rts(:)
            rts(1) = 0.0_dp
        end subroutine root_0

        subroutine root_1(cof, rts)  ! linear,    b + ax = 0
            complex(dp), intent(in)  :: cof(:)
            complex(dp), intent(out) :: rts(:)
            rts(1) = -cof(2) / cof(1)
        end subroutine root_1

        subroutine root_2(cof, rts)  ! quadratic, c + bx + ax^2 = 0
            complex(dp), intent(in)  :: cof(:)
            complex(dp), intent(out) :: rts(:)
            complex(dp) :: p, q

            p = sqrt(cof(2) * cof(2) - 4.0_dp * cof(1) * cof(3))
            q = -2.0_dp * cof(3)
            rts(1) = (cof(2) - p) / q
            rts(2) = (cof(2) + p) / q
        end subroutine root_2

        subroutine root_3(cof, rts)  ! cubic,     d + cx + bx^2 + ax^3 = 0
            complex(dp) :: abc, bbb, aad, dd, d0, d1, cin, cc, p, x1
            complex(dp), intent(in)  :: cof(:)
            complex(dp), intent(out) :: rts(:)

            abc = cof(4) * cof(3) * cof(2)
            bbb = cof(3) * cof(3) * cof(3)
            aad = cof(4) * cof(4) * cof(1)

            dd = 18.0_dp * abc * cof(1) - 4.0_dp * bbb * cof(1)
            dd = dd + cof(3) * cof(3) * cof(2) * cof(2) - 4.0_dp * cof(4) * cof(2) * cof(2) * cof(2)
            dd = dd - 27.0_dp * aad * cof(1)
            d0 = cof(3) * cof(3) - 3.0_dp * cof(4) * cof(2)

            if ((dd == 0).and.(d0 == 0)) then  ! triple root
                x1 = -cof(3) / (3.0_dp * cof(4))
                rts(1) = x1
                rts(2) = x1
                rts(3) = x1
            else if ((dd == 0).and.(d0 /= 0)) then  ! double root, simple root
                x1 = ((9.0_dp * cof(4) * cof(1) - cof(3) * cof(2)) / (2.0_dp * d0))
                rts(1) = x1
                rts(2) = x1
                rts(3) = (4.0_dp * abc - 9.0_dp * aad - bbb) / (cof(4) * d0)
            else
                d1 = 2.0_dp * bbb - 9.0_dp * abc + 27.0_dp * aad

                if (d0 == 0) then
                    cin = d1
                else
                    cin = (d1 - sqrt(-27.0_dp * cof(4) * cof(4) * dd)) / 2.0_dp
                end if
                cc = cin ** (1.0_dp / 3.0_dp)
                p = (-1.0_dp / (3.0_dp * cof(4)))

                rts(1) = p * (cof(3) + cc + d0 / cc)
                rts(2) = p * (cof(3) + cu2 * cc + d0 / (cu2 * cc))
                rts(3) = p * (cof(3) + cu3 * cc + d0 / (cu3 * cc))
            end if
        end subroutine root_3

        subroutine root_4(cof, rts)  ! quartic,   e + dx + cx^2 + bx^3 + ax^4 = 0
            complex(dp) :: a, b, c, d, a2, b2, bq, c1
            complex(dp) :: usq, usq2, u4, blkp, blkm
            complex(dp) :: p1, v, u, uu, v3, p2, p3
            complex(dp), intent(in)  :: cof(:)
            complex(dp), intent(out) :: rts(:)

            a = cof(4) / cof(5)
            b = cof(3) / cof(5)
            c = cof(2) / cof(5)
            d = cof(1) / cof(5)

            a2 = a * a; b2 = b * b

            bq = -(2.0_dp * b2 * b) + 9.0_dp * a * b * c
            bq = bq - 27.0_dp * (c * c + a2 * d) + 72.0_dp * b * d
            c1 = b2 - 3.0_dp * a * c + 12.0_dp * d

            p1 = sqrt(bq * bq - (4.0_dp * c1 * c1 * c1))
            v = -(bq - p1) / 2.0_dp  ! solve quadratic
            if (v == 0) v = -(bq + p1) / 2.0_dp  ! choose non zero quad root

            u = a2 / 4.0_dp - (2.0_dp * b) / 3.0_dp
            if (v == 0) then  ! both quad roots zero, uu simplifies to u
                uu = u
            else
                v3 = (v ** (1.0_dp / 3.0_dp)) * cu2
                uu = u + (1.0_dp / 3.0_dp) * (v3 + c1 / v3)
            end if

            p1 = - a / 4.0_dp

            if (uu == 0) then  ! degenerate, quadruple root
                rts(1) = p1
                rts(2) = p1
                rts(3) = p1
                rts(4) = p1
            else
                p2 = 3.0_dp * a2 - 8.0_dp * b - 4.0_dp * uu
                p3 = -(a2 * a) + 4.0_dp * a * b - 8.0_dp * c

                usq = sqrt(uu)
                usq2 = usq / 2.0_dp
                u4 = uu / 4.0_dp

                blkp = (1.0_dp / 4.0_dp) * sqrt(p2 + p3 / usq)
                blkm = (1.0_dp / 4.0_dp) * sqrt(p2 - p3 / usq)

                rts(1) = p1 + usq2 + blkp
                rts(2) = p1 - usq2 + blkm
                rts(3) = p1 + usq2 - blkp
                rts(4) = p1 - usq2 - blkm
            end if
        end subroutine root_4

end module analytic

!program example
!use analytic  ! ascending order
!complex(8) :: cofs(5), rts(4)
!
!cofs = [3, 0, 0, 0, 0]
!call root_0(cofs, rts)
!print *,'root_0 in : ',cofs(1:1)
!print *,'root_0 out: ', rts(1:1)
!print *, ''
!
!cofs = [2, 3, 0, 0, 0]
!call root_1(cofs, rts)
!print *,'root_1 in : ',cofs(1:2)
!print *,'root_1 out: ', rts(1:1)
!print *, ''
!
!cofs = [7, 2, 3, 0, 0]
!call root_2(cofs, rts)
!print *,'root_2 in : ',cofs(1:3)
!print *,'root_2 out: ', rts(1:2)
!print *, ''
!
!cofs = [13, 7, 2, 3, 0]
!call root_3(cofs, rts)
!print *,'root_3 in : ',cofs(1:4)
!print *,'root_3 out: ', rts(1:3)
!print *, ''
!
!cofs = [5, 13, 7, 2, 3]
!call root_4(cofs, rts)
!print *,'root_4 in : ',cofs(1:5)
!print *,'root_4 out: ', rts(1:4)
!end program example


!	module fourier_mod
!	contains

	subroutine convert3d_row_kron( row_in, row_out,
     &    nnodex,nnodey,nnodephi,
     &    nkx1,nkx2, nky1,nky2, nphi1,nphi2,
     &    inv_xx_t, inv_yy_t,inv_zz_t,
     &    xlen,ylen,philen, dx,dy,dphi )

	use kron3_mod
	implicit none

	integer nnodex,nnodey,nnodephi
	integer nkx1,nkx2,  nky1,nky2
	integer nphi1,nphi2

	complex inv_xx_t(1:nnodex,nkx1:nkx2)
	complex inv_yy_t(1:nnodey,nky1:nky2)
	complex inv_zz_t(1:nnodephi,nphi1:nphi2)

	real xlen,ylen,philen,   dx,dy,dphi

	complex row_in(3*nnodex*nnodey*nnodephi)
	complex row_out(3*nnodex*nnodey*nnodephi)

!	---------------
!	local variables
!	---------------
	complex*16, dimension(:,:,:), allocatable :: ealphak,ealpha
	complex*16, dimension(:,:,:), allocatable :: ebetak,ebeta
	complex*16, dimension(:,:,:), allocatable :: ebk,eb

	real*8 :: fact

	integer :: i,j,k,  m,n,nphi

        integer :: i_lo,i_hi,i_size
        integer :: j_lo,j_hi,j_size
        integer :: k_lo,k_hi,k_size
        integer :: m_lo,m_hi,m_size
        integer :: n_lo,n_hi,n_size
        integer :: nphi_lo,nphi_hi,nphi_size
	integer :: icnc , irnc

!       -----------------------------------------------------
!       Forward transformation  used in sft3d_kron from (i,j,k)
!	configuration space to (n,m,nphi) fourier space is
!
!	F( (n,m,nphi); (i,j,k) ) =
!          kron( inv_zz(nphi,k), kron(inv_yy(m,j),inv_xx(n,i) ))
!
!	ealphak(n,m,nphi) = F( (n,m,phi); (i,j,k) ) * ealpha(i,j,k)
!
!	Now we need the action of transpose(F) * vector
!
!	Let F_t( (i,j,k); (n,m,nphi) ) = transpose( F )
!
!	Let inv_xx_t(i,n) = transpose( inv_xx(n,i) )
!	Let inv_yy_t(j,m) = transpose( inv_yy(m,j) )
!	Let inv_zz_t(k,nphi) = transpose( inv_zz(nphi,k) )
!
!	Then F_t( (i,j,k); (n,m,nphi) ) =
!	kron( inv_zz_t(k,nphi), kron(inv_yy_t(j,m), inv_xx_t(i,n)) )
!
!	ealpha(i,j,k) = F_t( (i,j,k); (n,m,nphi) ) * ealphak(n,m,nphi)
!
!       -----------------------------------------------------

	fact = (dx/xlen)*(dy/ylen)*(dphi/philen)

	nphi_lo = nphi1
	nphi_hi = nphi2

	n_lo = nkx1
	n_hi = nkx2

	m_lo = nky1
	m_hi = nky2

	k_lo = 1
	k_hi = nnodephi

	i_lo = 1
	i_hi = nnodex

	j_lo = 1
	j_hi = nnodey

	nphi_size = nphi_hi-nphi_lo+1
	n_size = n_hi-n_lo+1
	m_size = m_hi-m_lo+1

	k_size = k_hi-k_lo+1
	i_size = i_hi-i_lo+1
	j_size = j_hi-j_lo+1

!	---------------------------------------------------------
!	copy to ealphak(n,m,nphi), ebetak(n,m,nphi),ebk(n,m,nphi)
!	---------------------------------------------------------

	allocate( ealphak(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) )
	allocate( ebetak(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) )
	allocate( ebk(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) )

	allocate( ealpha(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi) )
	allocate( ebeta(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi) )
	allocate( eb(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi) )

	do nphi=nphi_lo,nphi_hi
	do m=m_lo,m_hi
	do n=n_lo,n_hi

                icnc = (nphi - nphi_lo) * 3 * m_size * n_size
     &             + (n - n_lo) * 3 * m_size
     &             + (m - m_lo) * 3 + 1

	  ealphak(n,m,nphi) = row_in(icnc)
	  ebetak(n,m,nphi)  = row_in(icnc+1)
	  ebk(n,m,nphi)     = row_in(icnc+2)

	enddo
	enddo
	enddo

        call kron3(
     &              i_size,j_size,k_size,
     &              n_size, m_size, nphi_size,
     &              inv_zz_t,size(inv_zz_t,1),
     &              inv_yy_t,size(inv_yy_t,1),
     &              inv_xx_t,size(inv_xx_t,1),
     &              ealphak, size(ealphak,1), size(ealphak,2),
     &              ealpha, size(ealpha,1), size(ealpha,2) )

        call kron3(
     &              i_size,j_size,k_size,
     &              n_size, m_size, nphi_size,
     &              inv_zz_t,size(inv_zz_t,1),
     &              inv_yy_t,size(inv_yy_t,1),
     &              inv_xx_t,size(inv_xx_t,1),
     &              ebetak, size(ebetak,1), size(ebetak,2),
     &              ebeta, size(ebeta,1), size(ebeta,2) )

        call kron3(
     &              i_size,j_size,k_size,
     &              n_size, m_size, nphi_size,
     &              inv_zz_t,size(inv_zz_t,1),
     &              inv_yy_t,size(inv_yy_t,1),
     &              inv_xx_t,size(inv_xx_t,1),
     &              ebk, size(ebk,1), size(ebk,2),
     &              eb, size(eb,1), size(eb,2) )

!	---------------------------------------------------------
!	copy to ealpha(i,j,k), ebeta(i,j,k), eb(i,j,k) to row_out
!	---------------------------------------------------------

       do k = k_lo,k_hi
       do j = j_lo,j_hi
       do i = i_lo,i_hi

                irnc = (k-k_lo) * j_size * i_size * 3
     &             + (i-i_lo) * j_size * 3
     &             + (j-j_lo) * 3
     &             + 1

		row_out(irnc)   = fact * ealpha(i,j,k)
		row_out(irnc+1) = fact * ebeta(i,j,k)
		row_out(irnc+2) = fact * eb(i,j,k)
             end do
          end do
       end do

	deallocate( ealphak )
	deallocate( ebetak )
	deallocate( ebk )

	deallocate( ealpha )
	deallocate( ebeta )
	deallocate( eb )

	return
	end subroutine convert3d_row_kron

      subroutine sft3d_kron(f, fk,
     .   nxdim, nydim, nphidim,
     .   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2, xx, yy, zz,
     .   xlen, ylen, philen, dx, dy, dphi)

      use kron3_mod
      implicit none

	integer, parameter :: idebug = 0

      integer nxdim, nydim, nphidim, n, m, nphi, i, j, k,
     .   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2

      real x(nxdim), y(nydim), phi(nphidim)
      real xlen, ylen, philen, dx, dy, dphi, fact

      real xkx(nkdim1 : nkdim2),
     .     xky(mkdim1 : mkdim2),
     .     xkz(nphidim1 : nphidim2)
      complex f(nxdim, nydim, nphidim),
     .   fk(nkdim1 : nkdim2, mkdim1 : mkdim2, nphidim1 : nphidim2)
      complex cexpkxkykz

      complex xx(nkdim1 : nkdim2, 1 : nxdim),
     .        yy(mkdim1 : mkdim2, 1 : nydim),
     .        zz(nphidim1 : nphidim2, 1 : nphidim)

	integer :: i_lo,i_hi,i_size
	integer :: j_lo,j_hi,j_size
	integer :: k_lo,k_hi,k_size
	integer :: m_lo,m_hi,m_size
	integer :: n_lo,n_hi,n_size
	integer :: nphi_lo,nphi_hi,nphi_size
	complex*16, dimension(:,:), allocatable :: inv_xx
	complex*16, dimension(:,:), allocatable :: inv_yy
	complex*16, dimension(:,:), allocatable :: inv_zz
	complex*16, dimension(:,:,:), allocatable :: f_in
	complex*16, dimension(:,:,:), allocatable :: fk_kron
	real*8 :: err

	nphi_lo = nphi1
	nphi_hi = nphi2

	n_lo = nkx1
	n_hi = nkx2

	m_lo = nky1
	m_hi = nky2

	k_lo = 1
	k_hi = nnodephi

	i_lo = 1
	i_hi = nnodex

	j_lo = 1
	j_hi = nnodey

	nphi_size = nphi_hi-nphi_lo+1
	n_size = n_hi-n_lo+1
	m_size = m_hi-m_lo+1

	k_size = k_hi-k_lo+1
	i_size = i_hi-i_lo+1
	j_size = j_hi-j_lo+1

      fact = dx/xlen * dy/ylen * dphi/philen

      	allocate( fk_kron(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) )
	allocate( f_in(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi) )

	allocate( inv_xx(n_lo:n_hi,i_lo:i_hi) )
	allocate( inv_yy(m_lo:m_hi,j_lo:j_hi) )
	allocate( inv_zz(nphi_lo:nphi_hi,k_lo:k_hi) )

	do i=i_lo,i_hi
	do n=n_lo,n_hi
	  inv_xx(n,i) = 1.0d0/xx(n,i)
	enddo
	enddo

	do j=j_lo,j_hi
	do m=m_lo,m_hi
	  inv_yy(m,j) = 1.0d0/yy(m,j)
	enddo
	enddo

	do k=k_lo,k_hi
	do nphi=nphi_lo,nphi_hi
	  inv_zz(nphi,k) = 1.0d0/zz(nphi,k)
	enddo
	enddo

!	-----------------------------------------------------
!	fk(n,m,nphi) =
!	kron( inv_zz(nphi,k), kron(inv_yy(m,j),inv_xx(n,i) )) * f(i,j,k)
!	-----------------------------------------------------
	f_in(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi) =
     &            f(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi)

  	call kron3( n_size, m_size, nphi_size,
     &              i_size,j_size,k_size,
     &		    inv_zz,size(inv_zz,1),
     &		    inv_yy,size(inv_yy,1),
     &		    inv_xx,size(inv_xx,1),
     &		    f_in, size(f_in,1), size(f_in,2),
     &              fk_kron, size(fk_kron,1), size(fk_kron,2) )

	fk(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) =
     &        fact * fk_kron(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi)

!debug-begin
	if (idebug.ge.2) then
!	------------
!	double check
!	------------
	fk_kron(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) =
     &          fk(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi)

      do nphi = nphi1, nphi2
         do n = nkx1, nkx2
            do m = nky1, nky2

               fk(n, m, nphi) = 0.0

               do k = 1, nnodephi
                  do i = 1, nnodex
                     do j = 1, nnodey

                        cexpkxkykz = xx(n, i) * yy(m, j) * zz(nphi, k)

                        fk(n, m, nphi) = fk(n, m, nphi)
     .                                 + f(i,j,k) / cexpkxkykz * fact

                     end do
                  end do
               end do

            end do
         end do
      end do

	err = maxval( abs( fk(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) -
     &                 fk_kron(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) ))
        write(*,*) 'sft3d_kron: err fk = ', err
        write(15,*) 'sft3d_kron: err fk = ', err

	endif
!debug-end

	deallocate( f_in )
	deallocate( fk_kron )
	deallocate( inv_xx )
	deallocate( inv_yy )
	deallocate( inv_zz )

      return
      end subroutine sft3d_kron

      subroutine fftinv3d_kron(x, y, phi, xkx, xky, xkz, f, a,
     .   nxdim, nydim, nphidim,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     1   nkx1, nkx2, nky1, nky2, nphi1, nphi2, xx, yy, zz)

	use kron3_mod

      implicit none

	integer, parameter :: idebug = 0

      integer nxdim, nydim, nphidim, n, m, nphi, i, j, k,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     1   nkx1, nkx2, nky1, nky2, nphi1, nphi2

      real x(nxdim), y(nydim), phi(nphidim)
      real xkx(nkdim1 : nkdim2),
     .     xky(mkdim1 : mkdim2),
     .     xkz(nphidim1 : nphidim2)
      complex f(nxdim, nydim, nphidim),
     1   a(nkdim1 : nkdim2, mkdim1 : mkdim2, nphidim1 : nphidim2)
      complex cexpkxkykz

      complex xx(nkdim1 : nkdim2, 1 : nxdim),
     .        yy(mkdim1 : mkdim2, 1 : nydim),
     .        zz(nphidim1 : nphidim2, 1 : nphidim)

	integer :: i_lo,i_hi, j_lo,j_hi, k_lo,k_hi
	integer :: n_lo,n_hi, m_lo,m_hi, nphi_lo,nphi_hi
	integer :: i_size,j_size,k_size
	integer :: m_size,n_size,nphi_size

	complex*16, dimension(:,:), allocatable :: xx_t,yy_t,zz_t
	complex*16, dimension(:,:,:), allocatable :: Ain,Fout
	real*8 :: err

	i_lo = 1
	i_hi = nnodex

	j_lo = 1
	j_hi = nnodey

	k_lo = 1
	k_hi = nnodephi

	nphi_lo = nphi1
	nphi_hi = nphi2

	n_lo = nkx1
	n_hi = nkx2

	m_lo = nky1
	m_hi = nky2

	i_size = i_hi-i_lo+1
	j_size = j_hi-j_lo+1
	k_size = k_hi-k_lo+1

	m_size = m_hi-m_lo+1
	n_size = n_hi-n_lo+1
	nphi_size = nphi_hi - nphi_lo + 1

!	---------------------------
!	generate transpose matrices
!	---------------------------
	allocate(xx_t(i_lo:i_hi,n_lo:n_hi))
	allocate(yy_t(j_lo:j_hi,m_lo:m_hi))
	allocate(zz_t(k_lo:k_hi,nphi_lo:nphi_hi))

	do n=n_lo,n_hi
	do i=i_lo,i_hi
	   xx_t(i,n) = xx(n,i)
	enddo
	enddo

	do m=m_lo,m_hi
	do j=j_lo,j_hi
	  yy_t(j,m) = yy(m,j)
	enddo
	enddo

	do nphi=nphi_lo,nphi_hi
	do k=k_lo,k_hi
	   zz_t(k,nphi) = zz(nphi,k)
	enddo
	enddo

!	------------------------------------------------
!	The nested loop is performing  matrix-vector multiply
!	where the matrix is conceptually the kronecker product
!	of 3 matrices, zz_t, yy_t, xx_t
!
!	f(i,j,k) =
!	  kron(zz_t(k,nphi), kron( yy_t(j,m), xx_t(i,n))*a(n,m,nphi)
!	------------------------------------------------
	allocate( Fout(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi) )
	allocate( Ain(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) )

	Ain(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi) =
     &      a(n_lo:n_hi,m_lo:m_hi,nphi_lo:nphi_hi)

	call kron3(i_size,j_size,k_size,
     &             n_size,m_size,nphi_size,
     &		zz_t, size(zz_t,1),
     &          yy_t, size(yy_t,1),
     &          xx_t, size(xx_t,1),
     &          Ain, size(Ain,1), size(Ain,2),
     &          Fout, size(Fout,1), size(Fout,2) )

	f(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi) =
     &       Fout(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi)

!debug-begin
      if (idebug.ge.2) then
      do k = 1, nnodephi
         do i = 1, nnodex
            do j = 1, nnodey
               f(i,j,k) = 0.0

               do nphi = nphi1, nphi2
c                 if(nphi .ne. 0) then
                     do n = nkx1, nkx2
                        do m = nky1, nky2

c                           cexpkxkykz = exp(zi * (xkx(n) * x(i)
c     .                                          + xky(m) * y(j)
c     .                                          + xkz(nphi) * phi(k)))

                           cexpkxkykz = xx(n, i) * yy(m, j) * zz(nphi,k)

                           f(i,j,k) = f(i,j,k) + a(n,m,nphi) *cexpkxkykz
                        end do
                     end do
c                 end if
               end do

            end do
         end do
      end do

	err = maxval( abs(f(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi) -
     &                    Fout(i_lo:i_hi,j_lo:j_hi,k_lo:k_hi)))
	write(*,*) 'fftinv3d, err Fout ', err
	write(15,*) 'fftinv3d, err Fout ', err
      endif
!debug-end

	deallocate( xx_t )
	deallocate( yy_t )
	deallocate( zz_t )
	deallocate( Ain )
	deallocate( Fout )

      return
      end subroutine fftinv3d_kron
c*******************************************************************
c
      subroutine sft3d(f, fk,
     .   nxdim, nydim, nphidim,
     .   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2, xx, yy, zz,
     .   xlen, ylen, philen, dx, dy, dphi)

      implicit none

      integer nxdim, nydim, nphidim, n, m, nphi, i, j, k,
     .   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2

      real x(nxdim), y(nydim), phi(nphidim)
      real xlen, ylen, philen, dx, dy, dphi, fact

      real xkx(nkdim1 : nkdim2),
     .     xky(mkdim1 : mkdim2),
     .     xkz(nphidim1 : nphidim2)
      complex f(nxdim, nydim, nphidim),
     .   fk(nkdim1 : nkdim2, mkdim1 : mkdim2, nphidim1 : nphidim2)
      complex cexpkxkykz

      complex xx(nkdim1 : nkdim2, 1 : nxdim),
     .        yy(mkdim1 : mkdim2, 1 : nydim),
     .        zz(nphidim1 : nphidim2, 1 : nphidim)

      fact = dx/xlen * dy/ylen * dphi/philen

      do nphi = nphi1, nphi2
         do n = nkx1, nkx2
            do m = nky1, nky2

               fk(n, m, nphi) = 0.0

               do k = 1, nnodephi
                  do i = 1, nnodex
                     do j = 1, nnodey

                        cexpkxkykz = xx(n, i) * yy(m, j) * zz(nphi, k)

                        fk(n, m, nphi) = fk(n, m, nphi)
     .                                 + f(i,j,k) / cexpkxkykz * fact

                     end do
                  end do
               end do

            end do
         end do
      end do

      return
      end subroutine sft3d

c
c*******************************************************************
c
      subroutine convert3d_row(row_in, row_out,
     .   xlen, ylen, philen,
     .   nnodex, nnodey, nnodephi,
     .   nxdim, nydim, nphidim,
     .   nkdim1, nkdim2, mkdim1, mkdim2, lkdim1, lkdim2,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2,
     .   xx, yy, zz, dx, dy, dphi,
     .   ndfmax, nmodesx, nmodesy, nmodesphi)

*     ---------------------------------------------------
*     performs a Fourier transform on one row of the 3-D
*     collocation matrix, p_amat
*        row_in(ncol)  = input  vector of dimension ndfmax
*        row_out(ncol) = output vector of dimension ndfmax
*     ----------------------------------------------------

      implicit none

      integer nnodex, nnodey, nnodephi, nxdim, nydim, nphidim,
     .   nkdim1, nkdim2,
     .   mkdim1, mkdim2, lkdim1, lkdim2, i, j, k, n, m, nphi,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2,
     .   ndfmax, ncol, nmodesx, nmodesy, nmodesphi

      real xlen, ylen, philen, dx, dy, dphi, fact

      complex ealphak(nkdim1 : nkdim2, mkdim1 : mkdim2, lkdim1 :lkdim2),
     .        ebetak(nkdim1 : nkdim2, mkdim1 : mkdim2, lkdim1 : lkdim2),
     .        ebk(nkdim1 : nkdim2, mkdim1 : mkdim2, lkdim1 : lkdim2)

      complex ealpha(nxdim, nydim, nphidim),
     .        ebeta(nxdim, nydim, nphidim),
     .        eb(nxdim, nydim, nphidim)

      complex xx(nkdim1 : nkdim2, 1 : nxdim),
     .        yy(mkdim1 : mkdim2, 1 : nydim),
     .        zz(lkdim1 : lkdim2, 1 : nphidim)

      complex row_in(ndfmax), row_out(ndfmax)

      do nphi = nphi1, nphi2
         do n = nkx1, nkx2
            do m = nky1, nky2

               ncol = (nphi - nphi1) * 3 * nmodesx * nmodesy
     .             + (n - nkx1) * 3 * nmodesy
     .             + (m - nky1) * 3 + 1

               ealphak(n, m, nphi) = row_in(ncol)
               ebetak(n, m, nphi)  = row_in(ncol + 1)
               ebk(n, m, nphi)     = row_in(ncol + 2)

            end do
         end do
      end do

      call fftmat3d(ealphak, ealpha,
     .   nxdim, nydim, nphidim,
     .   nkdim1, nkdim2, mkdim1, mkdim2, lkdim1, lkdim2,
     .   nnodex, nnodey, nnodephi,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2, xx, yy, zz)

      call fftmat3d(ebetak, ebeta,
     .   nxdim, nydim, nphidim,
     .   nkdim1, nkdim2, mkdim1, mkdim2, lkdim1, lkdim2,
     .   nnodex, nnodey, nnodephi,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2, xx, yy, zz)

      call fftmat3d(ebk, eb,
     .   nxdim, nydim, nphidim,
     .   nkdim1, nkdim2, mkdim1, mkdim2, lkdim1, lkdim2,
     .   nnodex, nnodey, nnodephi,
     .   nkx1, nkx2, nky1, nky2, nphi1, nphi2, xx, yy, zz)

      fact = dx / xlen * dy / ylen * dphi / philen

      do i = 1, nnodex
         do j = 1, nnodey
            do k = 1, nnodephi

               ncol = (k - 1) * nnodex * nnodey * 3
     .              + (i - 1) * nnodey * 3
     .              + (j - 1) * 3
     .              + 1

               row_out(ncol)     = ealpha(i, j, k) * fact
               row_out(ncol + 1) = ebeta(i, j, k)  * fact
               row_out(ncol + 2) = eb(i, j, k)     * fact

            end do
         end do
      end do

      return
      end subroutine convert3d_row

c
c*******************************************************************
c
      subroutine fftmat3d(a, f,
     .   nxdim, nydim, nphidim,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     1   nkx1, nkx2, nky1, nky2, nphi1, nphi2, xx, yy, zz)

      implicit none

      integer nxdim, nydim, nphidim, n, m, nphi, i, j, k,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     1   nkx1, nkx2, nky1, nky2, nphi1, nphi2

      real x(nxdim), y(nydim), phi(nphidim)
      real xkx(nkdim1 : nkdim2),
     .     xky(mkdim1 : mkdim2),
     .     xkz(nphidim1 : nphidim2)
      complex f(nxdim, nydim, nphidim),
     1   a(nkdim1 : nkdim2, mkdim1 : mkdim2, nphidim1 : nphidim2)
      complex cexpkxkykz

      complex xx(nkdim1 : nkdim2, 1 : nxdim),
     .        yy(mkdim1 : mkdim2, 1 : nydim),
     .        zz(nphidim1 : nphidim2, 1 : nphidim)

      do k = 1, nnodephi
         do i = 1, nnodex
            do j = 1, nnodey

               f(i,j,k) = 0.0

               do nphi = nphi1, nphi2
                  do n = nkx1, nkx2
                     do m = nky1, nky2

                        cexpkxkykz = xx(n, i) * yy(m, j) * zz(nphi, k)
                        f(i,j,k) = f(i,j,k) + a(n,m,nphi) / cexpkxkykz

                     end do
                  end do
               end do

            end do
         end do
      end do

      return
      end subroutine fftmat3d
c
c*******************************************************************
c
      subroutine fftinv3d(x, y, phi, xkx, xky, xkz, f, a,
     .   nxdim, nydim, nphidim,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     1   nkx1, nkx2, nky1, nky2, nphi1, nphi2, xx, yy, zz)

      implicit none

      integer nxdim, nydim, nphidim, n, m, nphi, i, j, k,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nphidim1, nphidim2,
     .   nnodex, nnodey, nnodephi,
     1   nkx1, nkx2, nky1, nky2, nphi1, nphi2

      real x(nxdim), y(nydim), phi(nphidim)
      real xkx(nkdim1 : nkdim2),
     .     xky(mkdim1 : mkdim2),
     .     xkz(nphidim1 : nphidim2)
      complex f(nxdim, nydim, nphidim),
     1   a(nkdim1 : nkdim2, mkdim1 : mkdim2, nphidim1 : nphidim2)
      complex cexpkxkykz

      complex xx(nkdim1 : nkdim2, 1 : nxdim),
     .        yy(mkdim1 : mkdim2, 1 : nydim),
     .        zz(nphidim1 : nphidim2, 1 : nphidim)

      do k = 1, nnodephi
         do i = 1, nnodex
            do j = 1, nnodey
               f(i,j,k) = 0.0

               do nphi = nphi1, nphi2
c                 if(nphi .ne. 0) then
                     do n = nkx1, nkx2
                        do m = nky1, nky2

c                           cexpkxkykz = exp(zi * (xkx(n) * x(i)
c     .                                          + xky(m) * y(j)
c     .                                          + xkz(nphi) * phi(k)))

                           cexpkxkykz = xx(n, i) * yy(m, j) * zz(nphi,k)

                           f(i,j,k) = f(i,j,k) + a(n,m,nphi) *cexpkxkykz
                        end do
                     end do
c                 end if
               end do

            end do
         end do
      end do

      return
      end subroutine fftinv3d
c
c*******************************************************************
c

      subroutine sftinv2d(x, y, xkx, xky, f, a, nxdim, nydim,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nnodex, nnodey,
     1   nkx1, nkx2, nky1, nky2)

      implicit none

      integer nxdim, nydim, n, m, i, j,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nnodex, nnodey,
     1   nkx1, nkx2, nky1, nky2

      real x(nxdim), y(nydim)
      real xkx(nkdim1 : nkdim2), xky(mkdim1 : mkdim2)
      complex zi, f(nxdim, nydim),
     1   a(nkdim1 : nkdim2, mkdim1 : mkdim2)
      complex cexpkxky

      zi = cmplx(0.0, 1.0)

      do i = 1, nnodex
         do j = 1, nnodey
            f(i, j) = 0.0
            do n = nkx1, nkx2
               do m = nky1, nky2
                  cexpkxky = exp(zi * (xkx(n) * x(i) + xky(m) * y(j)))
                  f(i,j) = f(i,j) + a(n,m) * cexpkxky
               end do
            end do
         end do
      end do

      return
      end subroutine sftinv2d
c
c*******************************************************************
c

      subroutine sftinv2dp(x, y, xkx, xky, f, a, nxdim, nydim,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nnodex, nnodey,
     1   nkx1, nkx2, nky1, nky2,
     1   myrow, mycol, nprow, npcol, icontxt, desc_amat, dlen_)

      implicit none

      logical ismine

      integer nxdim, nydim, n, m, i, j,
     1   nkdim1, nkdim2, mkdim1, mkdim2, nnodex, nnodey,
     1   nkx1, nkx2, nky1, nky2, irnc, icnc

      integer rsrc, csrc, myrow, mycol, nprow, npcol, lrindx, lcindx,
     .   icontxt
      integer dlen_, desc_amat(dlen_)

      real x(nxdim), y(nydim)
      real xkx(nkdim1 : nkdim2), xky(mkdim1 : mkdim2)
      complex zi, f(nxdim, nydim),
     1   a(nkdim1 : nkdim2, mkdim1 : mkdim2)
      complex cexpkxky

      zi = cmplx(0.0, 1.0)

c--   Loop over mode numbers and mesh:
      do i = 1, nnodex
         do j = 1, nnodey

            f(i, j) = 0.0

            irnc = (j-1) * 3 + (i-1) * nnodey * 3 + 1

            do n = nkx1, nkx2
               do m = nky1, nky2

                  icnc = (n - nkx1) * 3 * nnodey + (m - nky1) * 3 + 1

c                  call infog2l(irnc, icnc, desc_amat, nprow, npcol,
c     .                myrow, mycol, lrindx, lcindx, rsrc, csrc)

c                  ismine = (rsrc .eq. myrow) .and. (csrc .eq. mycol)

c                  if (ismine) then
                     cexpkxky = exp(zi * (xkx(n) * x(i) + xky(m) *y(j)))
                     f(i,j) = f(i,j) + a(n,m) * cexpkxky
c                  end if

               end do
            end do
         end do
      end do

c      call zgsum2d(icontxt, 'All', ' ', nnodex, nnodey, f, nxdim,-1, -1)

      return
      end subroutine sftinv2dp
c
c*******************************************************************
c

      subroutine sft(x, fx, xlen, nx, nxdim, nkdim1, nkdim2,
     1   nk1, nk2, an)

      implicit none

      integer nx, nxdim, nkdim1, nkdim2, n, nk, nk1, nk2
      real x(nxdim), xlen, xkn, pi
      complex fx(nxdim), an(nkdim1:nkdim2), zi, xint

      pi = 3.14159
      zi = cmplx(0.0, 1.0)

      do nk = nk1, nk2
         xkn = 2.0 * pi/xlen * nk
         an(nk) = 0.0
         do n = 1, nx-1
            xint = (fx(n)   * exp(-zi * xkn * x(n))
     1            + fx(n+1) * exp(-zi * xkn * x(n+1)) )/ 2.0
            an(nk) = an(nk) + xint * (x(n+1) - x(n)) / xlen
         end do
      end do

      return
      end subroutine sft
c
c*******************************************************************
c
      subroutine sftinv(x, fx, xlen, nx, nxdim, nkdim1, nkdim2,
     1    nk1, nk2, an)

      implicit none

      integer nx, nxdim, nkdim1, nkdim2, n, nk, nk1, nk2
      real x(nxdim), xlen, xkn, pi
      complex zi, fx(nxdim), an(nkdim1 : nkdim2)

      pi = 3.141592654
      zi = cmplx(0.0, 1.0)

      do n = 1, nx
         fx(n) = 0.0
         do nk = nk1, nk2
            xkn = 2.0 * pi * nk / xlen
            fx(n) = fx(n) + an(nk) * exp(zi * xkn * x(n))
         end do
      end do

      return
      end subroutine sftinv
c
c*******************************************************************
c
      subroutine fft(nx, work, nxdim, nkdim1, nkdim2, fx, an)

      implicit none

      integer nhalf, nx, nxdim, nkdim1, nkdim2, n, nk
      complex fx(nxdim), work(nxdim), an(nkdim1 : nkdim2)
      real x(nxdim)

      nhalf = nx / 2
      do n = 1, nx
         work(n) = fx(n)
      end do

      call four1(work, nx, -1)

      do n = 1, nx
         nk = n - 1
         if(nk .gt. nhalf) nk = nk - nx
         an(nk) = work(n) / nx
      end do
      return
      end subroutine fft
c
c***************************************************************************
c
      subroutine fftinv(nx, work, nxdim, nkdim1, nkdim2, fx, an)

      implicit none

      integer nhalf, nx, nxdim, nkdim1, nkdim2, n, nk
      complex fx(nxdim), work(nxdim), an(nkdim1 : nkdim2)
      real x(nxdim)

      nhalf = nx / 2

      do n = 1, nx
         nk = n - 1
         if(nk .gt. nhalf) nk = nk - nx
         work(n) = an(nk)
      end do

      call four1(work, nx, 1)

      do n = 1, nx
         fx(n) = work(n)
      end do

      return
      end subroutine fftinv
c
c***************************************************************************
c
      subroutine fftinvr(nx, work, nxdim, nkdim1, nkdim2, fx, an)

      implicit none

      integer nhalf, nx, nxdim, nkdim1, nkdim2, n, nk
      real fx(nxdim), an(nkdim1 : nkdim2)
      real x(nxdim)
      complex work(nxdim)

      nhalf = nx / 2

      do n = 1, nx
         nk = n - 1
         if(nk .gt. nhalf) nk = nk - nx
         work(n) = cmplx(an(nk), 0.0)
c         write(6, 100) n, work(n)
      end do

      call four1(work, nx, 1)

      do n = 1, nx
         fx(n) = real(work(n))
      end do
  100 format (i10, 1p8e12.4)
      return
      end subroutine fftinvr
c
c*******************************************************************
c
      subroutine four1(cdata, nn, isign)

      implicit none

      integer isign, nn
      complex cdata(nn)
      real data(2*nn)
c         Replaces data(1:2*nn) by its discrete Fourier transform, if isign
c         is input as 1; or replaces
c         data(1:2*nn) by nn times its inverse discrete Fourier transform,
c         if isign is input as -1.
c         data is a complex array of length nn or, equivatlently a real
c         array of length 2*nn.  nn MUST be an integer power of 2
c         (this is not checked for!).
      integer i, istep, j, m, mmax, n
      real tempi, tempr
c          Double precision for the trigonometic recurrences
      double precision theta, wi, wpi, wpr, wr, wtemp

      do j=1,nn
        data( 2*j-1 ) = real(cdata(j))
        data( 2*j   ) = aimag(cdata(j))
      enddo

      n = 2 * nn
      j = 1
c          This is bit reversal section of the routine
      do i = 1, n, 2
         if (j .gt. i) then
            tempr = data(j)
            tempi = data(j+1)
            data(j) = data(i)
            data(j+1) = data(i+1)
            data(i) = tempr
            data(i+1) = tempi
         endif
         m = n/2

    1    if (m.ge.2 .and. j.gt.m) then
            j = j - m
            m = m/2
            go to 1
         endif
         j = j + m
      end do
      mmax = 2
    2 if (n .gt. mmax) then
         istep = 2*mmax
         theta = 6.28318530717959D0/(isign*mmax)
         wpr=-2.d0*dsin(0.5d0*theta)**2

         wpi=dsin(theta)

         wr=1.d0

         wi=0.d0

         do m = 1, mmax, 2
            do i = m, n, istep
               j = i + mmax
               tempr=sngl(wr)*data(j)-sngl(wi)*data(j+1)
               tempi=sngl(wr)*data(j+1)+sngl(wi)*data(j)

               data(j) = data(i) - tempr
               data(j+1) = data(i+1) - tempi
               data(i) = data(i) + tempr
               data(i+1) = data(i+1) + tempi
            end do
            wtemp = wr
            wr = wr*wpr - wi*wpi + wr
            wi = wi*wpr + wtemp*wpi + wi
         end do
         mmax = istep
         go to 2
      endif

      do j=1,nn
        cdata(j) = cmplx( data(2*j-1), data(2*j) )
      enddo

      return
      end subroutine four1

!      end module fourier_mod


! RUN: %python %S/test_folding.py %s %flang_fc1
! Tests folding of CSHIFT (valid cases)
module m
  integer, parameter :: arr(2,3) = reshape([1, 2, 3, 4, 5, 6], shape(arr))
  logical, parameter :: test_sanity = all([arr] == [1, 2, 3, 4, 5, 6])
  logical, parameter :: test_cshift_0 = all(cshift([1, 2, 3], 0) == [1, 2, 3])
  logical, parameter :: test_cshift_1 = all(cshift([1, 2, 3], 1) == [2, 3, 1])
  logical, parameter :: test_cshift_2 = all(cshift([1, 2, 3], 3) == [1, 2, 3])
  logical, parameter :: test_cshift_3 = all(cshift([1, 2, 3], 4) == [2, 3, 1])
  logical, parameter :: test_cshift_4 = all(cshift([1, 2, 3], -1) == [3, 1, 2])
  logical, parameter :: test_cshift_5 = all([cshift(arr, 1, dim=1)] == [2, 1, 4, 3, 6, 5])
  logical, parameter :: test_cshift_6 = all([cshift(arr, 1, dim=2)] == [3, 5, 1, 4, 6, 2])
  logical, parameter :: test_cshift_7 = all([cshift(arr, [1, 2, 3])] == [2, 1, 3, 4, 6, 5])
  logical, parameter :: test_cshift_8 = all([cshift(arr, [1, 2], dim=2)] == [3, 5, 1, 6, 2, 4])
end module


subroutine setaux(mbc,mx,my,mz,xlower,ylower,zlower,dx,dy,dz,maux,aux)

    ! Called at start of computation before calling qinit, and
    ! when AMR is used, also called every time a new grid patch is created.
    ! Use to set auxiliary arrays
    !   aux(1:maux, 1-mbc:mx+mbc, 1-mbc:my+mbc, 1-mbc:mz+mbc).
    ! Note that ghost cell values may need to be set if the aux arrays
    ! are used by the Riemann solver(s).

    implicit none
    integer, intent(in) :: mbc,mx,my,mz,maux
    real(kind=8), intent(in) :: xlower,ylower,zlower,dx,dy,dz
    real(kind=8), intent(out) ::  aux(maux,1-mbc:mx+mbc,1-mbc:my+mbc,1-mbc:mz+mbc)

    real(kind=8) :: xcell, ycell, zcell
    real(kind=8) :: xpcorn(4), ypcorn(4), zpcorn(4), mag
    integer :: i,j,k

    ! Auxiliary variables:
    !       1 slip
    !       2 capacity function value

    ! Loop over all cells
    do k=1-mbc,mz + mbc
      zcell = zlower + (k-0.5d0)*dz
      do j=1-mbc,my + mbc
        ycell = ylower + (j-0.5d0)*dy
        do i=1-mbc,mx + mbc
          xcell = xlower + (i-0.5d0)*dx

          ! initialize slip to zero
          aux(1,i,j,k) = 0.d0

          ! compute capacity function value ( volume = (vol of lower face in y direction)*dy )
          call mapc2p(xcell - 0.5d0*dx, ycell - 0.5d0*dy, zcell - 0.5d0*dz, xpcorn(1), ypcorn(1), zpcorn(1))
          call mapc2p(xcell - 0.5d0*dx, ycell - 0.5d0*dy, zcell + 0.5d0*dz, xpcorn(2), ypcorn(2), zpcorn(2))
          call mapc2p(xcell + 0.5d0*dx, ycell - 0.5d0*dy, zcell + 0.5d0*dz, xpcorn(3), ypcorn(3), zpcorn(3))
          call mapc2p(xcell + 0.5d0*dx, ycell - 0.5d0*dy, zcell - 0.5d0*dz, xpcorn(4), ypcorn(4), zpcorn(4))
          ! only need area ratio from cross-product of diagonals, which will point in the y direction
          mag = dabs((xpcorn(3) - xpcorn(1))*(zpcorn(2) - zpcorn(4)) - (xpcorn(2) - xpcorn(4))*(zpcorn(3) - zpcorn(1)))
          aux(2,i,j,k) = 0.5d0*mag/(dx*dz)

        end do
      end do
    end do

end subroutine setaux


Module HealPix

  IMPLICIT none

  INTEGER(4), private, PARAMETER :: ns_max=8192 ! 2^13 : largest nside available

  ! Numerical Constant (Double precision)
  REAL(8), PARAMETER, public :: QUARTPI=0.785398163397448309615660845819875721049_8
  REAL(8), PARAMETER, public :: HALFPI= 1.570796326794896619231321691639751442099_8
  REAL(8), PARAMETER, public :: PI    = 3.141592653589793238462643383279502884197_8
  REAL(8), PARAMETER, public :: TWOPI = 6.283185307179586476925286766559005768394_8
  REAL(8), PARAMETER, public :: FOURPI=12.56637061435917295385057353311801153679_8
  REAL(8), PARAMETER, public :: SQRT2 = 1.41421356237309504880168872420969807856967_8
  REAL(8), PARAMETER, public :: EULER = 0.5772156649015328606065120900824024310422_8
  REAL(8), PARAMETER, public :: SQ4PI_INV = 0.2820947917738781434740397257803862929220_8
  REAL(8), PARAMETER, public :: TWOTHIRD = 0.6666666666666666666666666666666666666666_8

   integer(4), private, save, dimension(128) :: x2pix=0,y2pix=0

CONTAINS
  !=======================================================================
  subroutine ang2pix_ring(nside, phi, theta, ipix)
    !=======================================================================
    !     renders the pixel number ipix (RING scheme) for a pixel which contains
    !     a point on a sphere at coordinates theta and phi, given the map
    !     resolution parameter nside
    !=======================================================================
    INTEGER(4), INTENT(IN) :: nside
    INTEGER(4), INTENT(OUT) :: ipix
    REAL(8), INTENT(IN) ::  theta, phi

    INTEGER(4) ::  nl4, jp, jm
    REAL(8)    ::  z, za, tt, tp, tmp, temp1, temp2
    INTEGER(4) ::  ir, ip, kshift

    !-----------------------------------------------------------------------
    if (nside<1 .or. nside>ns_max) STOP "ang2pix_ring: nside out of range"
    if (theta<-HALFPI .or. theta>HALFPI)  then
       print *,"ANG2PIX_RING: theta : ",theta," is out of range [-Pi/2, +Pi/2]"
       STOP
    endif

    z = DCOS(HALFPI-theta)
    za = DABS(z)
    tt = MODULO( phi, twopi) / halfpi  ! in [0,4)

    if ( za <= twothird ) then ! Equatorial region ------------------
       temp1 = nside*(.5_8+tt)
       temp2 = nside*.75_8*z
       jp = int(temp1-temp2) ! index of  ascending edge line
       jm = int(temp1+temp2) ! index of descending edge line

       ir = nside + 1 + jp - jm ! in {1,2n+1} (ring number counted from z=2/3)
       kshift = 1 - modulo(ir,2) ! kshift=1 if ir even, 0 otherwise

       nl4 = 4*nside
       ip = INT( ( jp+jm - nside + kshift + 1 ) / 2 ) ! in {0,4n-1}
       if (ip >= nl4) ip = ip - nl4

       ipix = 2*nside*(nside-1) + nl4*(ir-1) + ip

    else ! North & South polar caps -----------------------------

       tp = tt - INT(tt)      !MODULO(tt,1.0_8)
       tmp = nside * DSQRT( 3.0_8*(1.0_8 - za) )

       jp = INT(tp          * tmp ) ! increasing edge line index
       jm = INT((1.0_8 - tp) * tmp ) ! decreasing edge line index

       ir = jp + jm + 1        ! ring number counted from the closest pole
       ip = INT( tt * ir )     ! in {0,4*ir-1}
       if (ip >= 4*ir) ip = ip - 4*ir

       if (z>0._8) then
          ipix = 2*ir*(ir-1) + ip
       else
          ipix = 12*nside**2 - 2*ir*(ir+1) + ip
       endif

    endif

    return
  end subroutine ang2pix_ring

  !=======================================================================
  subroutine pix2ang_ring(nside, ipix, phi, theta)
    !=======================================================================
    !     renders theta and phi coordinates of the nominal pixel center
    !     for the pixel number ipix (RING scheme)
    !     given the map resolution parameter nside
    !=======================================================================
    INTEGER(KIND=4), INTENT(IN) :: ipix, nside
    REAL(KIND=8), INTENT(OUT) ::  theta, phi

    INTEGER(KIND=4) ::  nl2, nl4, npix, ncap, iring, iphi, ip, ipix1
    REAL(KIND=8) ::  fodd, hip, fihip
    !-----------------------------------------------------------------------
    if (nside<1 .or. nside>ns_max) STOP "nside out of range"
    npix = 12*nside**2       ! total number of points
    if (ipix <0 .or. ipix>npix-1) STOP "ipix out of range"

    ipix1 = ipix + 1 ! in {1, npix}
    nl2 = 2*nside
    ncap = 2*nside*(nside-1) ! points in each polar cap, =0 for nside =1

    if (ipix1 <= ncap) then ! North Polar cap -------------

       hip   = ipix1*0.5_8
       fihip = AINT ( hip , kind=8)
       iring = INT( SQRT( hip - SQRT(fihip) ) ) + 1 ! counted from North pole
       iphi  = ipix1 - 2*iring*(iring - 1)

       theta = ACOS( 1.0_8 - iring**2 / (3.0_8*nside**2) )
       phi   = (real(iphi,8) - 0.5_8) * PI/(2.0_8*iring)

    elseif (ipix1 <= nl2*(5*nside+1)) then ! Equatorial region ------

       ip    = ipix1 - ncap - 1
       nl4   = 4*nside
       iring = INT( ip / nl4 ) + nside ! counted from North pole
       iphi  = MODULO(ip,nl4) + 1

       fodd  = 0.5_8 * (1 + MODULO(iring+nside,2))  ! 1 if iring+nside is odd, 1/2 otherwise
       theta = ACOS( (nl2 - iring) / (1.5_8*nside) )
       phi   = (real(iphi,8) - fodd) * PI /(2.0_8*nside)

    else ! South Polar cap -----------------------------------

       ip    = npix - ipix1 + 1
       hip   = ip*0.5_8
       fihip = AINT ( hip , kind=8)
       iring = INT( SQRT( hip - SQRT(fihip) ) ) + 1     ! counted from South pole
       iphi  = 4*iring + 1 - (ip - 2*iring*(iring-1))

       theta = ACOS( -1.0_8 + iring**2 / (3.0_8*nside**2) )
       phi   = (real(iphi,8) - 0.5_8) * PI/(2.0_8*iring)

    endif

	theta=HALFPI-theta

    return
  end subroutine pix2ang_ring ! pix2ang_ring

  END MODULE HEALPIX


c     Calculate LU decomposition of a matrix
c     Author: Pearu Peterson, March 2002
c
c     prefixes: d,z,s,c   (double,complex double,float,complex float)
c     suffixes: _c,_r     (column major order,row major order)

      subroutine dlu_c(p,l,u,a,m,n,k,piv,info,permute_l,m1)
      integer m,n,piv(k),i,j,k,permute_l,m1
      double precision l(m,k),u(k,n),a(m,n)
      double precision p(m1,m1)

      external dgetrf,dlaswp
      call dgetrf(m,n,a,m,piv,info)
      if (info.lt.0) then
         return
      endif
      do 20 i=1,m
         do 10, j=1,n
            if (j.le.k) then
               if (i.eq.j) then
                  l(i,j) = 1d0
               elseif (i.gt.j) then
                  l(i,j) = a(i,j)
               endif
            endif
            if (i.le.k.and.i.le.j) then
               u(i,j) = a(i,j)
            endif
 10      continue
 20   continue
      if (permute_l.ne.0) then
         call dlaswp(k,l,m,1,k,piv,-1)
      else
         do 25 i=1,m
            p(i,i)=1d0
 25       continue
         call dlaswp(m,p,m,1,k,piv,-1)
      endif
      end

      subroutine zlu_c(p,l,u,a,m,n,k,piv,info,permute_l,m1)
      integer m,n,piv(k),i,j,k,permute_l,m1
      complex*16 l(m,k),u(k,n),a(m,n)
      double precision p(m1,m1)

      external zgetrf,zlaswp,dlaswp
      call zgetrf(m,n,a,m,piv,info)
      if (info.lt.0) then
         return
      endif
      do 20 i=1,m
         do 10, j=1,n
            if (j.le.k) then
               if (i.eq.j) then
                  l(i,j) = 1d0
               elseif (i.gt.j) then
                  l(i,j) = a(i,j)
               endif
            endif
            if (i.le.k.and.i.le.j) then
               u(i,j) = a(i,j)
            endif
 10      continue
 20   continue
      if (permute_l.ne.0) then
         call zlaswp(k,l,m,1,k,piv,-1)
      else
         do 25 i=1,m
            p(i,i)=1d0
 25       continue
         call dlaswp(m,p,m,1,k,piv,-1)
      endif
      end

      subroutine slu_c(p,l,u,a,m,n,k,piv,info,permute_l,m1)
      integer m,n,piv(k),i,j,k,permute_l,m1
      real l(m,k),u(k,n),a(m,n)
      real p(m1,m1)

      external sgetrf,slaswp
      call sgetrf(m,n,a,m,piv,info)
      if (info.lt.0) then
         return
      endif
      do 20 i=1,m
         do 10, j=1,n
            if (j.le.k) then
               if (i.eq.j) then
                  l(i,j) = 1e0
               elseif (i.gt.j) then
                  l(i,j) = a(i,j)
               endif
            endif
            if (i.le.k.and.i.le.j) then
               u(i,j) = a(i,j)
            endif
 10      continue
 20   continue
      if (permute_l.ne.0) then
         call slaswp(k,l,m,1,k,piv,-1)
      else
         do 25 i=1,m
            p(i,i)=1e0
 25       continue
         call slaswp(m,p,m,1,k,piv,-1)
      endif
      end

      subroutine clu_c(p,l,u,a,m,n,k,piv,info,permute_l,m1)
      integer m,n,piv(k),i,j,k,permute_l,m1
      complex l(m,k),u(k,n),a(m,n)
      real p(m1,m1)

      external cgetrf,claswp,slaswp
      call cgetrf(m,n,a,m,piv,info)
      if (info.lt.0) then
         return
      endif
      do 20 i=1,m
         do 10, j=1,n
            if (j.le.k) then
               if (i.eq.j) then
                  l(i,j) = 1e0
               elseif (i.gt.j) then
                  l(i,j) = a(i,j)
               endif
            endif
            if (i.le.k.and.i.le.j) then
               u(i,j) = a(i,j)
            endif
 10      continue
 20   continue
      if (permute_l.ne.0) then
         call claswp(k,l,m,1,k,piv,-1)
      else
         do 25 i=1,m
            p(i,i)=1e0
 25       continue
         call slaswp(m,p,m,1,k,piv,-1)
      endif
      end


program analyse
use TrajReader
implicit none

!needed to open the trajectory file
logical :: file_opened =  .False.
character(len=128) :: filename
integer :: filename_len

!data to be read from the trajectory file
integer(c_int) timestep, N
real(c_double), dimension(3) :: box_lo, box_hi, prd
real(c_double), dimension(:,:), allocatable :: x, v, f
logical(c_bool) :: loop_continue = .True.

!file output control
integer, parameter :: output_file = 1
character(*), parameter :: output_filename = "temperatures.dat"

!loop counters
integer :: i, j

!values for temperature calculations
real(c_double), parameter :: mass = 1.0
real(c_double), parameter :: C6 = 4.0
real(c_double), parameter :: C12 = 4.0
real(c_double), parameter :: cutoff = 3.0
real(c_double), parameter :: cutoffsq = cutoff*cutoff

integer :: timesteps_read = 0

!interparticle distance
real(c_double), dimension(3) :: del
real(c_double) :: rsq

!intermediate variables
real(c_double) :: tconF_numerator = 0.0
real(c_double) :: tconF_denominator = 0.0
real(c_double) :: tcon1 = 0.0
real(c_double) :: tkin = 0.0
real(c_double) :: tconF = 0.0

!accumulators for averages over entire simulation
real(c_double) :: tkin_ave = 0.0
real(c_double) :: tconF_numerator_ave = 0.0
real(c_double) :: tconF_denominator_ave = 0.0
real(c_double) :: tcon1_ave = 0.0
real(c_double) :: tkinsq_ave = 0.0
real(c_double) :: tconF_numerator_sq_ave = 0.0
real(c_double) :: tconF_denominator_sq_ave = 0.0
real(c_double) :: tcon1sq_ave = 0.0

!error variables
real(c_double) :: tkin_err = 0.0
real(c_double) :: tconF_numerator_err = 0.0
real(c_double) :: tconF_denominator_err = 0.0
real(c_double) :: tconF_err = 0.0
real(c_double) :: tcon1_err = 0.0

call get_command_argument(1, filename, filename_len)

file_opened = OpenFile(filename, filename_len)

if (.not. file_opened) then
        print*, "Couldn't open file ", filename
else
        open(unit=output_file, file=output_filename, action="write", status="replace")
        write(output_file,*) "# Timestep   Tkin   Tcon"
        do
                loop_continue = readframe(timestep, N, box_lo, box_hi, x, v, f)
                if (.not. loop_continue) exit
                print*, timestep

                ! calculate the periodic replica distance
                ! needed for the minimum_image subroutune
                prd = box_hi - box_lo

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

                !!! START MODIFYING HERE !!!

                ! the following have been read from the trajectory for you:
                ! integer timestep
                ! integer N (number of atoms)
                ! real box_lo(3) minimum coordinates in all three directions
                ! real box_hi(3) maximum coordinates in all three directions
                ! real prd(3) periodic replica distance in all three directions
                ! real x(N,3) atomic positons
                ! real v(N,3) atomic velocities
                ! real f(N,3) total force on each atom F_i

                ! the following subprogrammes are provided:
                ! - minimum_image(vector, prd) : apply the minimum image convention
                ! - veclengthsq(vector) : return the squared length of the vector x**2 + y**2 + z**2

                ! You need to give values to the following variables:
                ! - tkin
                ! - tconF_denominator
                ! - tconF_numerator
                ! the existing code will take care of the averages and standard deviations

                ! two loop counters have been predefined for you: i, j
                ! you also have a predefined vector(3): del
                ! and a variable to store squared lengths: rsq

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

                !!! STOP MODIFYING HERE !!!
                ! (Unless you wish to add a function/subroutine) !

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

                tkin = mass * tkin / (3.0*N)
                tkin_ave = tkin_ave + tkin
                tkinsq_ave = tkinsq_ave + tkin*tkin

                tcon1 = tconF_denominator / tconF_numerator
                tcon1_ave = tcon1_ave + tcon1
                tcon1sq_ave = tcon1sq_ave + tcon1*tcon1

                tconF_numerator_ave = tconF_numerator_ave + tconF_numerator
                tconF_numerator_sq_ave = tconF_numerator_sq_ave + tconF_numerator*tconF_numerator
                tconF_denominator_ave = tconF_denominator_ave + tconF_denominator
                tconF_denominator_sq_ave = tconF_denominator_sq_ave + tconF_denominator*tconF_denominator
                !write the current temperatures to the output file
                write(output_file,*) timestep, tkin, tcon1

                timesteps_read = timesteps_read + 1

                !reset the instantaneous variables for next time
                tconF_numerator = 0.0
                tconF_denominator = 0.0
                tcon1 = 0.0
                tkin = 0.0
                tconF = 0.0
        end do
        close(output_file)

        tkin = tkin_ave / timesteps_read
        tkin_err = stderr(tkin, tkinsq_ave / timesteps_read, timesteps_read)
        tcon1 = tcon1_ave / timesteps_read
        tcon1_err = stderr(tcon1, tcon1sq_ave / timesteps_read, timesteps_read)
        tconF_numerator = tconF_numerator_ave / timesteps_read
        tconF_numerator_err = stderr(tconF_numerator, tconF_numerator_sq_ave / timesteps_read, timesteps_read)
        tconF_denominator = tconF_denominator_ave / timesteps_read
        tconF_denominator_err = stderr(tconF_denominator, tconF_denominator_sq_ave / timesteps_read, timesteps_read)

        tconF = tconF_denominator / tconF_numerator
        tconF_err = tconF**2 * ((tconF_numerator_err/tconF_numerator)**2 + (tconF_denominator_err/tconF_denominator)**2)

        print*, "Read", timesteps_read, "timesteps."
        print*, "Temperatures"
        print*, "Kinetic: ", tkin, "+/-", tkin_err
        print*, "Con1", tcon1, "+/-", tcon1_err
        print*, "ConF", tconF, "+/-", tconF_err
end if
contains
real(c_double) function stderr(X, X2, N)
        implicit none
        real(c_double), intent(in) :: X, X2
        integer(c_int), intent(in) :: N

        stderr = X2 - X**2
        stderr = sqrt(stderr/N)
end function stderr
real(c_double) function veclengthsq(vec)
        implicit none
        real(c_double), dimension(3), intent(in) :: vec

        veclengthsq = vec(1)*vec(1) + vec(2)*vec(2) + vec(3)*vec(3)
end function veclengthsq
subroutine minimum_image(vec, prd)
        real(c_double), dimension(3), intent(inout) :: vec
        real(c_double), dimension(3), intent(in) :: prd
        integer i

        do i=1,3
                if (vec(i) < -prd(i)/2.0) then
                        vec(i) = vec(i) + prd(i)
                else if (vec(i) > prd(i)/2.0) then
                        vec(i) = vec(i) - prd(i)
                end if
        end do
end subroutine minimum_image
end program analyse


      SUBROUTINE SPHBES(N,X,FJ)
      IMPLICIT REAL*8(A-H,O-Z)
      DIMENSION FJ(*)
      DATA XLIM/1.D0/,HF/.5D0/,ZERO/0.D0/,ONE/1.D0/,TNHF/1.05D1/, &
      FFT/1.5D1/,T25/1.D18/,TN25/1.D-18/,TN50/1.D-36/,two/2.D0/
      IF(N.EQ.0) THEN
      IF(X.EQ.0.D0) THEN
      FJ(1)=1.D0
      RETURN
      END IF
      FJ(1)=DSIN(X)/X
      RETURN
      ELSE IF(N.EQ.1) THEN
      IF(X.EQ.0.D0) THEN
      FJ(1)=1.D0
      FJ(2)=0.D0
      RETURN
      END IF
      FJ(1)=DSIN(X)/X
      FJ(2)=DSIN(X)/X**2-DCOS(X)/X
      RETURN
      END IF
      IF (N.GE. 0) GO TO 7
    1 WRITE(6,2)
    2 FORMAT (33H1 ERROR, N SHOULD NOT BE NEGATIVE  )
      GO TO 99
    7 IF (X.GE.ZERO)  GO TO 10
    8 WRITE(6,9)
    9 FORMAT (33H1 ERROR, X SHOULD NOT BE NEGATIVE  )
      GO TO 99
   10 IF (X.GT.XLIM) GO TO 25
!
!     THE FOLLOWING SECTION IS ADDED TO HANDLE THE X<1.0 CASES
!     USING THE POLYNOMIAL EXPANSION OF SPHERICAL BESSEL FUNCTION.
!
!                                                  JUN YE    05/31/90
      DO 100 I=1,N+1
      BRU=SPHBRU(I-1,X)
      COB=1.D0
      DO 110 J=1,I-1
      COB=COB*X/(2.D0*J+1)
  110 CONTINUE
      FJ(I)=BRU*COB
  100 CONTINUE
      RETURN
!
   25 CUFAC=4.2
      IF (X.LT.(N-2)) CUFAC=TNHF/(N+HF-X)
      NS=N+5+X*CUFAC
!
!     ADD ADDITIONAL FACTOR
!
      NS=NS + (FFT/(ONE+DSQRT(X)))
!
      IF ((NS+2).LE.100) GO TO 30
      WRITE(6,28) X,NS
  28  FORMAT (1X,'* WARNING. FOR X=',G13.6,'   BESSH WANTS TO START AT N=',I5)
      NS=98
      IF (X.GT.FFT) GO TO 99
  30  CONTINUE
 313  FFO=ZERO
      FFN=TN25
      M=NS-1
      XI=ONE/X
      FM=(M+M)+ONE
      SDR=FM*TN50
 314  FFP=FM*XI*FFN-FFO
      IF (DABS(FFP).LT.T25) GO TO 315
      SDR=SDR*TN50
      FFP=FFP*TN25
      FFN=FFN*TN25
 315  SDR=SDR + (FM-TWO)*FFP*FFP
      FFO=FFN
      FFN=FFP
      IF (M.LE.N) GO TO 316
      M=M-1
      FM=FM-TWO
      GO TO 314
 316  FJ(M)=FFN
      FJ(M+1)=FFO
      GO TO 33
  32  FJ(M)=FM*XI*FJ(M+1)-FJ(M+2)
      IF(DABS(FJ(M)).GE.T25) GO TO 56
      SDR=SDR + (FM-TWO)*FJ(M)*FJ(M)
      IF (M.LE.1) GO TO 34
   33 M = M-1
      FM=FM-TWO
      GO TO 32
   34 SER=ONE/DSQRT(SDR)
   39 MM = N+1
      DO 40 M=1,MM
      FJ(M)=FJ(M)*SER
   40 CONTINUE
      GO TO 98
   56 JJ= M+1
      NS=N+1
      DO 57 J = JJ,NS
      FJ(J)=FJ(J)*TN25
   57 CONTINUE
      SDR=SDR*TN50
      GO TO 32
   99 CALL EXIT(1)
   98 continue
      return
      END


      SUBROUTINE MMID(Y,DYDX,NVAR,XS,HTOT,NSTEP,YOUT,DERIVS)
      PARAMETER (NMAX=10)
      DIMENSION Y(NVAR),DYDX(NVAR),YOUT(NVAR),YM(NMAX),YN(NMAX)
      H=HTOT/NSTEP
      DO 11 I=1,NVAR
        YM(I)=Y(I)
        YN(I)=Y(I)+H*DYDX(I)
11    CONTINUE
      X=XS+H
      CALL DERIVS(X,YN,YOUT)
      H2=2.*H
      DO 13 N=2,NSTEP
        DO 12 I=1,NVAR
          SWAP=YM(I)+H2*YOUT(I)
          YM(I)=YN(I)
          YN(I)=SWAP
12      CONTINUE
        X=X+H
        CALL DERIVS(X,YN,YOUT)
13    CONTINUE
      DO 14 I=1,NVAR
        YOUT(I)=0.5*(YM(I)+YN(I)+H*YOUT(I))
14    CONTINUE
      RETURN
      END


module mod_parallel

  ! A module to provide parallel facilities
  ! to the shallow water solver.

  use iso_fortran_env, only: int32, real32

  implicit none

  private
  public :: num_tiles, tile_indices, &
            tile_neighbors_1d, tile_neighbors_2d

  interface tile_indices
    module procedure :: tile_indices_1d, tile_indices_2d
  end interface tile_indices

contains

  pure function denominators(n)
    ! Returns all common denominators of n.
    integer(int32), intent(in) :: n
    integer(int32), allocatable :: denominators(:)
    integer(int32) :: i
    denominators = [integer(int32) ::]
    do i = 1, n
      if (mod(n, i) == 0) denominators = [denominators, i]
    end do
  end function denominators

  pure function num_tiles(n)
    ! Returns the optimal number of tiles in 2 dimensions
    ! given total number of tiles n.
    !
    ! Examples:
    !   * num_tiles(1) = [1, 1]
    !   * num_tiles(2) = [2, 1]
    !   * num_tiles(3) = [3, 1]
    !   * num_tiles(4) = [2, 2]
    !   * num_tiles(5) = [5, 1]
    !   * num_tiles(6) = [3, 2]
    !
    integer(int32), intent(in) :: n
    integer(int32) :: num_tiles(2)
    integer(int32), allocatable :: denoms(:)
    integer(int32), allocatable :: dim1(:), dim2(:)
    integer(int32) :: i, j, n1, n2

    ! find all common denominators of the total number of images
    denoms = denominators(n)

    ! find all combinations of common denominators
    ! whose product equals the total number of images
    dim1 = [integer(int32) ::]
    dim2 = [integer(int32) ::]
    do j = 1, size(denoms)
      do i = 1, size(denoms)
        if (denoms(i) * denoms(j) == n) then
          dim1 = [dim1, denoms(i)]
          dim2 = [dim2, denoms(j)]
        end if
      end do
    end do

    ! pick the set of common denominators with the minimal norm
    ! between two elements -- rectangle closest to a square
    num_tiles = [dim1(1), dim2(1)]
    do i = 2, size(dim1)
      n1 = norm2([dim1(i), dim2(i)] - sqrt(real(n)))
      n2 = norm2(num_tiles - sqrt(real(n)))
      if (n1 < n2) num_tiles = [dim1(i), dim2(i)]
    end do

  end function num_tiles

  pure function tile_indices_1d(dims, i, n) result(indices)
    ! Given input global array size, return start and end index
    ! of a parallel 1-d tile that correspond to this image.
    integer(int32), intent(in) :: dims, i, n
    integer(int32) :: indices(2)
    integer(int32) :: offset, tile_size

    tile_size = dims / n

    ! start and end indices assuming equal tile sizes
    indices(1) = (i - 1) * tile_size + 1
    indices(2) = indices(1) + tile_size - 1

    ! if we have any remainder, distribute it to the tiles at the end
    offset = n - mod(dims, n)
    if (i > offset) then
      indices(1) = indices(1) + i - offset - 1
      indices(2) = indices(2) + i - offset
    end if

  end function tile_indices_1d

  pure function tile_indices_2d(dims) result(indices)
    ! Given an input x- and y- dimensions of the total computational domain [im, jm].
    ! returns an array of start and end indices in x- and y-, [is, ie, js, je].
    integer(int32), intent(in) :: dims(2)
    integer(int32) :: indices(4)
    integer(int32) :: tiles(2), tiles_ij(2)
    tiles = num_tiles(num_images())
    tiles_ij = tile_n2ij(this_image())
    indices(1:2) = tile_indices_1d(dims(1), tiles_ij(1), tiles(1))
    indices(3:4) = tile_indices_1d(dims(2), tiles_ij(2), tiles(2))
  end function tile_indices_2d

  pure function tile_neighbors_1d() result(neighbors)
    ! Returns the image indices corresponding
    ! to left and right neighbor tiles.
    integer(int32) :: neighbors(2)
    integer(int32) :: left, right
    if (num_images() > 1) then
      left = this_image() - 1
      right = this_image() + 1
      if (this_image() == 1) then
        left = num_images()
      else if (this_image() == num_images()) then
        right = 1
      end if
    else
      left = 1
      right = 1
    end if
    neighbors = [left, right]
  end function tile_neighbors_1d

  pure function tile_n2ij(n) result(ij)
    ! Given tile index in a 1-d layout, returns the
    ! corresponding tile indices in a 2-d layout.
    !
    !    +---+---+---+
    !  2 | 4 | 5 | 6 |
    !    +---+---+---+
    !  1 | 1 | 2 | 3 |
    !  j +---+---+---+
    !    i 1   2   3
    !
    ! Examples:
    !   * tile_n2ij(2) = [2, 1]
    !   * tile_n2ij(4) = [1, 2]
    !   * tile_n2ij(6) = [3, 2]
    !
    integer(int32), intent(in) :: n
    integer(int32) :: ij(2), i, j, tiles(2)
    if (n == 0) then
      ij = 0
    else
      tiles = num_tiles(num_images())
      j = (n - 1) / tiles(1) + 1
      i = n - (j - 1) * tiles(1)
      ij = [i, j]
    end if
  end function tile_n2ij

  pure function tile_ij2n(ij) result(n)
    ! Given tile indices in a 2-d layout, returns the
    ! corresponding tile index in a 1-d layout:
    !
    !    +---+---+---+
    !  2 | 4 | 5 | 6 |
    !    +---+---+---+
    !  1 | 1 | 2 | 3 |
    !  j +---+---+---+
    !    i 1   2   3
    !
    ! Examples:
    !   * tile_ij2n([2, 1]) = 2
    !   * tile_ij2n([1, 2]) = 4
    !   * tile_ij2n([3, 2]) = 6
    !
    integer(int32), intent(in) :: ij(2)
    integer(int32) :: n, tiles(2)
    if (any(ij == 0)) then
      n = 0
    else
      tiles = num_tiles(num_images())
      n = (ij(2) - 1) * tiles(1) + ij(1)
    end if
  end function tile_ij2n

  pure function tile_neighbors_2d(periodic) result(neighbors)
    ! Returns the neighbor image indices given.
    logical, intent(in) :: periodic
    integer(int32) :: neighbors(4)
    integer(int32) :: tiles(2), tiles_ij(2), itile, jtile
    integer(int32) :: left, right, down, up
    integer(int32) :: ij_left(2), ij_right(2), ij_down(2), ij_up(2)

    tiles = num_tiles(num_images())
    tiles_ij = tile_n2ij(this_image())
    itile = tiles_ij(1)
    jtile = tiles_ij(2)

    ! i, j tile indices for each of the neighbors
    ij_left = [itile - 1, jtile]
    ij_right = [itile + 1, jtile]
    ij_down = [itile, jtile - 1]
    ij_up = [itile, jtile + 1]

    if (periodic) then
      ! set neighbor to wrap around the edge
      if (ij_left(1) < 1) ij_left(1) = tiles(1)
      if (ij_right(1) > tiles(1)) ij_right(1) = 1
      if (ij_down(2) < 1) ij_down(2) = tiles(2)
      if (ij_up(2) > tiles(2)) ij_up(2) = 1
    else
      ! set neighbor to 0 -- no neighbor
      if (ij_left(1) < 1) ij_left = 0
      if (ij_right(1) > tiles(1)) ij_right = 0
      if (ij_down(2) < 1) ij_down = 0
      if (ij_up(2) > tiles(2)) ij_up = 0
    end if

    left = tile_ij2n(ij_left)
    right = tile_ij2n(ij_right)
    down = tile_ij2n(ij_down)
    up = tile_ij2n(ij_up)

    neighbors = [left, right, down, up]

  end function tile_neighbors_2d

end module mod_parallel