#!/bin/bash
#SBATCH -A IscrC_RICK
#SBATCH -p boost_usr_prod
##SBATCH --qos boost_qos_bprod
#SBATCH -J RICK
### number of nodes
#SBATCH -N 1
### number of MPI tasks per node
#SBATCH --ntasks-per-node=4
#SBATCH -n 4
### number of openmp threads
#SBATCH --cpus-per-task=8
### number of allocated GPUs per node
#SBATCH --gpus-per-node=4
#SBATCH --mem=450G
#SBATCH -o test.out
#SBATCH -e test.err
#SBATCH -t 03:00:00 


module load openmpi/
module load fftw/
module load nvhpc/23.5
module load cuda/

export LD_LIBRARY_PATH="/leonardo/prod/spack/03/install/0.19/linux-rhel8-icelake/gcc-11.3.0/nvhpc-23.5-pdmwq3k5perrhdqyrv2hspv4poqrb2dr/Linux_x86_64/23.5/comm_libs/11.8/nvshmem_cufftmp_compat/lib/:$LD_LIBRARY_PATH"

export LD_LIBRARY_PATH="/leonardo/prod/spack/03/install/0.19/linux-rhel8-icelake/gcc-11.3.0/nvhpc-23.5-pdmwq3k5perrhdqyrv2hspv4poqrb2dr/Linux_x86_64/23.5/math_libs/11.8/lib64/:$LD_LIBRARY_PATH"

export OMPI_CC=gcc
export OMPI_CXX=g++

export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK}
export OMP_PLACES=cores

cd ../
make -j1 clean
rm -f w-stacking_fftw_acc-omp_acc-fft
make -j1 w-stacking

export typestring=omp_gpu_cufftmp
export exe=w-stacking_fftw_acc-omp_acc-fft

OUT_SHM=result_${SLURM_NTASKS}_${typestring}_${SLURM_CPUS_PER_TASK}
OUT_SHM_RES=/leonardo_scratch/large/userexternal/glacopo0/hpc_imaging/scripts/Tests/times_${SLURM_NTASKS}_${typestring}_${SLURM_CPUS_PER_TASK}_large

rm -f ${OUT_SHM} ${OUT_SHM_RES}

export logdir=mpi_${SLURM_NTASKS}_${typestring}_${SLURM_CPUS_PER_TASK}
echo "Creating $logdir"
rm -fr $logdir
mkdir $logdir

for itest in {1..3}
do
  export logfile=test_${itest}_${logdir}.log
  echo "time mpirun -np ${SLURM_NTASKS} --bind-to core --map-by ppr:${SLURM_NTASKS_PER_NODE}:node:pe=${SLURM_CPUS_PER_TASK} -x OMP_NUM_THREADS data/paramfile.txt" > $logfile
  time mpirun -np ${SLURM_NTASKS} --bind-to core --map-by ppr:${SLURM_NTASKS_PER_NODE}:node:pe=${SLURM_CPUS_PER_TASK} -x OMP_NUM_THREADS --mca btl self,vader /leonardo_scratch/large/userexternal/glacopo0/hpc_imaging/${exe} data/paramfile.txt >> $logfile
  mv $logfile $logdir
  mv timings.dat ${logdir}/timings_${itest}.dat
  cat ${logdir}/timings_all.dat ${logdir}/timings_${itest}.dat >> ${logdir}/timings_all.dat

  Reduce_time=$( grep -w 'Reduce time :' $logdir/$logfile | gawk '{print $4}' )
  FFTW_time=$( grep -w 'cufftMP time :' $logdir/$logfile | gawk '{print $4}' )
  Composition_time=$( grep -w 'Array Composition time :' $logdir/$logfile | gawk '{print $5}' )
  Writing_time=$( grep -w ' Image writing time :' $logdir/$logfile | gawk '{print $5}' )
  Total_time=$( grep -w 'TOT time :' $logdir/$logfile | gawk '{print $4}' )

  #Not relevant for the paper
  Setup_time=$( grep -w 'Setup time:' $logdir/$logfile | gawk '{print $3}' )
  Kernel_time=$( grep -w 'Kernel time :' $logdir/$logfile | gawk '{print $4}' )
  Phase_time=$( grep -w 'Phase time :' $logdir/$logfile | gawk '{print $4}' )
  ##########################

  
  echo $itest $Reduce_time $FFTW_time $Composition_time $Writing_time $Total_time $Setup_time $Kernel_time $Phase_time >> ${OUT_SHM}
done

echo -e "\n\n" >> ${OUT_SHM}
avg_red=$( awk '{sum+=$2} END { print sum/3 }' ${OUT_SHM} )
avg_fftw=$( awk '{sum+=$3} END { print sum/3 }' ${OUT_SHM} )
avg_comp=$( awk '{sum+=$4} END { print sum/3 }' ${OUT_SHM} )
avg_write=$( awk '{sum+=$5} END { print sum/3 }' ${OUT_SHM} )
avg_tot=$( awk '{sum+=$6} END { print sum/3 }' ${OUT_SHM} )

std_red=$( awk '{if($2!=""){count++;sum+=$2};y+=$2^2} END{sq=sqrt(y/3-(sum/3)^2);sq=sq?sq:0;print sq}' ${OUT_SHM} )
std_fftw=$( awk '{if($3!=""){count++;sum+=$3};y+=$3^2} END{sq=sqrt(y/3-(sum/3)^2);sq=sq?sq:0;print sq}' ${OUT_SHM} )
std_comp=$( awk '{if($4!=""){count++;sum+=$4};y+=$4^2} END{sq=sqrt(y/3-(sum/3)^2);sq=sq?sq:0;print sq}' ${OUT_SHM} )
std_write=$( awk '{if($5!=""){count++;sum+=$5};y+=$5^2} END{sq=sqrt(y/3-(sum/3)^2);sq=sq?sq:0;print sq}' ${OUT_SHM} )
std_tot=$( awk '{if($6!=""){count++;sum+=$6};y+=$6^2} END{sq=sqrt(y/3-(sum/3)^2);sq=sq?sq:0;print sq}' ${OUT_SHM} )


#Not relevant for the paper
avg_setup=$( awk '{sum+=$7} END { print sum/3 }' ${OUT_SHM} )
avg_ker=$( awk '{sum+=$8} END { print sum/3 }' ${OUT_SHM} )
avg_phase=$( awk '{sum+=$9} END { print sum/3 }' ${OUT_SHM} )

std_setup=$( awk '{if($7!=""){count++;sum+=$7};y+=$7^2} END{sq=sqrt(y/3-(sum/3)^2);sq=sq?sq:0;print sq}' ${OUT_SHM} )
std_ker=$( awk '{if($8!=""){count++;sum+=$8};y+=$8^2} END{sq=sqrt(y/3-(sum/3)^2);sq=sq?sq:0;print sq}' ${OUT_SHM} )
std_phase=$( awk '{if($9!=""){count++;sum+=$9};y+=$9^2} END{sq=sqrt(y/3-(sum/3)^2);sq=sq?sq:0;print sq}' ${OUT_SHM} )
##########################


echo "Averages and standard deviations over 3 shots" >> ${OUT_SHM_RES}
echo -e "\n" ${OUT_SHM_RES}
echo "${SLURM_NTASKS} MPI tasks;  ${SLURM_CPUS_PER_TASK} OpenMP threads per task;  ${SLURM_GPUS_PER_NODE} GPUs per node;" >> ${OUT_SHM_RES}
echo -e "\n\n" ${OUT_SHM_RES}
echo $avg_red $std_red $avg_fftw $std_fftw $avg_comp $std_comp $avg_write $std_write $avg_tot $std_tot >> ${OUT_SHM_RES}
echo -e "\n\n" ${OUT_SHM_RES}
echo $avg_setup $std_setup $avg_ker $std_ker $avg_phase $std_phase >> ${OUT_SHM_RES} 
rm -f ${OUT_SHM}
