#!/bin/bash
#
# This script sets up and submits job of global WW3 simulation
# using multiple grids on Yellowstone/Cheyenne.
#
# Qing Li, 20170713
#		   20170717 set up for a grid set of 3

## load shared functions ------------------------------------------------------
source ./set_env
source ./share/bash_functions.sh

## Run control ----------------------------------------------------------------
# case name
# CURVILINEAR GRID SUPPORT NOT IMPLEMENTED FOR NRGRD > 1!
# Should try curvilinear at north polar region and regular elsewhere
gridtype="mww3a"  # grid type: mww3a
if [[ ${gridtype} == mww3a ]]; then
	grid1="gx16b_spr"
	grid2="ww3a_eq"
	grid3="gx16b_np"
	grids="${grid1} ${grid2} ${grid3}"  # combination of gx16b_sp, gx37b_eq and gx16b_np
	inputgrid="" # curvilinear not supported for input grid
	targetgrid="gx16b"
else
	echo "Grid type ${gridtype} not supported. Stop."
	exit 1
fi
casename="testMulti_200901_${gridtype}"

# run times
# (nonzero hhmmss not supported)
t_beg="20090101 000000"  # first output
t_end="20090131 000000"  # last output
dt="21600"               # output interval
dtrs="86400"             # restart interval

# restart run or not
f_restart="no"
# clean old exes and build new ones
f_new="yes"
# update grid
f_grid="yes"
# input file options: core2update, core2copy, cice, none
opt_input="cice"

# number of processors
nproc=64
# wall time (hh:mm)
walltime="01:00"
# queue (small, regular, economy)
queue="regular"
# project number
#projnum="UBRN0005"

# set switch
case_switch_base="F90 NOGRB NOPA LRB4 NC4 SCRIP SCRIPNC PR3 UQ FLX0 LN1 ST4 BT1 DB1 MLIM "\
"NL1 TR0 BS0 IC2 IS2 REF0 XX0 WNT1 WNX1 CRT1 CRX1 O0 O1 O2 O3 O4 O5 O6 O7 O11"
case_switch_ser="${case_switch_base} SHRD"
case_switch_mpi="${case_switch_base} MPI DIST"


# set directory
tmp_dir="/glade/scratch/aordonez/wwatch3v5p16"
toolbox_dir=${WW3_TOOLBOX_ROOT}

# set environment
export WWATCH3_NETCDF="NC4"
if [[ ${HOSTNAME} == yslogin* ]]; then
	ww3_dir="/glade/u/home/aordonez/wwatch3_v5.16"
	projnum="UWAS0039"
	export NETCDF_CONFIG="/glade/apps/opt/netcdf/4.3.0/intel/12.1.5/bin/nc-config"
    # specify compiler string (as per in ww3_dir/bin options) here:
    compstr="Intel"
elif [[ ${HOSTNAME} == cheyenne* ]]; then
	ww3_dir="/glade/u/home/aordonez/wwatch3_v5.16_cheyenne"
	projnum="UWAS0039"
	export NETCDF_CONFIG="/glade/u/apps/ch/opt/netcdf/4.4.1.1/intel/16.0.3/bin/nc-config"
    # specify compiler string (as per in ww3_dir/bin options) here:
    compstr="Cheyenne_Intel"
else
	echo "Machine ${HOSTNAME} not supported. Stop."
	exit 1
fi
export WWATCH3_ENV="${ww3_dir}/wwatch3.env"


## 0. Preprocessing -----------------------------------------------------------
print_title "Preprocessing"

# calculate the output file number
date_beg=${t_beg:0:8}
date_end=${t_end:0:8}
nsec=$(( $(date --date=${date_end} +%s) - $(date --date=${date_beg} +%s) ))
nday=$(( ${nsec}/86400 ))
nfile=$(( ${nsec}/${dt} ))
nres=$(( ${nsec}/${dtrs} ))

# spatial dimensions
if [[ ${targetgrid} == gx37 ]]; then
	nx=100
	ny=116
elif [[ ${targetgrid} == gx16* ]]; then
	nx=320
	ny=384
elif [[ ${targetgrid} == ww3a ]]; then
	nx=90
	ny=50
else
	echo "Grid type ${targetgrid} not supported"
	exit 1
fi

# set path
path_e="${ww3_dir}/exe"            # path for executables
path_b="${ww3_dir}/bin"            # path for binaries
path_a="${ww3_dir}/aux"            # path for aux files and scripts
path_work="${tmp_dir}/${casename}" # work directory
path_dat="${toolbox_dir}/data"     # input data directory
path_out="${path_work}/out"        # output directory
path_res="${path_work}/restart"    # restart directory
path_exe="${path_work}/exe"        # directory executables for each case
path_input="${tmp_dir}/data"       # input data directory for copy
path_scpt="${toolbox_dir}/bin"     # path for compiler and linker scripts

# setup work directory
mkdir -p ${path_work}
mkdir -p ${path_exe}
cd ${path_work}

# setup output directory
mkdir -p ${path_out}

# setup restart directory
mkdir -p ${path_res}

# setup compiler and linker script
cp ${path_scpt}/comp.${compstr} ${path_b}/comp
cp ${path_scpt}/link.${compstr} ${path_b}/link

# clean old files and create new case
if [[ ${f_new} == yes ]]; then
	${path_b}/w3_clean all
	${path_b}/w3_new
	# check error
	check_error "Error in w3_new. Stop..."
fi


## 1. Grid pre-processor ------------------------------------------------------
if [[ ${f_grid} == yes ]]; then
	print_title "Grid preprocessor"

	# compile appropriate code
	echo ${case_switch_ser} > ${path_b}/switch
	${path_b}/w3_make ww3_grid
	# check error
	check_error "Error in compiling ww3_grid. Stop..."
	cp ${path_e}/ww3_grid ${path_exe}/

	rm -f mod_def.*
	path_grid="${toolbox_dir}/grids/grid_inp"
	mkdir -p grid_inp

	# loop over grids
	for grd in ${grids} ${inputgrid} ${targetgrid}
	do
	if [ ! -f ./mod_def.${grd} ]; then
		cp ${path_grid}/${grd}_bottom.inp ./
		cp ${path_grid}/${grd}_mapsta.inp ./
		if [ -f ${path_grid}/${grd}_x.inp ]; then
			cp ${path_grid}/${grd}_x.inp ./
		fi
		if [ -f ${path_grid}/${grd}_y.inp ]; then
			cp ${path_grid}/${grd}_y.inp ./
		fi
		cp ${path_grid}/ww3_grid.inp.${grd} ./ww3_grid.inp
		echo "  Screen output routed to ${path_out}/ww3_grid_${grd}.out"
		${path_exe}/ww3_grid > ${path_out}/ww3_grid_${grd}.out
		mv ${grd}_bottom.inp ${grd}_mapsta.inp ww3_grid.inp ${path_work}/grid_inp
		if [ -f ${grd}_x.inp ]; then
			mv ${grd}_x.inp ${path_work}/grid_inp
		fi
		if [ -f ${grd}_y.inp ]; then
			mv ${grd}_y.inp ${path_work}/grid_inp
		fi
		mv mod_def.ww3 mod_def.${grd}
	fi
	done
fi


## 2. Initial conditions ------------------------------------------------------
print_title "Initial conditions"

if [[ ${f_restart} == yes ]]; then
	file_res="${path_res}/restart_${date_beg}.ww3"
	if [[ -e ${file_res} ]]; then
		echo "Restarting run from ${file_res}..."
		cp ${file_res} ./restart.ww3
	else
		echo "Restart file ${file_res} not found. Stop."
		exit 1
	fi
else
	echo "Starting from initial condition...Preparing..."
	# compile appropriate code
	echo ${case_switch_ser} > ${path_b}/switch
	${path_b}/w3_make ww3_strt
	# check error
	check_error "Error in compiling ww3_strt. Stop..."
	cp ${path_e}/ww3_strt ${path_exe}/
	path_strt="${toolbox_dir}/data/strt_inp"
	cp ${path_strt}/ww3_strt.inp ./

	for grd in ${grids}
	do
		rm -f mod_def.ww3
		ln -s mod_def.${grd} mod_def.ww3
		echo "  Screen output routed to ${path_out}/ww3_strt.${grd}.out"
		${path_exe}/ww3_strt > ${path_out}/ww3_strt.${grd}.out
		mv restart.ww3 restart.${grd}
	done
	mkdir -p strt_inp
	mv ww3_strt.inp strt_inp/
fi


## 3. Input fields ------------------------------------------------------------
# update input files
if [[ ${opt_input} == core2update ]]; then
	print_title "Input fields"

	# compile appropriate code
	echo ${case_switch_ser} > ${path_b}/switch
	${path_b}/w3_make ww3_prnc
	# check error
	check_error "Error in compiling ww3_prnc. Stop..."
	cp ${path_e}/ww3_prnc ${path_exe}
	path_prnc="${toolbox_dir}/data/prnc_inp"

	# prepare input fields
	rm -f mod_def.ww3
	ln -s mod_def.${inputgrid} mod_def.ww3
	# wind
	ln -s ${path_dat}/wind.nc ./
	cp ${path_prnc}/ww3_prnc.inp.wind ./ww3_prnc.inp
	echo "  Screen output routed to ${path_out}/ww3_prnc_wind.out"
	${path_exe}/ww3_prnc > ${path_out}/ww3_prnc_wind.out
	mkdir -p prnc_inp
	mv ww3_prnc.inp prnc_inp/ww3_prnc_wind.inp
	mv wind.ww3 wind.${inputgrid}
	rm -f wind.nc

	# ice
	ln -s ${path_dat}/ice.nc ./
	cp ${path_prnc}/ww3_prnc.inp.ice ./ww3_prnc.inp
	echo "  Screen output routed to ${path_out}/ww3_prnc_ice.out"
	${path_exe}/ww3_prnc > ${path_out}/ww3_prnc_ice.out
	mv ww3_prnc.inp prnc_inp/ww3_prnc_ice.inp
	mv ice.ww3 ice.${inputgrid}
	rm -f ice.nc mod_def.ww3
elif [[ ${opt_input} == core2copy ]]; then
	for grd in ${grids}
	do
		rm -f wind.${grd} ice*.${grd}
		ln -s ${path_input}/wind.${grd} ./
		ln -s ${path_input}/ice*.${grd} ./
	done
elif [[ ${opt_input} == cice ]]; then
	print_title "Input fields"

	# compile appropriate code
	echo ${case_switch_ser} > ${path_b}/switch
	${path_b}/w3_make ww3_prnc
	# check error
	check_error "Error in compiling ww3_prnc. Stop..."
	cp ${path_e}/ww3_prnc ${path_exe}
	path_prnc="${toolbox_dir}/data/prnc_inp"
	mkdir -p prnc_inp

	# wind
	ln -s ${path_dat}/wind.nc ./
	cp ${path_prnc}/ww3_prnc.inp.wind ./ww3_prnc.inp
	for grd in ${grids}
	do
		rm -f mod_def.ww3
		ln -s mod_def.${grd} mod_def.ww3
		echo "  Screen output routed to ${path_out}/ww3_prnc_wind.${grd}.out"
		${path_exe}/ww3_prnc > ${path_out}/ww3_prnc_wind.${grd}.out
		mv wind.ww3 wind.${grd}
	done
	mv ww3_prnc.inp prnc_inp/ww3_prnc_wind.inp
	rm -f wind.nc

	# ice fraction
	ln -s ${path_dat}/cice.nc ./
	cp ${path_prnc}/ww3_prnc.inp.cice.ice ./ww3_prnc.inp
	for grd in ${grids}
	do
		rm -f mod_def.ww3
		ln -s mod_def.${grd} mod_def.ww3
		echo "  Screen output routed to ${path_out}/ww3_prnc_cice_ice.${grd}.out"
		${path_exe}/ww3_prnc > ${path_out}/ww3_prnc_cice_ice.${grd}.out
		mv ice.ww3 ice.${grd}
	done
	mv ww3_prnc.inp prnc_inp/ww3_prnc_cice_ice.inp

	# ice thickness
	cp ${path_prnc}/ww3_prnc.inp.cice.ic1 ./ww3_prnc.inp
	for grd in ${grids}
	do
		rm -f mod_def.ww3
		ln -s mod_def.${grd} mod_def.ww3
		echo "  Screen output routed to ${path_out}/ww3_prnc_cice_ic1.${grd}.out"
		${path_exe}/ww3_prnc > ${path_out}/ww3_prnc_cice_ic1.${grd}.out
		mv ice1.ww3 ice1.${grd}
	done
	mv ww3_prnc.inp prnc_inp/ww3_prnc_cice_ic1.inp
	rm -f cice.nc mod_def.ww3
elif [[ ${opt_input} == none ]]; then
	echo "Do nothing on input file. Please make sure wind.ww3 and ice.ww3 exist."
else
	echo "Invalid input option ${opt_input}. Stop."
	exit 1
fi


## 4. Main program ------------------------------------------------------------
print_title "Main program"

# compile appropriate code
echo ${case_switch_mpi} > ${path_b}/switch
${path_b}/w3_make ww3_multi
# check error
check_error "Error in compiling ww3_multi. Stop..."
cp ${path_e}/ww3_multi ${path_exe}/

# processing input file
path_multi="${toolbox_dir}/run/multi_inp"
cp ${path_multi}/ww3_multi.inp ./
find_replace LB_STARTING_TIME "${t_beg}" ww3_multi.inp
find_replace LB_ENDING_TIME "${t_end}" ww3_multi.inp
find_replace LB_TIME_INT ${dt} ww3_multi.inp
find_replace LB_TIME_RES ${dtrs} ww3_multi.inp
#find_replace LB_INPUT_GRID ${inputgrid} ww3_multi.inp
find_replace LB_GRID1 ${grid1} ww3_multi.inp
find_replace LB_GRID2 ${grid2} ww3_multi.inp
find_replace LB_GRID3 ${grid3} ww3_multi.inp
mkdir -p multi_inp


## 5. Grid integration --------------------------------------------------------
print_title "Grid integration"

# compile appropriate code
echo ${case_switch_ser} > ${path_b}/switch
${path_b}/w3_make ww3_gint
# check error
check_error "Error in compiling ww3_gint. Stop..."
cp ${path_e}/ww3_gint ${path_exe}/

path_ounf="${toolbox_dir}/postprocess/gint_inp"
cp ${path_ounf}/ww3_gint.inp ./
find_replace LB_STARTING_TIME "${t_beg}" ww3_gint.inp
find_replace LB_TIME_INT ${dt} ww3_gint.inp
find_replace LB_NFILE ${nfile} ww3_gint.inp
find_replace LB_GRID1 ${grid1} ww3_gint.inp
find_replace LB_GRID2 ${grid2} ww3_gint.inp
find_replace LB_GRID3 ${grid3} ww3_gint.inp
find_replace LB_TARGET_GRID ${targetgrid} ww3_gint.inp
mkdir -p gint_inp


## 6. Gridded netCDF output post-processor ------------------------------------
print_title "Gridded netCDF output post-processor"

# compile appropriate code
echo ${case_switch_ser} > ${path_b}/switch
${path_b}/w3_make ww3_ounf
# check error
check_error "Error in compiling ww3_ounf. Stop..."
cp ${path_e}/ww3_ounf ${path_exe}/

# processing input file
path_ounf="${toolbox_dir}/postprocess/ounf_inp"
cp ${path_ounf}/ww3_ounf.inp ./
find_replace LB_STARTING_TIME "${t_beg}" ww3_ounf.inp
find_replace LB_TIME_INT ${dt} ww3_ounf.inp
find_replace LB_NFILE ${nfile} ww3_ounf.inp
find_replace LB_NX ${nx} ww3_ounf.inp
find_replace LB_NY ${ny} ww3_ounf.inp
mkdir -p ounf_inp

# restart file
nnn_res=`printf %03d ${nres}`
sav_res="restart${nnn_res}.ww3"
rname_res="restart_${date_end}.ww3"


## 7. Submitting job ----------------------------------------------------------
print_title "Submitting job"

echo "  Running WW3 with ${nproc} processors..."

# generate machine-specific submission script
if [[ ${HOSTNAME} == yslogin* ]]; then
	exempi="mpirun.lsf"
# write LSF header for yellowstone
	cat > submit.sh << EOF
#!/bin/bash
#BSUB -n ${nproc}
#BSUB -q ${queue}
#BSUB -R "span[ptile=15]"
#BSUB -N
#BSUB -W ${walltime}
#BSUB -x
#BSUB -J ${casename}
#BSUB -e ${casename}.err
#BSUB -o ${casename}.out
#BSUB -P ${projnum}

module load intel
module load impi

EOF
# write postprocessing script for yellowstone
	cat > postproc.sh << EOF
#!/bin/bash

EOF

elif [[ ${HOSTNAME} == cheyenne* ]]; then
	nnode=$(( ${nproc} / 36 + 1 ))
	if [ ${nnode} -eq 1 ]; then
		ncpu=${nproc}
	else
		ncpu=$(( ${nproc} / ${nnode} ))
	fi
	exempi="mpiexec_mpt"
# write PBS header for cheyenne
	cat > submit.sh << EOF
#!/bin/bash
#PBS -N ${casename}
#PBS -j oe
#PBS -q ${queue}
#PBS -l walltime=${walltime}:00
#PBS -A ${projnum}
#PBS -l select=${nnode}:ncpus=${ncpu}:mpiprocs=${ncpu}

module load intel
module load mpt

export MPI_SHEPHERD=true
EOF
# write postprocessing script for cheyenne
	cat > postproc.sh << EOF
#!/bin/bash
#PBS -N pp_${casename}
#PBS -j oe
#PBS -q ${queue}
#PBS -l walltime=00:15:00
#PBS -A ${projnum}
#PBS -l select=1:ncpus=1:mpiprocs=1

EOF
fi # check hostname

# write command
# mpi run
if [[ ${HOSTNAME} == yslogin* ]]; then
	cat >> submit.sh << EOF
cd ${path_work}
${exempi} ${path_exe}/ww3_multi > ${path_out}/ww3_multi.out
./postproc.sh
EOF
elif [[ ${HOSTNAME} == cheyenne* ]]; then
	cat >> submit.sh << EOF
cd ${path_work}
${exempi} ${path_e}/ww3_multi > ${path_out}/ww3_multi.out
qsub -W depend=afterok:\$PBS_JOBID ./postproc.sh
EOF
fi

# command for postprocessing script
cat >> postproc.sh << EOF
echo "  Log files routed to ${path_out}"
mv log.* ${path_out}/
mv ww3_multi.inp multi_inp/

${path_exe}/ww3_gint > ${path_out}/ww3_gint.out
mv ww3_gint.inp gint_inp/
${path_exe}/ww3_ounf > ${path_out}/ww3_ounf.out
mv ww3_ounf.inp ounf_inp/
mv ww3.*.nc out/
mv ${sav_res} restart/${rname_res}

EOF


# submitting job
chmod 755 submit.sh
chmod 755 postproc.sh
if [[ ${HOSTNAME} == yslogin* ]]; then
	bsub < submit.sh
elif [[ ${HOSTNAME} == cheyenne* ]]; then
	qsub submit.sh
fi