#!/bin/bash

# Submission script for Lemaitre2
#SBATCH --job-name=NEMO2r0
#SBATCH --time=10:00:00
#
#SBATCH --ntasks=65
#SBATCH --mem-per-cpu=3072
#SBATCH --partition=Def
#
#SBATCH --mail-user=pierre-yves.barriat@uclouvain.be
#SBATCH --mail-type=ALL

set -ue

# =============================================================================
# *** BEGIN User configuration
# =============================================================================

script_name="run.sh"

# -----------------------------------------------------------------------------
# *** General configuration
# -----------------------------------------------------------------------------

# Component configuration
# (for syntax of the $config variable, see librunscript.sh)
config="nemo lim3 xios:detached"

# Experiment name (exactly 4 letters!)
exp_name=EXP9

# Simulation start and end date. Use any (reasonable) syntax you want.
run_start_date="1990-01-01"
run_end_date="${run_start_date} + 4 years"

# Set $force_run_from_scratch to 'true' if you want to force this run to start
# from scratch, possibly ignoring any restart files present in the run
# directory. Leave set to 'false' otherwise.
# NOTE: If set to 'true' the run directory $run_dir is cleaned!
force_run_from_scratch=false

# Resolution
#nem_grid=ORCA1L75
nem_grid=ORCA2L31

# Restart frequency. Use any (reasonable) number and time unit you want.
# For runs without restart, leave this variable empty
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
ece_info_file="ece.info"

# -----------------------------------------------------------------------------
# *** Read platform dependent configuration
# -----------------------------------------------------------------------------

# librunscript defines some helper functions
source ${start_dir}/librunscript.sh

source ${start_dir}/nemo.cfg

configure

# -----------------------------------------------------------------------------
# *** NEMO/LIM configuration
# -----------------------------------------------------------------------------

nem_time_step_sec=2700
lim_time_step_sec=2700

nem_restart_offset=0

nem_res_hor=$(echo ${nem_grid} | sed 's:ORCA\([0-9]\+\)L[0-9]\+:\1:')

nem_forcing_set=fill

nem_config_name=ORCA2_LIM3

nem_exe_file=${nemo_src_dir}/CONFIG/${nem_config_name}/BLD/bin/nemo.exe

nem_numproc=64

# -----------------------------------------------------------------------------
# *** XIOS configuration
# -----------------------------------------------------------------------------

xio_exe_file=${nemo_src_dir}/EXTERNAL/xios-1.0/bin/xios_server.exe

xio_numproc=1

# =============================================================================
# *** END of User configuration
# =============================================================================

# =============================================================================
# *** This is where the code begins ...
# =============================================================================

# -----------------------------------------------------------------------------
# *** Make sure to clean up on exit
# -----------------------------------------------------------------------------

trap 'cleanup' EXIT SIGHUP SIGINT SIGTERM

# -----------------------------------------------------------------------------
# *** Create the run dir if necessary and go there
#     Everything is done from here.
# -----------------------------------------------------------------------------
if [ ! -d ${run_dir} ]
then
    mkdir -p ${run_dir}
    force_run_from_scratch=true
else
    force_run_from_scratch=false
    if ! [ -r ${run_dir}/${ece_info_file} ]
    then
        info "Problem: ${ece_info_file} doesn't exist. Exiting."
        exit 0
    fi
fi
cd ${run_dir}

# -----------------------------------------------------------------------------
# *** Determine the time span of this run and whether it's a restart leg
# -----------------------------------------------------------------------------

# Regularise the format of the start and end date of the simulation
run_start_date=$(date -uR -d "${run_start_date}")
run_end_date=$(date -uR -d "${run_end_date}")

# Loop over the number of legs
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 ${force_run_from_scratch} || ! [ -r ${ece_info_file} ]
    then
        leg_is_restart=false
        leg_start_date=${run_start_date}
        leg_number=1
    else
        leg_is_restart=true
        . ./${ece_info_file}
        leg_start_date=${leg_end_date}
        leg_number=$((leg_number+1))
    fi

    # Compute the end date of the current leg
    if [ -n "${rst_freq}" ]
    then
        leg_end_date=$(date -uR -d "${leg_start_date} + ${rst_freq}")
    else
        leg_end_date=${run_end_date}
    fi

    if [ $(date -d "${leg_end_date}" +%s) -gt $(date -d "${run_end_date}" +%s) ]
    then
        leg_end_date=${run_end_date}
    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 ))

    # 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
        # ---------------------------------------------------------------------
        if $(ls * >& /dev/null)
        then
            if ${force_run_from_scratch}
            then
                rm -fr ${run_dir}/*
            else
                error "Run directory not empty and \$force_run_from_scratch not set."
            fi
        fi

        # ---------------------------------------------------------------------
        # *** Copy executables of model components
        # *** Additionally, create symlinks to the original place for reference
        # ---------------------------------------------------------------------
        cp    ${nem_exe_file} .
        ln -s ${nem_exe_file} $(basename ${nem_exe_file}).lnk

        cp    ${xio_exe_file} .
        ln -s ${xio_exe_file} $(basename ${xio_exe_file}).lnk

        # ---------------------------------------------------------------------
        # *** Files needed for NEMO (linked)
        # ---------------------------------------------------------------------

        # Various initialisation and forcing files
        configure_opa 

        # Write fake file for previous fresh water budget adjustment (nn_fwb==2 in namelist)
        echo "                               0  0.0000000000000000E+00  0.0000000000000000E+00" > EMPave_old.dat

        # XIOS files
        configure_xios

    else # i.e. $leg_is_restart == true

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

        # NEMO output files
        rm -f ${exp_name}_??_????????_????????_{grid_U,grid_V,grid_W,grid_T,icemod,SBC,SBC_scalar,scalar}.nc
        rm -f ${exp_name}_??_{grid_U,grid_V,grid_W,grid_T,icemod,SBC,SBC_scalar,scalar}.nc

    fi # ! $leg_is_restart

    # -------------------------------------------------------------------------
    # *** Create some control files
    # -------------------------------------------------------------------------

    # NEMO and LIM namelists 
    cp ${ctrl_file_dir}/namelist_ref                 .
    cp ${ctrl_file_dir}/namelist_ice_lim3_ref        ./namelist_ice_ref
    cp ${ctrl_file_dir}/namelist_ice_lim3_cfg        ./namelist_ice_cfg
    . ${ctrl_file_dir}/namelist_${nem_grid}_cfg.sh   > namelist_cfg

    # -------------------------------------------------------------------------
    # *** Link the appropriate NEMO restart files of the previous leg
    # -------------------------------------------------------------------------
    if $leg_is_restart
    then
        ns=$(printf %08d $(( leg_start_sec / nem_time_step_sec - nem_restart_offset )))
        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

    # -------------------------------------------------------------------------
    # *** Start the run
    # -------------------------------------------------------------------------

    # Use the launch function from the platform configuration file
    t1=$(date +%s)
    launch \
        ${xio_numproc} ${xio_exe_file} -- \
        ${nem_numproc} ${nem_exe_file}
    t2=$(date +%s)

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

    # -------------------------------------------------------------------------
    # *** Check for signs of success
    #     Note the tests provide no guarantee that things went fine! They are
    #     just based on the IFS and NEMO log files. More tests (e.g. checking
    #     restart files) could be implemented.
    # -------------------------------------------------------------------------
    # Check for NEMO success
    if [ -f ocean.output ]
    then
        if [ "$(awk '/New day/{d=$10}END{print d}' ocean.output)" == "$(date -d "${leg_end_date} - 1 day" +%Y/%m/%d)" ]
        then
            info "Leg successfully completed according to NEMO log file 'ocean.output'."
        else
            error "Leg not completed according to NEMO log file 'ocean.output'."
        fi

    else
        error "NEMO log file 'ocean.output' not found after run."
    fi

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

    for v in grid_U grid_V grid_W grid_T icemod SBC SBC_scalar scalar
    do
        for f in ${exp_name}_??_????????_????????_${v}.nc
        do
            test -f $f && mv $f $outdir/
        done
        for f in ${exp_name}_??_${v}.nc
        do
            test -f $f && mv $f $outdir/
        done
    done

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

        ls .
        ns=$(printf %08d $(( leg_start_sec / nem_time_step_sec - nem_restart_offset )))
        for f in oce ice
        do
            mv ${exp_name}_${ns}_restart_${f}_????.nc ${outdir}
        done
        echo "PEDRO"
        ls .
    fi

    # -------------------------------------------------------------------------
    # *** Move log files to archive directory
    # -------------------------------------------------------------------------
    outdir="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

    # -------------------------------------------------------------------------
    # *** 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
    force_run_from_scratch=false

done # loop over legs

# -----------------------------------------------------------------------------
# *** Platform dependent finalising of the run
# -----------------------------------------------------------------------------
finalise

exit 0