! **************************************************************

program mainprogram
        implicit none
        
        integer :: Nx,Ny
        parameter (Nx = 201, Ny = 101)
               
        double precision :: dx,dy,x_val,y_val,Lx,Ly,pi=3.1415926535
        double precision :: ax,ay,diag_val,relax,tempval,err1,err2,A0
        double precision, dimension(1:Nx,1:Ny) :: f_source,C,C_old
        
        integer :: i,j,im1,ip1,iteration,max_iteration=1000
        
        Lx=2.d0
        Ly=1.d0
        
        A0=10.d0
        
        dx=Lx/dble(Nx-1)
        dy=Ly/dble(Ny-1)

! **************************************************************
! compute source, initialise guess

        f_source=0.d0
        C=0.d0
        C_old=0.d0
        
        write (*,*) 'getting source'
        call get_f_periodic(f_source,Nx,Ny,dx,dy,Lx,Ly,A0)
        write (*,*) 'done'
        
! **************************************************************
! sor steps
        
        do iteration=1,max_iteration
          
          ! for keeping track of the error
!$omp parallel workshare
            C_old=C
!$omp end parallel workshare

            ! First sweep
            call do_sor_C(C,f_source,dx,dy,Nx,Ny,0)
            ! Second sweep
            call do_sor_C(C,f_source,dx,dy,Nx,Ny,1)
            
            ! Implement Dirichlet conditions at y=0,y=L_y.
            C(:,1)=C(:,2)
            C(:,Ny)=C(:,Ny-1)

          if(mod(iteration,100)==0)then
            call get_diff(C,C_old,Nx,Ny,err1)
            write(*,*) iteration, ' Difference is ', err1
          end If
        
        end do
        
        write(*,*) ' Difference is ', err1
        
        
! **************************************************************
! write result to file
        
          write(*,*) 'writing to file'
          open(unit=20,file='poisson.dat',status='UNKNOWN')
          
          do j=1,Ny
            do i=1,Nx
              x_val=(i-1)*dx
              y_val=(j-1)*dy
              Write(20,*) x_val, y_val, C(i,j)
            end do
          end do
          close(unit=20, status='KEEP')
          write(*,*) 'done'
      
end program mainprogram


! **************************************************************
! **************************************************************

  subroutine get_f_periodic(f_src,Nx,Ny,dx,dy,Lx,Ly,A0)
  implicit none
  
  integer ::i,j,Nx,Ny
  double precision :: dx,dy,Lx,Ly,x_val,y_val, pi=3.1415926535
  double precision :: kx0,ky0,kx,ky,A0
  double precision :: f_src(1:Nx,1:Ny)
  
  kx0=2.d0*pi/Lx
  ky0=pi/Ly
  
  kx=kx0
  ky=3.d0*ky0
  
  f_src=0.d0
  do j=1,Ny
    do i=1,Nx
      x_val=(i-1)*dx
      y_val=(j-1)*dy
      f_src(i,j)=A0*cos(kx*x_val)*cos(ky*y_val)
    end do
  end do
  
  return
  end subroutine get_f_periodic
  
! **************************************************************

  subroutine get_diff(C,C_old,Nx,Ny,diff)
  implicit none

  double precision :: diff,sum
  integer :: Nx,Ny,i,j
  double precision, dimension(1:Nx,1:Ny) :: C, C_old

  sum = 0.0D0
!$omp parallel do default(none), &
!$omp   private(i,j), shared(C,C_old,Ny,Nx), &
!$omp   reduction(+:sum)
  Do j = 1, Ny
    Do i = 1, Nx
      sum = sum + (C(i,j)-C_old(i,j))**2
    End Do
  End Do
!$omp end parallel do
  diff = sum
  
  Return
  End subroutine get_diff

! **************************************************************
! **************************************************************
! ***************************************************************

        subroutine do_sor_C(C,f_source,dx,dy,Nx,Ny,flag)
        implicit none
  
        integer ::i,j,ip1,im1,Nx,Ny,flag
        double precision :: dx,dy,ax,ay
        double precision :: f_source(1:Nx,1:Ny),C(1:Nx,1:Ny)
        
        double precision :: relax,diag_val,tempval

        ax=1.d0/(dx*dx)
        ay=1.d0/(dy*dy)
        diag_val=2.d0*ax+2.d0*ay
        relax=1.5d0

!$omp parallel default(shared), private(i,j,im1,ip1,tempval)
!$omp do      
        do j=2,Ny-1
          do i=mod(j+flag,2)+1,Nx,2
          
             if(i.eq.1)then
                im1=Nx-1
             else
                im1=i-1
             end if

              if(i.eq.Nx)then
                ip1=2
              else
                ip1=i+1
              end if
          
              tempval=ax*(C(ip1,j)+C(im1,j))+ay*(C(i,j+1)+C(i,j-1))-f_source(i,j)
              C(i,j)=(1-relax)*C(i,j)+relax*tempval/diag_val
         
            end do
          end do
!$omp end do
!$omp end parallel

        end subroutine do_sor_C