# Function leap days calculates the number of leap days (29th of Februrary) in
# a time intervall between two dates.
#
# Usage leap_days START_DATE END_DATE
function leap_days()
{
    local ld=0
    local frstYYYY=$(date -ud "$1" +%Y)
    local lastYYYY=$(date -ud "$2" +%Y)

    set +e

    # Check first year for leap day between start and end date
    $(date -ud "${frstYYYY}-02-29" > /dev/null 2>&1) \
    && (( $(date -ud "$1" +%s) < $(date -ud "${frstYYYY}-03-01" +%s) )) \
    && (( $(date -ud "$2" +%s) > $(date -ud "${lastYYYY}-02-28" +%s) )) \
    && (( ld++ ))

    # Check intermediate years for leap day
    for (( y=(( ${frstYYYY}+1 )); y<=(( ${lastYYYY}-1 )); y++ ))
    do
        $(date -ud "$y-02-29" > /dev/null 2>&1) && (( ld++ ))
    done

    # Check last year (if different from first year) for leap day between start
    # and end date
    (( $lastYYYY > $frstYYYY )) \
    && $(date -ud "${lastYYYY}-02-29" > /dev/null 2>&1) \
    && (( $(date -ud "$1" +%s) < $(date -ud "${frstYYYY}-03-01" +%s) )) \
    && (( $(date -ud "$2" +%s) > $(date -ud "${lastYYYY}-02-28" +%s) )) \
    && (( ld++ ))

    set -e

    echo "$ld"
}


[[ $@ == *verbose* ]] && set -x

module load ${module_list:?}
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:+${LD_LIBRARY_PATH}:}"${extralibs_list}"

mkdir -p ${run_dir}
cd ${run_dir}


run_start_epoch=$(date -u -d"${run_start_date}" +%s)
run_start_date_cos="${run_start_date//-}00"
run_start_date=$(date -uR -d "${run_start_date}")
run_end_date="${run_start_date} + ${run_duration:?}"
run_end_date=$(date -uR -d "${run_end_date}")
run_end_epoch=$(date -u -d"${run_end_date}" +%s)
run_end_date_cos="$(date -u -d"${run_end_date}" +%Y%m%d%H)"


if [[ ${special_rs_nemo} == "true" && ${special_restart} == "true" ]]; then
    echo "ERROR: both special_rs_nemo and special_restart are true. only can be true (at most). exit".
    exit 2
fi


for (( ; run_num_legs>0 ; run_num_legs-- ))
do

    # Check for restart information file and set the current leg start date
    #   Ignore restart information file if force_run_from_scratch is true
    if ! [ -r ${info_file} ]
    then
        leg_is_restart=false
        leg_start_date=${run_start_date}
        leg_start_date_cos=${run_start_date_cos}
        leg_number=1
        echo "leg_is_restart=false"
    else
        leg_is_restart=true
        . ./${info_file}
        leg_start_date=${leg_end_date}
        leg_start_date_cos="$(date -u -d"${leg_start_date}" +%Y%m%d%H)"
        leg_number=$((leg_number+1))
        echo "leg_is_restart=true"
    fi

    # Compute the end date of the current leg
    if [ -n "${rst_freq}" ]
    then
        if (($(date -u -d "${leg_start_date:?} + ${rst_freq}" +%s) > $(date -u -d "${leg_start_date:?} + ${run_duration}" +%s) )); then
           rst_freq=$run_duration
        fi
        leg_end_date=$(date -uR -d "${leg_start_date} + ${rst_freq}")
        leg_end_date_cos="$(date -u -d"${leg_end_date}" +%Y%m%d%H)"
    else
        leg_end_date=${run_end_date}
        leg_end_date_cos=${run_end_date_cos}
    fi

    if [ $(date -u -d "${leg_end_date}" +%s) -gt $(date -u -d "${run_end_date}" +%s) ]
    then
        leg_end_date=${run_end_date}
        leg_end_date_cos=${run_end_date_cos}
    fi

    # Some time variables needed later
    leg_start_epoch=$(date -u -d "${leg_start_date}" +%s)
    leg_end_epoch=$(date -u -d "${leg_start_date:?} + ${rst_freq:=$run_duration}" +%s)
    leg_length_sec=$(( $(date -u -d "${leg_end_date}" +%s) - $(date -u -d "${leg_start_date}" +%s) ))
    leg_start_sec=$(( $(date -u -d "${leg_start_date}" +%s) - $(date -u -d "${run_start_date}" +%s) ))
    leg_end_sec=$(( $(date -u -d "${leg_end_date}" +%s) - $(date -u -d "${run_start_date}" +%s) ))
    leg_start_date_yyyymmdd=$(date -u -d "${leg_start_date}" +%Y%m%d)
    leg_start_date_yyyy=$(date -u -d "${leg_start_date}" +%Y)
    leg_end_date_yyyy=$(date -u -d "${leg_end_date}" +%Y)

    echo "leg_number: $leg_number"
    echo "leg_start_date: $leg_start_date"
    echo "leg_end_date: $leg_end_date"

    # Check whether there's actually time left to simulate - exit otherwise
    if [ ${leg_length_sec} -le 0 ]
    then
        info "Leg start date equal to or after end of simulation."
        info "Nothing left to do. Exiting."
        exit 0
    fi

    # -------------------------------------------------------------------------
    # *** Prepare the run directory for a run from scratch
    # -------------------------------------------------------------------------
    if ! $leg_is_restart
    then
        # ---------------------------------------------------------------------
        # *** Check if run dir is empty. If not, and if we are allowed to do so
        #     by ${force_run_from_scratch}, remove everything
        # ---------------------------------------------------------------------
        echo "pwd: $PWD"
        if $(ls * >& /dev/null)
        then
            if ${force_run_from_scratch}
            then
                rm -fr ${run_dir}/*
            else
                echo "error: run directory not empty and \$force_run_from_scratch not set."
                exit 0
            fi
        fi

        # ---------------------------------------------------------------------
        # *** Creates some directories
        # ---------------------------------------------------------------------
        mkdir -p ${out_dir}
	
        # ---------------------------------------------------------------------
        # *** Copy executables and coupling-related files
        # ---------------------------------------------------------------------
        cd ${start_dir}
        cp -u xios_config/*xml "${run_dir}"
        cp -u namelists/* "${run_dir}"
        cd ${run_dir}
        cp -u "${nem_exe_file:?}" .
        cp -u "${xio_exe_file:?}" .

        # ---------------------------------------------------------------------
        # *** Link ic, forcing and shared files for NEMO
        # ---------------------------------------------------------------------
        [[ ! -f EMPave_old.dat ]] && echo "                               0  0.0000000000000000E+00  0.0000000000000000E+00" > EMPave_old.dat
        for file in "${ic_files[@]}"; do
            [[ ! -e ${file#*> } ]] && ln -sf $(sed 's/ *=> */ /' <<< "${ini_data_dir}/${ic_subdir}/${nem_grid}/$file")
        done
        for file in "${forcing_files[@]}"; do
            [[ ! -e ${file#*> } ||  "$file" == \** ]] && ln -sf $(sed 's/ *=> */ /' <<< "${ini_data_dir}/${forcing_subdir}/${nem_forcing_set}/$file")
        done
        for file in "${shared_files[@]}"; do
            [[ ! -e ${file#*> } ]] && ln -sf $(sed 's/ *=> */ /' <<< "${shared_dir}/$file")
        done

        # ---------------------------------------------------------------------
        # *** Copy potential special restarts 
        # ---------------------------------------------------------------------
        if $special_restart
        then

            ### Getting the right output folders 
            leg_start_date_tmp=$leg_start_date
            leg_end_date_tmp=$leg_end_date
            leg_number_tmp=$leg_number
            rsync -avz  ${run_dir}/../../${special_restart_from}/rundir/${info_file} ${run_dir}
            special_date=$(date -uR -d "${special_restart_date}")
            sed -i "/$special_date/q" ${run_dir}/${info_file}
            . ${run_dir}/${info_file}
            special_restart_leg=$(printf %03d $((leg_number)))

            echo "special_date: $special_date"
            echo "special_restart_leg: $special_restart_leg"
            echo "special_restart_from: $special_restart_from"

            ### NEMO
            indir="${run_dir}/../../${special_restart_from}/archive/restart/$(printf %03d $((special_restart_leg)))"
            cd ${indir}
            for f in *.nc; do
                nf=`echo "${f/$special_restart_from/$exp_name}"`
                echo "copy $f to ${run_dir}/$nf" 
               cp $f ${run_dir}/$nf
            done
            cd -
            cd ${run_dir}
            echo "pwd: $PWD"
            for f in ${exp_name}_????????_restart_???_????.nc; do
                nf=`echo ${f: -19}`
                echo "linking $f to ${run_dir}/$nf" 
                ln -s $f $nf
            done

            cd ${run_dir}
            rm -rf ${info_file}
            leg_start_date=$leg_start_date_tmp
            leg_end_date=$leg_end_date_tmp
            leg_number=$leg_number_tmp

	elif ${special_rs_nemo}
	then

	    dir_rs=`dirname ${special_rs_nemo_stencil}`
	    base_rs=`basename ${special_rs_nemo_stencil}`

	    if ${special_rs_nemo_rebuilt}
	    then

		if [[ -f ${special_rs_nemo_stencil}oce.nc && -f ${special_rs_nemo_stencil}ice.nc ]]
		then
		    ln -fs ${special_rs_nemo_stencil}oce.nc restart_oce.nc
		    ln -fs ${special_rs_nemo_stencil}ice.nc restart_ice.nc
		else
		    echo "EXITING. ${special_rs_nemo_stencil}oce.nc or ${special_rs_nemo_stencil}ice.nc not found"
		    exit 2
		fi

	    else
	        for (( n=0 ; n<nem_numproc ; n++ ))
		do
		    np=$(printf %04d ${n})
		    if [[ -f ${special_rs_nemo_stencil}oce_${np}.nc && -f ${special_rs_nemo_stencil}ice_${np}.nc ]]
		    then
			ln -fs ${special_rs_nemo_stencil}oce_${np}.nc restart_oce_${np}.nc
			ln -fs ${special_rs_nemo_stencil}ice_${np}.nc restart_ice_${np}.nc
		    else
			echo "EXITING. ${special_rs_nemo_stencil}oce_${np}.nc or ${special_rs_nemo_stencil}ice_${np}.nc not found"
			exit 2
		    fi
		done
	    fi

        fi # if $special_restart

    else # if ! $leg_is_restart

        # ---------------------------------------------------------------------
        # *** Creates some directories
        # ---------------------------------------------------------------------
        mkdir -p ${out_dir}

        # ---------------------------------------------------------------------
        # *** Copy restart files in run directory
        # ---------------------------------------------------------------------
        cd ${run_dir}
        ### NEMO
        outdir="${archive_dir}/restart/$(printf %03d $((leg_number-1)))"

	
        ns=$(printf %08d $(( leg_start_sec / nem_time_step_sec )))
        for f in oce ice
        do
	    echo "${outdir}/${exp_name}_${ns}_restart_${f}_????.nc"
	    ln -sf ${outdir}/${exp_name}_${ns}_restart_${f}_????.nc ${run_dir}
        done

        for (( n=0 ; n<nem_numproc ; n++ ))
        do
	    np=$(printf %04d ${n})
	    ln -fs ${exp_name}_${ns}_restart_oce_${np}.nc restart_oce_${np}.nc
	    ln -fs ${exp_name}_${ns}_restart_ice_${np}.nc restart_ice_${np}.nc
        done

    fi # ! $leg_is_restart

    # -------------------------------------------------------------------------
    # *** Update ocean namelist
    # -------------------------------------------------------------------------
    
    source build_namelist_cfg.sh > namelist_cfg


    # -------------------------------------------------------------------------
    # *** Launch experiment
    # -------------------------------------------------------------------------

    echo "START of mpirun"
    time_begin=$(date +%s)
    mpirun -np "${xio_numproc:?}" "${xio_exe_file:?}" : -np "${nem_numproc:?}" "${nem_exe_file:?}"
    time_end=$(date +%s)
    echo "END of: mpirun"

    # -------------------------------------------------------------------------
    # *** Move output files to archive directory
    # -------------------------------------------------------------------------
    ### NEMO
    outdir="${archive_dir}/output/$(printf %03d $((leg_number)))"
    mkdir -p ${outdir}

    for v in grid_U grid_V grid_W grid_T icemod SBC scalar SBC_scalar diad_T \
             grid_T_2D grid_U_2D grid_V_2D grid_W_2D grid_T_3D grid_U_3D grid_V_3D grid_W_3D \
             grid_1point grid_T_3D_ncatice vert_sum \
             grid_ptr_W_3basin_3D grid_ptr_T_3basin_2D grid_ptr_T_2D \
             zoom_700_sum zoom_300_sum zoom_2000_sum
    do
        for f in ${exp_name}_*_????????_????????_*${v}*.nc
        do
            test -f $f && mv $f $outdir/
        done
    done

    
    # -------------------------------------------------------------------------
    # *** Move restart files to archive directory
    # -------------------------------------------------------------------------
    ### NEMO
    outdir="${archive_dir}/restart/$(printf %03d $((leg_number)))"
    mkdir -p ${outdir}

    ns=$(printf %08d $(( leg_end_sec / nem_time_step_sec )))
    for f in oce ice
    do
        mv ${exp_name}_${ns}_restart_${f}_????.nc ${outdir}
    done

    # -------------------------------------------------------------------------
    # *** Move log files to archive directory
    # -------------------------------------------------------------------------
    outdir="${archive_dir}/log/$(printf %03d $((leg_number)))"
    mkdir -p ${outdir}

    for f in \
        ocean.output time.step solver.stat
    do
        test -f ${f} && mv ${f} ${outdir}
    done

    # ---------------------------------------------------------------------
    # *** Remove all leftover output files from previous legs
    # ---------------------------------------------------------------------

    ### NEMO files
    rm -f ${exp_name}_??_????????_????????_{grid_U,grid_V,grid_W,grid_T,icemod,SBC,scalar,SBC_scalar}.nc
    rm -f ${exp_name}_*restart_*nc
    ns=$(printf %08d $(( leg_start_sec / nem_time_step_sec )))
    for (( n=0 ; n<nem_numproc ; n++ ))
    do
        np=$(printf %04d ${n})
        rm -f restart_oce_${np}.nc
        rm -f restart_ice_${np}.nc
    done

    # -------------------------------------------------------------------------
    # *** Write the restart control file
    # -------------------------------------------------------------------------
    shopt -u nullglob

    tr=$(date -d "0 -$time_begin sec + $time_end sec" +%T)
    current_date=$(date +'%F %T')
    {
      echo "#"
      echo "# Finished leg at ${current_date} after ${tr} (hh:mm:ss)" 
      echo "leg_number=${leg_number}"
      echo "leg_start_date=\"${leg_start_date}\""
      echo "leg_end_date=\"${leg_end_date}\""
    } | tee -a "${info_file}"

    # Need to reset force_run_from_scratch in order to avoid destroying the next leg
    force_run_from_scratch=false
    special_restart=false

done # Loop over legs

cd ${start_dir}

cd - >/dev/null
[[ $@ == *noresubmit* ]] && exit 0
if (( leg_end_epoch < run_end_epoch )) ; then
    echo "Leg end earlier than end of simulation."
    echo "Submitting another job."
    if [[ "$@" == *"local"* ]] ; then
        exec "$0" "$@"
    elif [ "qsub" == *sbatch* ] ; then
        sbatch $0 $@ | awk '{print $4}' >> ${run_dir}/.coral_jobs
    elif [[ $USER == "vsc"* ]]; then
        if [ -f ${PBS_O_WORKDIR-$PWD}/tmp/credits_file ]; then
            credits=`cat ${PBS_O_WORKDIR-$PWD}/tmp/credits_file`
        else
            credits_tmp=`grep "account string:" ${PBS_O_WORKDIR-$PWD}/*.o* -R | sed 's/^.*: //'`
            credits=`echo $credits_tmp | awk '{ print $1 }'`

            cd ${start_dir}
            mkdir -p tmp
            echo ${credits} > tmp/credits_file
        fi
        if [ -f ${PBS_O_WORKDIR-$PWD}/tmp/script_name ]; then
            script_name=`cat ${PBS_O_WORKDIR-$PWD}/tmp/script_name`
        else
            script_name="run.sh"
        fi
        ssh login2 "cd $start_dir; qsub -A $credits $script_name | tee -a coral_jobs"
        sleep 2
        jobid=`cat coral_jobs`
        rm -f coral_jobs
        jobid=${jobid%%.*}
        echo "${jobid}" >> "${run_dir}"/.coral_jobs ;
    else
        "qsub" -v PBS_OPTIONS="$@" "$0" | tee -a coral_jobs
        sleep 2
        jobid=`cat coral_jobs`
        rm -f coral_jobs
        jobid=${jobid%.*}
        echo "${jobid}" >> "${run_dir}"/.coral_jobs
    fi
else
    if [[ $USER == "vsc"* ]]; then
        rm -rf ${start_dir}/tmp
    fi

    rm -rf ${run_dir}/../cclm_out_tmp
    echo "Nothing left to do. Exiting." 

cat << "EOF" 
   _____ _                 _       _   _                                         _      _           _ 
  / ____(_)               | |     | | (_)                                       | |    | |         | |
 | (___  _ _ __ ___  _   _| | __ _| |_ _  ___  _ __     ___ ___  _ __ ___  _ __ | | ___| |_ ___  __| |
  \___ \| | '_ ` _ \| | | | |/ _` | __| |/ _ \| '_ \   / __/ _ \| '_ ` _ \| '_ \| |/ _ \ __/ _ \/ _` |
  ____) | | | | | | | |_| | | (_| | |_| | (_) | | | | | (_| (_) | | | | | | |_) | |  __/ ||  __/ (_| |
 |_____/|_|_| |_| |_|\__,_|_|\__,_|\__|_|\___/|_| |_|  \___\___/|_| |_| |_| .__/|_|\___|\__\___|\__,_|
                                                                          | |                         
                                                                          |_|                         
EOF

fi

exit 0