! Checkpoint and restart demonstration (May 2002)
! research.support@ualberta.ca

      program demo
      implicit none
      integer NX, NY, NZ, MAXSTEP, CPT
      parameter (NX=128, NY=128, NZ=128)
      parameter (MAXSTEP=500000, CPT=1000)
      real*8  data(NX,NY,NZ)
      integer step, laststep
      integer get_laststep

      laststep = get_laststep()
      if (laststep .eq. 0) then
         call initialize_data(data,NX,NY,NZ)
      else
         call read_checkpoint(laststep,CPT,data,NX,NY,NZ)
      end if

      do step = laststep + 1, MAXSTEP
         call do_work(data,NX,NY,NZ)
         if (mod(step,CPT) .eq. 0) then
            call write_checkpoint(step,CPT,data,NX,NY,NZ)
         end if
      end do

      call final_output(data,NX,NY,NZ)
      call clean_up()
      stop
      end


      integer function get_laststep()
      implicit none
      integer ierr
      open(99,file='restart',form='FORMATTED')
      read(99,*,iostat=ierr) get_laststep
      if (ierr .lt. 0) then
         get_laststep = 0
      end if
      close(99,status='KEEP')
      return
      end


      subroutine initialize_data(data,nx,ny,nz)
      implicit none
      integer nx, ny, nz
      real*8  data(nx,ny,nz)
      integer i, j, k
      real*8 rand
      do k = 1, nz
         do j = 1, ny
            do i = 1, nx
               data(i,j,k) = rand()
            end do
         end do
      end do
      return
      end


      subroutine read_checkpoint(laststep,cpt,data,nx,ny,nz)
      implicit none
      integer laststep, cpt, nx, ny, nz
      real*8  data(nx,ny,nz)
      integer phase, step
      character*4 cprfile
      phase = mod(laststep/cpt,2) + 1
      cprfile = 'cpr'
      write(cprfile(4:4),'(i1)') phase
      open(99,file=cprfile,form='UNFORMATTED')
      read(99) step, data
      close(99,status='KEEP')
      if (step .ne. laststep) then
         stop '### inconsistent checkpoint ###'
      end if
      return
      end


      subroutine write_checkpoint(step,cpt,data,nx,ny,nz)
      implicit none
      integer step, cpt, nx, ny, nz
      real*8  data(nx,ny,nz)
      integer phase
      character*4 cprfile
      phase = mod(step/cpt,2) + 1
      cprfile = 'cpr'
      write(cprfile(4:4),'(i1)') phase
      open(99,file=cprfile,form='UNFORMATTED')
      write(99) step, data
      close(99,status='KEEP')
      open(99,file='restart',form='FORMATTED')
      write(99,*) step
      close(99,status='KEEP')
      return
      end


      subroutine final_output(data,nx,ny,nz)
      implicit none
      integer nx, ny, nz
      real*8  data(nx,ny,nz)
      open(99,file='output',form='UNFORMATTED')
      write(99) data
      close(99,status='KEEP')
      return
      end


      subroutine clean_up()
      open(99,file='restart',form='FORMATTED')
      close(99,status='DELETE')
      open(99,file='cpr1',form='UNFORMATTED')
      close(99, status='DELETE')
      open(99,file='cpr2',form='UNFORMATTED')
      close(99, status='DELETE')
      return
      end


      subroutine do_work(data,nx,ny,nz)
      implicit none
      integer nx, ny, nz
      real*8  data(nx,ny,nz)
      integer i, j, k
!$omp parallel shared(data,nx,ny,nz) private(i,j,k)
!$omp do
      do k = 1, nz
         do j = 1, ny
            call xsweep(data(1,j,k),1,nx)
         end do
      end do
!$omp end do nowait
!$omp do
      do k = 1, nz
         do i = 1, nx
            call ysweep(data(i,1,k),nx,ny)
         end do
      end do
!$omp end do nowait
!$omp do
      do j = 1, ny
         do i = 1, nx
            call zsweep(data(i,j,1),nx*ny,nz)
         end do
      end do
!$omp end do
!$omp end parallel
      return
      end


      subroutine xsweep(v,is,n)
      implicit none
      integer is, n
      real*8  v(1+is*(n-1))
      integer i
      real*8 half
      parameter (half = 0.5d0)
      do i = 2, n
         v(1+is*(i-1)) = v(1+is*(i-1)) + half*v(1+is*(i-2))
      end do
      do i = n-1, 1, -1
         v(1+is*(i-1)) = v(1+is*(i-1)) - half*v(1+is*i)
      end do
      return
      end


      subroutine ysweep(v,is,n)
      implicit none
      integer is, n
      real*8  v(1+is*(n-1))
      integer i
      real*8 half
      parameter (half = 0.5d0)
      do i = 2, n
         v(1+is*(i-1)) = v(1+is*(i-1)) + half*v(1+is*(i-2))
      end do
      do i = n-1, 1, -1
         v(1+is*(i-1)) = v(1+is*(i-1)) - half*v(1+is*i)
      end do
      return
      end


      subroutine zsweep(v,is,n)
      implicit none
      integer is, n
      real*8  v(1+is*(n-1))
      integer i
      real*8 half
      parameter (half = 0.5d0)
      do i = 2, n
         v(1+is*(i-1)) = v(1+is*(i-1)) + half*v(1+is*(i-2))
      end do
      do i = n-1, 1, -1
         v(1+is*(i-1)) = v(1+is*(i-1)) - half*v(1+is*i)
      end do
      return
      end