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

program mainprogram
        implicit none
        
        integer :: Nx,Ny,Nz
        parameter (Nx = 201, Ny = 201, Nz=101)
               
        double precision :: dx,dy,dz,x_val,y_val,z_val,Lx,Ly,Lz,pi=3.1415926535
        double precision :: A0,err1
        double precision, allocatable, dimension(:,:,:) :: f_source,C,C_old
        
        integer :: i,j,k,iteration,max_iteration=1000
        
        real(8) :: start_time, end_time, OMP_get_wtime
        
        Lx=2.d0
        Ly=2.d0
        Lz=1.d0
        
        A0=10.d0
        
        dx=Lx/dble(Nx-1)
        dy=Ly/dble(Ny-1)
        dz=Lz/dble(Nz-1)

! **************************************************************
! Allocate variables

        allocate(f_source(1:Nx,1:Ny,1:Nz),C(1:Nx,1:Ny,1:Nz),C_old(1:Nx,1:Ny,1:Nz))
        
        f_source=0.d0
        C=0.d0
        C_old=0.d0
        
! **************************************************************
! compute source, initialise guess

        write (*,*) 'getting source'
        call get_f_periodic(f_source,Nx,Ny,Nz,dx,dy,dz,Lx,Ly,Lz,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,dz,Nx,Ny,Nz,0)
            ! Second sweep
            call do_sor_C(C,f_source,dx,dy,dz,Nx,Ny,Nz,1)
            
            ! Implement Neumann conditions at z=0,z=L_z.
            C(:,:,1) =C(:,:,2)
            C(:,:,Nz)=C(:,:,Nz-1)

          if(mod(iteration,100)==0)then
            call get_diff(C,C_old,Nx,Ny,Nz,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 3d data to file'
          open(unit=20,file='poisson.dat',status='UNKNOWN')
          
          do k=1,Nz
            do j=1,Ny
              do i=1,Nx
                x_val=(i-1)*dx
                y_val=(j-1)*dy
                z_val=(k-1)*dz
                Write(20,*) x_val, y_val, z_val,C(i,j,k)
              end do
            end do
          end do
          close(unit=20, status='KEEP')
          write(*,*) 'done'
          
          j=floor(dble(Nz)/2.d0)
          write(*,*) 'outputting slice to file, with j= ',j
          open(unit=20,file='poisson_slice.dat',status='UNKNOWN')
          do k=1,Nz
            do i=1,Nx
              x_val=(i-1)*dx
              z_val=(k-1)*dz
              Write(20,*) x_val, z_val,C(i,j,k)
            end do
          end do
          close(unit=20, status='KEEP')
          write(*,*) 'done'

! **************************************************************
! De-allocate variables:

        deallocate(f_source,C,C_old)
      
end program mainprogram


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

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

        subroutine get_diff(C,C_old,Nx,Ny,Nz,diff_val)
        implicit none

        integer :: Nx,Ny,Nz,i,j,k
        double precision, dimension(1:Nx,1:Ny,1:Nz) :: 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,k,tid,diff_val)
        tid=OMP_GET_THREAD_NUM()
        diff_val=0.d0
!$omp do  
        do k=1, Nz
          do j = 1, Ny
            do i = 1, Nx
              diff_val = diff_val + (C(i,j,k)-C_old(i,j,k))**2
            end do
          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,dz,Nx,Ny,Nz,flag)
        implicit none
  
        integer ::i,j,k,ip1,im1,jp1,jm1,Nx,Ny,Nz,flag
        double precision :: dx,dy,dz,ax,ay,az
        double precision :: f_source(1:Nx,1:Ny,1:Nz),C(1:Nx,1:Ny,1:Nz)
        
        double precision :: relax,diag_val,tempval

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

!$omp parallel default(shared), private(i,j,k,im1,ip1,jm1,jp1,tempval)
!$omp do
        do k=2,Nz-1
          do j=1,Ny
          
            if(j.eq.1)then
              jm1=Ny-1
            else
              jm1=j-1
            end if

            if(j.eq.Ny)then
              jp1=2
            else
              jp1=j+1
            end if
              
            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,k)+C(im1,j,k))+ay*(C(i,j+1,k)+C(i,j-1,k))-f_source(i,j,k)+&
                   az*(C(i,j,k+1)+C(i,j,k-1))
              C(i,j,k)=(1-relax)*C(i,j,k)+relax*tempval/diag_val
            end do
          end do
        end do
!$omp end do
!$omp end parallel

        end subroutine do_sor_C