module mysubs

contains

  subroutine init(a)
    ! initial values for a: a(i,j, k) = i+j+k

    use kinds
    implicit none
    
    real(kind=REAL8), dimension(:,:,:)  :: a
    integer                             :: n, i, j, k
 
    n = size(a, 1)
    
    do i=1,n
       do j=1,n
          do k=1,3
             a(i,j,k) = i + j + k
          enddo
       enddo
    enddo
    return
  end subroutine init

  subroutine check(a)
    ! check results
    ! here: print total sum
    
    use kinds
    implicit none
    
    real(kind=REAL8), dimension(:,:,:)  :: a
    real(kind=REAL8)                    :: sum1, sum2
    integer                             :: n
    
    n = size(a, 1)

    sum1 = sum(a)
    sum2 = 3*n**3 + 9*n**2
    
    print *, 'sum is       : ', sum1
    print *, 'and should be: ', sum2
    
    return
  end subroutine check

  function f(x, work)
    ! time consuming way of setting f = x + eps
    use kinds
    implicit none

    real(kind=REAL8) :: f, x
    integer          :: work
    integer          :: i
    real(kind=REAL8) :: t1, t2
    
    t1 = x
    do i= 1, work
       t1 = cos(t1)
    end do
    
    t2 = x
    do i= 1, work
       t2 = sin(t2)
    end do
  
    f = x + t1 + t2 - 0.7390851332151607 
  end function f


  subroutine smooth(a, direction, work)
    ! smoothes part of a
    use kinds
    implicit none

    real(kind=REAL8), dimension(:,:,:) :: a
    integer          :: direction, work
    
    integer          :: i, j, k, n
    
    n = size(a, 1)
    k = direction

    !$OMP PARALLEL DO
    do i=2, n-1
       do j=2, n-1
          a(i,j,k) = (a(i-1,j,k) + f(a(i,j,k),work) + a(i+1,j,k) + &
                      a(i,j-1,k) + a(i,j+1,k))/5.0
       end do
    end do
    !$OMP END PARALLEL DO

    return
  end subroutine smooth


end module mysubs


program section2

  use kinds
  use mysubs
  implicit none

  integer, parameter                 :: n = 20
  integer, parameter                 :: work = 100000

  real(kind=REAL8), dimension(n,n,3) :: a
      
  call init(a)   ! initial values for a

  !$OMP PARALLEL SECTIONS
  !$OMP SECTION
  call smooth(a, 1, work)
  !$OMP SECTION
  call smooth(a, 2, work)
  !$OMP SECTION
  call smooth(a, 3, work)
  !$OMP END PARALLEL SECTIONS

  call check(a)  ! check result

end program section2