! gfortran -O3 -fopenmp -o mwe mwe.f08
program mwe
	use, intrinsic :: iso_c_binding
	use omp_lib
	implicit none
	
	integer, parameter :: N = 729, threadNum = 1, iterations = 1000
	integer :: i, j, k
	double precision :: start, finish
	complex(C_DOUBLE_COMPLEX), parameter :: im = (0.d0, 1.d0)
	double precision, parameter :: dx = 1.d0/dble(N), dy = dx, dt = dx
	double precision :: getFreq
	
	complex(C_DOUBLE_COMPLEX), allocatable, dimension(:,:) :: A, T
	
	allocate(A(N,N))
	allocate(T(N,N))
	
	A = 1.d0
	
	call omp_set_num_threads(threadNum)
	
	! T(x,y) = e^(-i/2 (px^2 + py^2) dt)
	!$OMP PARALLEL DO
	do j = 1, N
		do i = 1, N
			T(i,j) = exp(-im/2.d0 * (getFreq(i, dx)**2.d0 + getFreq(j, dy)**2.d0) * dt)
		end do
	end do
	!$OMP END PARALLEL DO
	
	start = omp_get_wtime()
	do k=1,iterations
		!$OMP PARALLEL DO
		do j = 1, N
			do i = 1, N
				A(i,j) = T(i,j) * A(i,j)
			end do
		end do
		!$OMP END PARALLEL DO
		!A = T * A
	end do
	finish = omp_get_wtime()
	print '("Time = ",f12.3," seconds.")',finish-start
end program

! returns appropriate frequency for given index:
! i/N, if i <= N/2
! (i-N)/N = i/N - 1, if N/2 < i <= N
double precision function getFreq(i, d)
	integer, intent(in) :: i
	double precision, intent(in) :: d
	if (i.le.N/2) then
		getFreq = 2.d0 * pi * dble(i-1)/(d * dble(N))
	else
		getFreq = 2.d0 * pi * dble(i-N-1)/(d * dble(N))
	endif
end function