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

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
        
        real(8) :: start_time, end_time, OMP_get_wtime
        
        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
        
        ! Count time for SOR inversion
        start_time=OMP_get_wtime()
        
        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
        
        end_time=OMP_get_wtime()
        
        write(*,*) ' Difference is ', err1
        write(*,*) ' Walltime is ', end_time-start_time
        
        
! **************************************************************
! 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_val)
        implicit none

        integer :: Nx,Ny,i,j
        double precision, dimension(1:Nx,1:Ny) :: C, C_old
  
        integer :: large,tid,numthreads,OMP_GET_THREAD_NUM,OMP_GET_NUM_THREADS
        parameter (large=100)
        double precision :: diff_vec(0:large),diff_val

!$omp parallel
        numthreads = OMP_GET_NUM_THREADS()
!$omp end parallel

!$omp parallel default(shared), private(i,j,tid,diff_val)
        tid=OMP_GET_THREAD_NUM()
        diff_val=0.d0
!$omp do  
        Do j = 1, Ny
          Do i = 1, Nx
            diff_val = diff_val + (C(i,j)-C_old(i,j))**2
          End Do
        End Do
!$omp end do
        diff_vec(tid)=diff_val
        write(*,*) 'tid= ',tid,'diff= ',diff_val, 'num threads= ',numthreads
!$omp end parallel

        diff_val=0.d0
        do tid=0,numthreads-1
          if(diff_vec(tid).gt.diff_val)diff_val=diff_vec(tid)
        end do
  
        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