#================================================================================
# save_ic functions - these are used to save EC-Earth states during a run,
# which can be used as Initial Conditions (ICs) for other runs.
#
# save_ic_get_config : define the config to use based on the save_ic config parameter
#
# save_ic has two options available:
#   save_ic:end_leg (save ICs at the end of every leg)
#   save_ic:end_run (save ICs at the end of the run)
#
# By default, without an option, ICs are saved at given offsets from the leg
# start, possibly if conditions are met. These offset(s) and condition(s) must
# be coded in the "else" clause below. See commented example.
#
function save_ic_get_config ()
{
    do_save_ic=false
    save_ic_custom=false
    if $(has_config save_ic)
    then
        if $(has_config save_ic:end_leg)
        then
            # save ICs at the end of every leg
            do_save_ic=true
            save_ic_date_offset=( "+ ${leg_length_sec} sec" )
        elif $(has_config save_ic:end_run)
        then
            # save ICs at the end of the run
            leg_end_date_yyyymmdd=$(date -u -d "${leg_end_date}" +%Y%m%d)
            run_end_date_yyyymmdd=$(date -u -d "${run_end_date}" +%Y%m%d)
            if [ ${leg_end_date_yyyymmdd} -eq ${run_end_date_yyyymmdd} ]
            then
                do_save_ic=true && save_ic_date_offset=( "+ ${leg_length_sec} sec" )
            fi
        else
            echo "WARNING: save_ic requires to define condition and offset here or in the runscript"
            save_ic_custom=true
            do_save_ic=false
            save_ic_date_offset=( )
            # edit here or in the runscript the conditional statement to decide which legs to save ICs and the offset from leg start
            # e.g. to save at the end of every 10 legs:
            #EXAMPLE if (( leg_number%10 == 0 ))
            #EXAMPLE then
            #EXAMPLE     do_save_ic=true
            #EXAMPLE     save_ic_date_offset=( "+ ${leg_length_sec} sec" )
            #EXAMPLE fi
        fi
    fi
}

# save_ic_define_vars : define variables used in the save_ic functions
#
# The following variables must be defined in the runscript before calling this function
#
# 1) do_save_ic : set to true/false to enable save_ic feature
#
# 2) save_ic_date_offset : array of dates (offset from start of chunk, at 00:00h) to save ICs.
#    e.g. save_ic_date_offset=( "+1 month" "+2 month" )
#    The following restrictions apply, it is up to the user to ensure compliance
#     - first timestep of run is not supported
#     - if requesting first timestep of leg start, NEMO restarts from previous legs will
#       be used if present
#     - only one IC per month is supported (to simplify the IFS filter)
#     - maximum 9 ICs per leg (because of current restrictions in NEMO namelist)
#     - if requesting model-level output for IFS (e.g. for PRIMAVERA),
#       there will be no model-level variables in output on the timestep which
#       ICs are requested, unless requesting the last timestep of the leg
#
function save_ic_define_vars ()
{
    if true
    then
        # save original value of ifs_lastout
        has_config ifs && [[ -z "${save_ic_ifs_lastout_orig:-}" ]] && save_ic_ifs_lastout_orig=${ifs_lastout}
        save_ic_ifs_lastout=false

        # sanity checks
        [ ${#save_ic_date_offset[@]} -gt 9 ] && error "Cannot create more than 9 ICs in a given leg"

        # fill vars
        for (( i=0; i< ${#save_ic_date_offset[@]}; i++ ))
        do
            #save_ic_date1[$i]=${save_ic_date[$i]}
            save_ic_date1[$i]=$(date -u -d "${leg_start_date} ${save_ic_date_offset[$i]}" +%Y%m%d)
            save_ic_date[$i]=$(date -uR -d "${save_ic_date1[$i]}")
            save_ic_sec[$i]=$(( $(date -d "${save_ic_date[$i]}" +%s) - $(date -d "${run_start_date}" +%s) ))
            # Correct for leap days because NEMO standalone uses no-leap calendar
            has_config nemo && ! has_config ifs && save_ic_sec[$i]=$(( save_ic_sec[$i] - $(leap_days "${run_start_date}" "${save_ic_date[$i]}")*24*3600 ))
            save_ic_day[$i]=$(( ${save_ic_sec[$i]} / 86400 ))
            has_config ifs && save_ic_ppt_files[$i]="ppt"$(printf %06d ${save_ic_day[$i]})"0000"
            has_config nemo && save_ic_nemo_ts[$i]=$(( save_ic_sec[$i] / nem_time_step_sec ))
            # sanity checks
            if [ ${save_ic_sec[$i]} -eq 0 ]
            then
                info "Initial conditions cannot be created for first timestep of run"
                has_config ifs && save_ic_ppt_files[$i]=""
                has_config nemo && save_ic_nemo_ts[$i]=-1
            elif [ ${save_ic_sec[$i]} -eq $leg_start_sec ]
            then
                if has_config nemo
                then
                    info "Cannot create NEMO ICs for first timestep of leg at save_ic_date1[$i], will use restarts from previous leg"
                    save_ic_nemo_ts[$i]=-1
                fi
            fi
            # if saving ICs at the end of the leg, set ifs_lastout=true
            if [ ${save_ic_sec[$i]} -eq $leg_end_sec ]
            then
                if $(has_config ifs) && ! ${save_ic_ifs_lastout_orig}
                then
                    ifs_lastout=true
                    save_ic_ifs_lastout=true
                fi
            fi
        done
    fi
}

# prepare output configuration (NEMO namelist, IFS ppt file)
function save_ic_prepare_output()
{
    # For NEMO this is handled in the namelist variables ln_rst_list and nn_stocklist in namelist.nemo.ref.sh
    if has_config nemo
    then
        nemo_stocklist=""
        #for (( i=0; i< ${#save_ic_date[@]}; i++ )) ; do nemo_stocklist+=$(( save_ic_sec[$i] / nem_time_step_sec )), ; done
        for (( i=0; i< ${#save_ic_date[@]}; i++ ))
        do
            [ ${save_ic_nemo_ts[$i]} -ne -1 ] && nemo_stocklist+=${save_ic_nemo_ts[$i]},
        done
        nemo_stocklist+=$(( leg_end_sec / nem_time_step_sec )),
        for (( i=${#save_ic_date[@]}; i<9 ; i++ )) ; do nemo_stocklist+=0, ; done
        nemo_stocklist+=$'\n'
    fi

    # For IFS we need to create links to special ppt file
    if has_config ifs
    then
        # create the special ppt file
        save_ic_ppt_filename=pptdddddd0000_save_ic
        save_ic_ppt_ifile=pptdddddd0000 # change this for e.g. primavera
        # these are the codes which must be added to the special ppt file used for generating
        # required IFS model-level output at each requested timestep
        save_ic_MFP2DF_param_ids="129,152"
        save_ic_MFPPHY_param_ids="139,170,183,236,39,40,41,42,141,198,235,31,238,32,33,34,35,36,37,38,148,172,173,174,15,17,16,18,66,67,74,43,160,161,162,163,234,28,27,30,29" #all variables needed
        #save_ic_MFPPHY_param_ids="198,36,37,38,148,172,173,174,15,17,16,18,66,67,74,43,160,161,162,163,234,28,27,30,29" #all variables (old)
        #save_ic_MFPPHY_param_ids="198,36,37,38,148,172,173,174,15,17,16,18,74,43,160,161,162,163,234" #trunk
        #save_ic_MFPPHY_param_ids="173,174,15,17,16,18,66,67,74,43,160,161,162,163,234,28,27,30,29" #primavera
        save_ic_filter_param_ids=""
        # get requested output codes as a single line without space and without the last comma
        NFPPHY=`grep NFPPHY postins/$save_ic_ppt_ifile | sed -e 's/^\s*NFPPHY\s*=\s*//' -e 's/[,]$//'`
        MFPPHY=$(sed -e '1h;2,$H;$!d;g' -re 's/.*MFPPHY *=([0-9,\. \n]+).*/\1/' -e 's/[ \n]//g' -e 's/,$//' <postins/$save_ic_ppt_ifile)
        [[ $MFPPHY == "" ]] && error "IFS ppt file postins/$save_ic_ppt_ifile does not contain any entries in MFPPHY."

        NFP2DF=`grep NFP2DF postins/$save_ic_ppt_ifile | sed -e 's/^\s*NFP2DF\s*=\s*//' -e 's/[,]$//'`
        MFP2DF=$(sed -e '1h;2,$H;$!d;g' -re 's/.*MFP2DF *=([0-9,\. \n]+).*/\1/' -e 's/[ \n]//g' -e 's/,$//' <postins/$save_ic_ppt_ifile)
        [[ $MFP2DF == "" ]] && NFP2DF=0 && MFP2DF=","

        IFS=","
        old_MFPPHY_params=( ${MFPPHY} )
        old_MFP2DF_params=( ${MFP2DF} )
        new_MFPPHY_params=( ${save_ic_MFPPHY_param_ids} )
        new_MFP2DF_params=( ${save_ic_MFP2DF_param_ids} )
        unset IFS
        # update requested output codes by adding necessary variables
        for new_param_id in ${new_MFPPHY_params[@]}
        do
            new_param_id=`printf "%0*d" 3 $new_param_id`
            found=false
            for old_param_id in ${old_MFPPHY_params[@]}
            do
                old_param_id=`printf "%03g" $old_param_id`
                [ ${old_param_id} -eq ${new_param_id} ] && found=true && break
            done
            if ! $found
            then
                NFPPHY=$(( $NFPPHY + 1 ))
                MFPPHY="${MFPPHY},$new_param_id"
                save_ic_filter_param_ids="${save_ic_filter_param_ids}$new_param_id,"
            fi
        done
        for new_param_id in ${new_MFP2DF_params[@]}
        do
            new_param_id=`printf "%0*d" 3 $new_param_id`
            found=false
            for old_param_id in ${old_MFP2DF_params[@]}
            do
                old_param_id=`printf "%03g" $old_param_id`
                [ ${old_param_id} -eq ${new_param_id} ] && found=true && break
            done
            if ! $found
            then
                NFP2DF=$(( $NFP2DF + 1 ))
                MFP2DF="${MFP2DF},$new_param_id"
            fi
        done

        # create special ppt file with modified codes from original ppt file
        sed -r -e '1h;2,$H;$!d;g' \
            -e 's|(NFPPHY\s*=)[0-9 ]+|\1'${NFPPHY}'|' \
            -e 's|(MFPPHY\s*=)[0-9, \n]+\n|\1'${MFPPHY}',\n|' \
            -e 's|(NFP2DF\s*=)\s*[0-9]+|\1'${NFP2DF}'|' \
            -e 's|(MFP2DF\s*=)[0-9, \n]+\n|\1'${MFP2DF}',\n|' \
            -e 's|NFP3DFS\s*=\s*[0-9]+,||' \
            -e 's|MFP3DFS\s*=[0-9, \n]+\n|\n|' \
            -e 's|NRFP3S\s*=\s*\-?[0-9]+,||' \
            -e 's/&NAMFPC/\&NAMFPC\n    NFP3DFS = 9,\n    MFP3DFS = 133,75,76,246,247,248,138,155,130,\n    NRFP3S  = -99\n    NFP2DF = 2,\n    MFP2DF = 152,129,/' \
            <postins/$save_ic_ppt_ifile >postins/${save_ic_ppt_filename}

        # cleanup existing links from previous legs (required when save_ic_ifs_lastout=true)
        ls -ltr postins
        find -L postins/ -xtype l -samefile postins/${save_ic_ppt_filename} -exec rm {} \;
        # create links to ppt file for all requested timesteps
        for (( i=0; i< ${#save_ic_date[@]}; i++ ))
        do
            if [ ! -f postins/$save_ic_ppt_filename ]
            then
                error "IFS IC file $save_ic_ppt_filename does not exist."
            elif [ "${save_ic_ppt_files[$i]}" != "" ]  ; then
                ln -sf $save_ic_ppt_filename postins/${save_ic_ppt_files[$i]}
            fi
        done
        /bin/ls -1 postins/* > dirlist
        ls -ltr postins
    fi
}

# post-process save_ic results
function save_ic_postproc()
{
    # cleanup existing links for subsequent legs (see issue #345)
    if has_config ifs
    then
        ls -ltr postins
        find -L  postins/ -xtype l -samefile postins/${save_ic_ppt_filename} -exec rm {} \;
        /bin/ls -1 postins/* > dirlist
        ls -ltr postins
    fi
    for (( i=0; i< ${#save_ic_date[@]}; i++ ))
    do
        outdir="save_ic/"${save_ic_date1[$i]}

        # save IFS ICs
        if has_config ifs
        then
            [ "${save_ic_ppt_files[$i]}" != "" ] && save_ic_ifs_out2init $i
        fi
        # copy oasis restarts
        if has_config oasis
        then
            # copy oasis restarts if this is the last or first timestep of the leg
            # oasis files depend on the components and configuration
            # so they go in their own sub-directory, easier to separate from ifs and nemo files
            if [ ${save_ic_sec[$i]} -eq $leg_end_sec ]
            then
                mkdir -p ${outdir}/oasis
                for f in ${oas_rst_files} ; do test -f ${f} && cp ${f} ${outdir}/oasis ; done
            elif [ ${save_ic_sec[$i]} -eq $leg_start_sec ]
            then
                mkdir -p ${outdir}/oasis
                oas_rst_dir="restart/oasis/$(printf %03d $((leg_number)))"
                for f in ${oas_rst_files} ; do test -f ${oas_rst_dir}/${f} && cp ${oas_rst_dir}/${f} ${outdir}/oasis ; done
            fi
        fi
        # copy NEMO restarts
        if has_config nemo
        then
            nemo_ts=$(( save_ic_sec[$i] / nem_time_step_sec ))
            ns=$(printf %08d $(( nemo_ts - nem_restart_offset )))
	    # copy each restart type seperately, depending on config
	    extensions="oce"
	    has_config lim3 && extensions+=" ice"
	    has_config pisces && extensions+=" trc"
	    for ext in $extensions
	    do
		if ls ${exp_name}_${ns}_restart_${ext}_????.nc > /dev/null 2>&1
		then
                    mkdir -p ${outdir}/nemo
                    if [[ ${save_ic_sec[$i]} -eq $leg_start_sec ]] || [[ ${save_ic_sec[$i]} -eq $leg_end_sec ]]
                    then
			cp -f ${exp_name}_${ns}_restart_${ext}_????.nc ${outdir}/nemo
                    else
			mv -f ${exp_name}_${ns}_restart_${ext}_????.nc ${outdir}/nemo
                    fi
		else
                    info "Cannot find NEMO ($ext) initial conditions for date ${save_ic_date1[$i]}"
		fi
	    done
        fi
        # copy TM5 restarts
        leg_end_date_yyyymmdd=$(date -u -d "${leg_end_date}" +%Y%m%d)
        if has_config tm5 && [[ $leg_end_date_yyyymmdd == ${save_ic_date1[$i]} ]]
        then
	    found=false
	    for f in TM5_restart_${leg_end_date_yyyymmdd}_0000_glb300x200.nc save_${leg_end_date_yyyymmdd}00_glb300x200.hdf
	    do
		if [[ -f "$f" ]] ; then
		    mkdir -p ${outdir}/tm5
		    cp -f $f ${outdir}/tm5
		    found=true
		fi
	    done
            ! $found && info "Cannot find TM5 restarts for date ${save_ic_date1[$i]}" || true
	fi
    done
}

# Function to create IC files for IFS from special output based on a script from K. Wyser
function save_ic_ifs_out2init()
{
    icdate=${save_ic_date1[$1]}
    #srcdir=$2
    filter_output=true
    grib_filter=${GRIB_BIN_PATH}${GRIB_BIN_PATH:+/}grib_filter

    # temporary directory
    tmpdir=tmp_save_ic
    [ -d $tmpdir ] && rm -rf ${tmpdir}
    mkdir -p $tmpdir

    # first IFS timestep of a month is saved in last output file
    if [[ "$(date -d "${icdate}" +%d)" == "01" ]] ; then
        yyyymm=$(date -d "${icdate} - 6 hours" +%Y%m)
    else
        yyyymm=$(date -d "${icdate}" +%Y%m)
    fi

    # find the ICMSH/GG files containing the last timestep, which are in $srcdir/output/ifs/???
    #local ifile_sh=`find $srcdir -name ICMSH${exp_name}+$yyyymm -print -quit`
    #local ifile_gg=`find $srcdir -name ICMGG${exp_name}+$yyyymm -print -quit`
    ifile_sh=ICMSH${exp_name}+$yyyymm
    ifile_gg=ICMGG${exp_name}+$yyyymm
    if [[ ! -f "$ifile_sh" ]] ; then info "cannot find ICMSH file $ifile_sh to create initial conditions at $icdate!" ; return 0 ; fi
    if [[ ! -f "$ifile_gg" ]] ; then info "cannot find ICMGG file $ifile_gg to create initial conditions at $icdate!" ; return 0 ; fi

    if $filter_output
    then
        # define grib filters to separate data for IC from normal output
        # filter_time_init is to define which timestep contains the ICs, currently only one per month is supported
        filter_time_init="dataDate == ${icdate} && dataTime == 0"
        # filter_time_out is to filter any timesteps which might contain ICs
        filter_time_out="dataTime == 0 && ( 0 "
        for (( j=0; j< ${#save_ic_date[@]}; j++ )) ; do filter_time_out+=" || ( dataDate == ${save_ic_date1[$j]} )"; done
        filter_time_out+=" )"
        $save_ic_ifs_lastout && filter_last="( dataTime == 0 ) && ( dataDate == ${save_ic_date1[$((j-1))]} )" || filter_last="0"
        # build expression to filter the paramIDs which we have to filter out, set them in ini_filter_param_ids in the runscript
        #save_ic_filter_param_ids="198,36,37,38,148,172,173,174,15,17,16,18,74,43,160,161,162,163,234"
        filter_params="1"
        IFS=","
        for p in $save_ic_filter_param_ids ; do filter_params+=" && paramId != $p" ; done
        unset IFS

        # SH file
        ofile_sh_init=${tmpdir}/sh_init.grb
        ofile_sh_out=${tmpdir}/sh_out.grb
        ofile_sh_last=${tmpdir}/sh_last.grb
        filter_sh=${tmpdir}/filter_sh
        write_sh_init="if ( $filter_time_init ) { write \"${ofile_sh_init}\"; }"
        write_sh_out="if ( $filter_last ) { write \"${ofile_sh_last}\"; } else { write \"${ofile_sh_out}\"; }"

        cat > $filter_sh << EOT
if ( ! ( $filter_time_out ) ) { write "${ofile_sh_out}"; }
else {
if ( typeOfLevel is "hybrid" ) { ${write_sh_init}; }
else { ${write_sh_out}; }
}
EOT

        #GG file
        ofile_gg_init=${tmpdir}/gg_init.grb
        ofile_gg_out=${tmpdir}/gg_out.grb
        ofile_gg_last=${tmpdir}/gg_last.grb
        filter_gg=${tmpdir}/filter_gg
        write_gg_init="if ( $filter_time_init ) { write \"${ofile_gg_init}\"; }"
        write_gg_out="if ( $filter_last ) { write \"${ofile_gg_last}\"; } else { write \"${ofile_gg_out}\"; }"

        cat > $filter_gg << EOT
if ( ! ( $filter_time_out ) ) { write "${ofile_gg_out}"; }
else {
if ( levelType is "ml" ) { ${write_gg_init}; }
else {
  if ( levelType is "pl" ) { ${write_gg_out}; }
  else {
    if ( $filter_params ) {
      ${write_gg_init};
      ${write_gg_out};
    }
    else {
      ${write_gg_init};
    }
  }
}
}
EOT

        #run grib_filter on SH and GG files
        $grib_filter $filter_sh $ifile_sh
        $grib_filter $filter_gg $ifile_gg
        if [[ ! -f "$ofile_sh_init" ]] ; then info "ICMSH file $ifile_sh does not contain data to create initial conditions at $icdate" ; return 0 ; fi
        if [[ ! -f "$ofile_gg_init" ]] ; then info "ICMGG file $ifile_gg does not contain data to create initial conditions at $icdate" ; return 0 ; fi

        #rename files
        ifile_sh=${tmpdir}/$(basename $ifile_sh)
        ifile_gg=${tmpdir}/$(basename $ifile_gg)
        mv ${ofile_sh_init} ${ifile_sh}
        mv ${ofile_gg_init} ${ifile_gg}
        #srcdir=${tmpdir}

    fi # $filter_output

    if [[ ! -f "$ifile_sh" ]] ; then info "cannot find ICMSH file $ifile_sh to create initial conditions at $icdate!" ; return 0 ; fi
    if [[ ! -f "$ifile_gg" ]] ; then info "cannot find ICMGG file $ifile_gg to create initial conditions at $icdate!" ; return 0 ; fi

    tgtdir=save_ic/$icdate/ifs
    ofile_sh=${tgtdir}/ICMSH${exp_name}INIT
    ofile_gg_init=${tgtdir}/ICMGG${exp_name}INIT
    ofile_gg_iniua=${tgtdir}/ICMGG${exp_name}INIUA

    # make sure the output folder is created and delete any existing files
    #[ -d $tgtdir ] && rm -rf ${tgtdir}/* || mkdir -p $tgtdir
    mkdir -p $tgtdir
    rm -f $ofile_sh $ofile_gg_init $ofile_gg_iniua

    cat > ${tmpdir}/gf1 << EOT
if ( typeOfLevel is "hybrid" ) { write "${tmpdir}/shinit.[shortName].[level]"; }
EOT
    $grib_filter ${tmpdir}/gf1 $ifile_sh

    cp -f ${tmpdir}/shinit.lnsp.1 $ofile_sh
    for lev in {1..91}
    do
        for var in vo d t
        do
            cat ${tmpdir}/shinit.$var.$lev >> $ofile_sh
        done
    done
    cat ${tmpdir}/shinit.z.1 >> $ofile_sh

    cat > ${tmpdir}/gf2 << EOT
write "${tmpdir}/gginit.[shortName]";
EOT
    $grib_filter ${tmpdir}/gf2 $ifile_gg

    for var in stl1 stl2 stl3 stl4 swvl1 swvl2 swvl3 swvl4 sd src skt ci tsn asn \
        rsn sst istl1 istl2 istl3 istl4 chnk lsm sr al aluvp alnip aluvd alnid \
        lai_lv lai_hv sdfor slt sdor isor anor slor lsrh cvh cvl tvh tvl
    do
        cat ${tmpdir}/gginit.$var >> $ofile_gg_init
    done

    cat > ${tmpdir}/gf3 << EOT
write "${tmpdir}/gginiua.[shortName].[level]";
EOT
    $grib_filter ${tmpdir}/gf3 $ifile_gg

    #   ${tmpdir}/gginiua.o3 \ check this!
    for lev in {1..91}
    do
        for var in q
        do
            cat ${tmpdir}/gginiua.$var.$lev >> $ofile_gg_iniua
        done
    done

    for lev in {1..91}
    do
        for var in crwc cswc clwc ciwc cc
        do
            cat ${tmpdir}/gginiua.$var.$lev >> $ofile_gg_iniua
        done
    done

    # copy filtered files to the appropriate location
    if $filter_output
    then
        # move filter files to $tgtdir, in case something went wrong
        mv -f $filter_sh $filter_gg $tgtdir
        ifile_sh=$(basename $ifile_sh)
        ifile_gg=$(basename $ifile_gg)
        # keep the full output files as a backup
        mv -f $ifile_sh $tgtdir/${ifile_sh}-out+init
        mv -f $ifile_gg $tgtdir/${ifile_gg}-out+init
        # move the filtered output file to the  runtime folder
        [ -f $ofile_sh_out ] && mv -f $ofile_sh_out $ifile_sh || error "save_ic_ifs_out2init - warning! SH file $ofile_sh_out missing!"
        [ -f $ofile_gg_out ] && mv -f $ofile_gg_out $ifile_gg || error "save_ic_ifs_out2init - warning! GG file $ofile_gg_out missing!"
        # keep a copy of the last timestep
        if $save_ic_ifs_lastout
        then
            [ -f $ofile_sh_last ] && mv -f $ofile_sh_last $tgtdir/${ifile_sh}-ifs_lastout
            [ -f $ofile_gg_last ] && mv -f $ofile_gg_last $tgtdir/${ifile_gg}-ifs_lastout
        fi
    fi #$filter_output

    # delete tmp folder
    rm -rf $tmpdir

    echo "save_ic_ifs_out2init ended successfully, results are in $tgtdir"
}