#!/bin/sh -e

MYSCRIPT=`basename $0 .sh`

source /cluster/installations/lmod/lmod/init/sh 
module --quiet restore system
module load StdEnv
module load iompi/2020b
module load netCDF-Fortran/4.5.3-iompi-2020b 

cat <<EOF> ${MYSCRIPT}.f
      program test
      use netcdf
      implicit none
      include 'mpif.h'
      integer i,rank, size, ierror, tag, status(MPI_STATUS_SIZE)
      integer ncid, dimid, varid
      integer, parameter :: NX=1000
      real(kind=8) ::  x(NX)
      character*12 message
        
      call MPI_INIT(ierror)
      call MPI_COMM_SIZE(MPI_COMM_WORLD, size, ierror)
      call MPI_COMM_RANK(MPI_COMM_WORLD, rank, ierror)
      write(*,*) MPI_REAL, MPI_INTEGER
c     maximum allowable tag found by try&error 
      tag = 2097151

      if (rank .eq. 0) then
         x=1.  
         status = nf90_create('test.nc', NF90_CLOBBER, ncid) 
         status = nf90_def_dim(ncid, 'x', NX, dimid)
         status = nf90_def_var(ncid, 'data', NF90_DOUBLE, dimid, varid)
         status = nf90_enddef(ncid) 
         status = nf90_put_var(ncid, varid, x)
         status = nf90_close(ncid) 
         message = 'Hello, world'
         do i=1, size-1
            call MPI_SEND(message, 12, MPI_CHARACTER, i, tag,
     &                    MPI_COMM_WORLD, ierror)
         enddo
      else
         call MPI_RECV(message, 12, MPI_CHARACTER, 0, tag,
     &                    MPI_COMM_WORLD, status, ierror)
         status = nf90_open('test.nc', NF90_NOWRITE, ncid) 
         status = nf90_inq_varid(ncid, 'data', varid)
         status = nf90_get_var(ncid, varid, x)
         status = nf90_close(ncid) 
      endif
 
      print*, 'Task', rank, ':', message, x(1)
      call MPI_FINALIZE(ierror)
      end
EOF

mpifort -O2 -lnetcdf -lnetcdff -o ${MYSCRIPT} ${MYSCRIPT}.f 


cat <<EOF> ${MYSCRIPT}.slurm
#! /bin/sh -
#SBATCH --account=nn9039k
#SBATCH --ntasks=512
#SBATCH --nodes=4
#SBATCH --time=00:10:00
#SBATCH --error=${MYSCRIPT}.err
#SBATCH --output=${MYSCRIPT}.out


source /cluster/installations/lmod/lmod/init/sh
module --quiet restore system
module load StdEnv
module load iompi/2020b
module load netCDF-Fortran/4.5.3-iompi-2020b
ulimit -s unlimited 

cd `pwd`

time mpirun -n 512 ./${MYSCRIPT}
EOF

rm -f ${MYSCRIPT}.out ${MYSCRIPT}.err
chmod +x ${MYSCRIPT}.slurm
sbatch ${MYSCRIPT}.slurm