module myshared
  use kinds
  implicit none

  real(kind=REAL8) :: total_error
end module myshared

module mysubs

contains

  subroutine getsrc(q)
    use kinds
    implicit none
    
    real(kind=REAL8), dimension(:,:)  :: q
    integer  :: n
    integer  :: d1, d2
    integer  :: x1, x2, y1, y2

    !
    !
    !                             |         |
    !                   U = 1  -> |   D2    | <-  U = -1
    !          Y2  ---------------|         |----------------
    !                                                      D1
    !
    !          Y1  ---------------|         |----------------
    !                   U = -1 -> |         | <-  U =  1
    !                             |         |  
    !             0               X1        X2
    !

    n  = size(q, 1)
    d1 = n/2
    d2 = n/4
    
    q = 0.0
    
    x1 = (n - d2)/2     
    x2 = (n + d2)/2
    y1 = (n - d1)/2     
    y2 = (n + d1)/2

    q(1:x1, y1) = -1.0
    q(1:x1, y2) =  1.0
    
    q(x2:n, y1) =  1.0
    q(x2:n, y2) = -1.0
      
    q(x1, 1:y1) = -1.0
    q(x2, 1:y1) =  1.0
    
    q(x1, y2:n) =  1.0
    q(x2, y2:n) = -1.0

    return
  end subroutine getsrc

  subroutine store(u)
    use kinds
    implicit none
    
    real(kind=REAL8), dimension(:,:)  :: u
    character(len=80), parameter      :: ofile = 'poisson.out'
    integer                           :: n, status, i, j

    n  = size(u, 1)
    open(1, file=ofile,iostat=status)
    if (status /= 0) then
       write (*,*) 'could not open file ', ofile
       stop
    end if
    write (1, *) n-2
    write (1,'(F10.5)') ((u(i,j), j=2, n-1), i=2, n-1)

    close(1)
    return
  end subroutine store

  function sweep(a, f, b)
    use kinds
    use myshared
    implicit none

    real(kind=REAL8)                  :: sweep
    real(kind=REAL8), dimension(:,:)  :: a, b, f
    real(kind=REAL8)                  :: resid, error
    integer                           :: n, i, j

    n  = size(a, 1)
    error = 0.0

    !$OMP DO
    do i = 2, n-1
       do j = 2, n-1
          resid = a(i-1,j) + a(i+1,j) + a(i,j+1) + a(i,j-1) + &
                 (-4.0)*a(i,j) - f(i,j)
          error = error + abs(resid)
          b(i,j) = a(i,j) + 0.25*resid
       end do
    end do
    !$OMP END DO

    ! sum local errors
    !$OMP SINGLE
    total_error = 0.0
    !$OMP END SINGLE
    
    !$OMP CRITICAL
    total_error = total_error + error
    !$OMP END CRITICAL
    !$OMP BARRIER

    sweep = total_error
    return
  end function sweep

end module mysubs


program poisson4
  ! solve Poisson equation using Jacobi relaxation

  use kinds
  use mysubs
  implicit none

  integer, parameter          :: n      = 200
  integer, parameter          :: maxits = 1000
  real(kind=REAL8), parameter :: eps    = 1.0e-4

  real(kind=REAL8), dimension(n,n) :: q, u1, u2
  real(kind=REAL8)                 :: error
  integer                          :: it

  call getsrc(q)        ! compute matrix of sources
  u1 = 0.0              ! initial solution
  u2 = 0.0

  !$OMP PARALLEL PRIVATE(error)
  do it = 1, maxits
     error = sweep(u1, q, u2)
     error = sweep(u2, q, u1)
     if (error/(n*n) < eps) exit
  end do

  !$OMP SINGLE
  print *, 'error after ', it, ' iterations: ', error/(n*n)
  !$OMP END SINGLE
  !$OMP END PARALLEL

  call store(u1)          ! store solution

  stop
end program poisson4