#!/bin/bash 

# 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"
}

set -ueo pipefail


nemo_src_dir=${HOME}/modeles/NEMO/3.6
ini_data_dir=/scratch/pbarriat/DATA/nemo
nem_grid=ORCA2L31
# Simulation start and end date. Use any (reasonable) syntax you want.
run_start_date="1990-01-01"
run_duration="4 years"

# Restart frequency. Use any (reasonable) number and time unit you want.
# For runs without restart, leave this variable empty or set to 0
rst_freq="1 year"

# Number of restart legs to be run in one go
run_num_legs=2

# Directories
#start_dir=${SLURM_SUBMIT_DIR}
start_dir=${PWD}
ctrl_file_dir=${start_dir}/ctrl

# This file is used to store information about restarts
run_dir="./rundir"
ece_info_file="ece.info"

nemo_src_dir=$HOME/modeles/NEMO/3.6
nem_config_name="ORCA2_LIM3"
nem_exe_file=${nemo_src_dir}/CONFIG/${nem_config_name}/BLD/bin/nemo.exe
xio_exe_file=${nemo_src_dir}/EXTERNAL/xios-1.0/bin/xios_server.exe

# Make sure run directory exists with the necessary stuff and move there
mkdir -p $run_dir && cd $_
cp ${nem_exe_file} .
cp ${xio_exe_file} .
# Configure paths for NEMO IC /!\ filenames cannot contain white spaces
ic_files=("coordinates.nc" 
          "bathy_meter.nc" 
          "runoff_core_monthly.nc runoffs_depth"
          "ahmcoef.nc" 
          "mask_itf.nc" 
          "K1rowdrg.nc" 
          "M2rowdrg.nc" 
          "data_1m_potential_temperature_nomask.nc temperature.nc"
          "data_1m_salinity_nomask.nc salinity" 
          "subbasins.nc" 
          "resto.nc" 
          "chlorophyll.nc" 
          "geothermal_heating.nc")
for file in "${ic_files[@]}"
do 
ln -sf ${ini_data_dir}/INIT/${nem_grid}/$file
done


# Normalize date formats
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}")

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

# Initialize variables  using restart file if it exists
[[ -r ${ece_info_file} ]] && source ${ece_info_file}
leg_start_date=${leg_end_date:-$run_start_date}
leg_number=$((${leg_number:=0}+1))

# Compute leg end-date and trim if necessary
leg_end_date=$(date -uR -d "${leg_start_date} + ${rst_freq:=$run_duration}")
(( $(date -d "${leg_end_date}" +%s) > $(date -d "${run_end_date}" +%s) )) && leg_end_date=${run_end_date}

# Check whether there is some work left to do
if (( $(date -d "${leg_start_date}" +%s) >= $(date -d "${run_end_date}" +%s) )) ;
then
	echo "Leg start date equal to or after end of simulation."
	echo "Nothing left to do. Exiting."
	exit 0
fi

    # Some time variables needed later
    leg_length_sec=$(( $(date -d "${leg_end_date}" +%s) - $(date -d "${leg_start_date}" +%s) ))
    leg_start_sec=$(( $(date -d "${leg_start_date}" +%s) - $(date -d "${run_start_date}" +%s) ))
    leg_end_sec=$(( $(date -d "${leg_end_date}" +%s) - $(date -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)

    # Correct for leap days because NEMO standalone uses no-leap calendar
    leg_length_sec=$(( leg_length_sec - $(leap_days "${leg_start_date}" "${leg_end_date}")*24*3600 ))
    leg_start_sec=$(( leg_start_sec - $(leap_days "${run_start_date}" "${leg_start_date}")*24*3600 ))
    leg_end_sec=$(( leg_end_sec - $(leap_days "${run_start_date}" "${leg_end_date}")*24*3600 ))

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


    t1=$(date +%s)
	sleep 1
    t2=$(date +%s)

    tr=$(date -d "0 -$t1 sec + $t2 sec" +%T)

    # -------------------------------------------------------------------------
    # *** Write the restart control file
    # -------------------------------------------------------------------------

    echo "#"                                             | tee -a ${ece_info_file}
    echo "# Finished leg at `date '+%F %T'` after ${tr} (hh:mm:ss)" \
                                                         | tee -a ${ece_info_file}
    echo "leg_number=${leg_number}"                      | tee -a ${ece_info_file}
    echo "leg_start_date=\"${leg_start_date}\""          | tee -a ${ece_info_file}
    echo "leg_end_date=\"${leg_end_date}\""              | tee -a ${ece_info_file}

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

done # loop over legs

exit 0