#!/bin/bash
#####################################################################
### Launching script for a chained simulation of PEM and PCM runs ###
#####################################################################

echo "The launching script is starting!"
echo "The output file is \"loglaunch.txt\"."

# Re-spawn the script as a background process and create a bash file to kill it in case
if [ "$1" = "bg" ]; then 
    date
    echo '#!/bin/bash' > kill_launch_pem.sh
    echo 'kill' $$ >> kill_launch_pem.sh
    chmod +x kill_launch_pem.sh
else
    nohup "$0" bg > loglaunch.txt 2>&1 &
    exit 1
fi

# A few parameters that might need be changed depending on your setup:
# Path to the arch.env to source:
source ../trunk/LMDZ.COMMON/arch.env

# Save the current value of LC_NUMERIC and set it to a locale that uses a dot as the decimal separator
OLD_LC_NUMERIC=$LC_NUMERIC
LC_NUMERIC=en_US.UTF-8


#####################################################################
#---------- Modify here the number of years to be simulated ------------
## set the number of years, either Martian or Earth years:
n_mars_years=1000
#n_earth_years=300

#------------- Modify here the number of initial PCM calls -------------
## set the number of initial PCM calls between each PEM call:
nPCM_ini=4

#---------------- Modify here the number of PCM calls ------------------
## set the number of PCM calls between each PEM call:
nPCM=2

#------------------ Modify here the name of PEM exe --------------------
## fill in the name of executable for PEM:
exePEM="pem_29_phymars_seq.e"

#-------------- Modify here the name of reshape XIOS exe ---------------
## fill in the name of executable for reshape XIOS:
exeReshape="reshape_XIOS_output_64x48x29_phymars_seq.e"
#####################################################################


#------ Check if files/directories necessary for the script exist ------
dir=`pwd`
machine=`hostname`
address=`whoami`
if [ ! -f "$exePEM" ]; then
    echo "Error: file \"$exePEM\" does not exist in $dir!"
    exit 1
fi
if [ ! -f "$exeReshape" ]; then
    echo "Error: file \"$exeReshape\" does not exist in $dir!"
    exit 1
fi
if [ ! -f "exePCM.sh" ]; then
    echo "Error: file \"exePCM.sh\" does not exist in $dir!"
    exit 1
fi
if [ ! -f "run_PEM.def" ]; then
    echo "Error: file \"run_PEM.def\" does not exist in $dir!"
    exit 1
fi
if [ ! -f "run_PCM.def" ]; then
    echo "Error: file \"run_PCM.def\" does not exist in $dir!"
    exit 1
fi
if [ ! -f "context_lmdz_physics.xml" ]; then
    echo "Error: file \"context_lmdz_physics.xml\" does not exist in $dir!"
    exit 1
fi
if [ ! -f "field_def_physics_mars.xml" ]; then
    echo "Error: file \"field_def_physics_mars.xml\" does not exist in $dir!"
    exit 1
fi
if [ ! -f "file_def_physics_mars.xml" ]; then
    echo "Error: file \"file_def_physics_mars.xml\" does not exist in $dir!"
    exit 1
fi
if [ ! -f "iodef.xml" ]; then
    echo "Error: file \"iodef.xml\" does not exist in $dir!"
    exit 1
fi
if [ ! -d "out_PCM" ]; then
    mkdir out_PCM
fi
if [ ! -d "out_PEM" ]; then
    mkdir out_PEM
fi
if [ ! -d "starts" ]; then
    mkdir starts
fi
if [ ! -d "diags" ]; then
    mkdir diags
fi

#---------------------------- Initialization ---------------------------
dir=`pwd`
machine=`hostname`
address=`whoami`
echo "This is a chained simulation for PEM and PCM runs in $dir on $machine."
myear=686.9725      # Number of Earth days in Martian year
eyear=365.256363004 # Number of days in Earth year
convert_years=$(echo "$myear/$eyear" | bc -l)
convert_years=$(printf "%.4f" $convert_years) # Rounding to the 4th decimal to respect the precision of Martian year
if [ -v n_mars_years ] && [ ! -z "$n_mars_years" ]; then
    n_myear=$n_mars_years
    echo "Number of years to be simulated: $n_myear Martian years."
elif [ -v n_earth_years ] && [ ! -z "$n_earth_years" ]; then
    n_myear=$(echo "($n_earth_years/$convert_years + 0.999999)/1" | bc) # Ceiling of n_earth_years/convert_years
    echo "Number of years to be simulated: $n_earth_years Earth years = $n_myear Martian years."
else
    echo "No number of years to be simulated has been set!"
    exit 1
fi
i_myear=0
iPEM=1
((iPCM = ($iPEM - 1)*$nPCM + 1))
cp startfi.nc starts/
if [ -f "start.nc" ]; then
    cp start.nc starts/
elif [ -f "star1D.nc" ]; then
    cp star1D.txt starts/
fi

# Create a file to manage years of the chained simulation and store some info from the PEM runs
if [ -f "info_PEM.txt" ]; then
    rm info_PEM.txt
fi
echo $i_myear $n_myear $convert_years > info_PEM.txt

#-------------------------- Initial PCM runs ---------------------------
cp run_PCM.def run.def
for ((i = 1; i <= $nPCM_ini; i++)); do
    echo "Run PCM $iPCM (initial): call $i/$nPCM_ini..."
    sed -i "s/#SBATCH --job-name=runPCM.*/#SBATCH --job-name=runPCM${iPCM}/" exePCM.sh
    sed -i "s/out_runPCM[0-9]\+/out_runPCM${iPCM}/" exePCM.sh
    sbatch -W exePCM.sh
    if [ ! -f "restartfi.nc" ]; then # Check if run ended abnormally
        echo "Error: the run PCM $iPCM crashed!"
        exit 1
    fi
    # Copy data files and prepare the next run
    mv out_runPCM${iPCM} out_PCM/run${iPCM}
    mv diagfi.nc diags/diagfi${iPCM}.nc
    if [ -f "diagsoil.nc" ]; then
        mv diagsoil.nc diags/diagsoil${iPCM}.nc
    fi
    if [ -f "stats.nc" ]; then
        mv stats.nc diags/stats${iPCM}.nc
    fi
    k=$(echo "$i + 2 - $nPCM_ini" | bc -l)
    if [ $(echo "$k < 1" | bc -l) -eq 1 ]; then # Only the last 2 years are taken for the PEM
        mv Xdiurnalave.nc diags/data2reshape${iPCM}.nc
    else
        cp Xdiurnalave.nc diags/data2reshape${iPCM}.nc
        mv Xdiurnalave.nc data2reshape_Y${k}.nc
    fi
    cp restartfi.nc starts/startfi${iPCM}.nc
    mv restartfi.nc startfi.nc
    if [ -f "restart.nc" ]; then
        cp restart.nc starts/restart${iPCM}.nc
        mv restart.nc start.nc
    elif [ -f "restart1D.txt" ]; then
        cp restart1D.txt starts/restart1D${iPCM}.txt
        mv restart1D.txt start1D.txt
    fi
    ((iPCM++))
    ((i_myear++))
    echo "Done!"
done
sed -i "1s/.*/$i_myear $n_myear $convert_years/" info_PEM.txt

#---------------------- Main loop to call PEM/PCM ----------------------
while [ $i_myear -lt $n_myear ]; do
    #--- Reshaping PCM data with XIOS
    echo "Reshaping PCM data with XIOS..."
    ./$exeReshape
    echo "Done!"

    #--- Running PEM
    echo "Run PEM $iPEM..."
    cp run_PEM.def run.def
    mv startfi.nc startfi_evol.nc
    if [ -f "start.nc" ]; then
        mv start.nc start_evol.nc
    elif [ -f "start1D.txt" ]; then
        mv start1D.txt start1D_evol.txt
    fi
    ./$exePEM > out_runPEM${iPEM} 2>&1
    if [ ! -f "restartfi_evol.nc" ]; then # Check if run ended abnormally
        echo "Error: the run PEM $iPEM crashed!"
        exit 1
    fi
    # Copy data files and prepare the next run
    mv out_runPEM${iPEM} out_PEM/run${iPEM}
    mv diagpem.nc diags/diagpem${iPEM}.nc
    if [ -f "diagsoilpem.nc" ]; then
        mv diagsoilpem.nc diags/diagsoilpem${iPEM}.nc
    fi
    cp restartpem.nc starts/startpem${iPEM}.nc
    mv restartpem.nc startpem.nc
    cp restartfi_evol.nc starts/startfi_postPEM${iPEM}.nc
    mv restartfi_evol.nc startfi.nc
    if [ -f "restart_evol.nc" ]; then
        cp restart_evol.nc starts/restart_postPEM${iPEM}.nc
        mv restart_evol.nc start.nc
    elif [ -f "restart1D_evol.txt" ]; then
        cp restart1D_evol.txt starts/restart1D_postPEM${iPEM}.txt
        mv restart1D_evol.txt start1D.txt
    fi
    ((iPEM++))
    read i_myear n_myear convert_years < info_PEM.txt
    echo "Done!"

    #--- Loop to run PCM year by year
    cp run_PCM.def run.def
    for ((i = 1; i <= $nPCM; i++)); do
        echo "Run PCM $iPCM: call $i/$nPCM..."
        sed -i "s/#SBATCH --job-name=runPCM.*/#SBATCH --job-name=runPCM${iPCM}/" exePCM.sh
        sed -i "s/out_runPCM[0-9]\+/out_runPCM${iPCM}/" exePCM.sh
        sbatch -W exePCM.sh
        if [ ! -f "restartfi.nc" ]; then # Check if run ended abnormally
            echo "Error: the run PCM $iPCM crashed!"
            exit 1
        fi
        # Copy data files and prepare the next run
        mv out_runPCM${iPCM} out_PCM/run${iPCM}
        mv diagfi.nc diags/diagfi${iPCM}.nc
        if [ -f "diagsoil.nc" ]; then
            mv diagsoil.nc diags/diagsoil${iPCM}.nc
        fi
        if [ -f "stats.nc" ]; then
            mv stats.nc diags/stats${iPCM}.nc
        fi
        k=$(echo "$i + 2 - $nPCM" | bc -l)
        if [ $(echo "$k < 1" | bc -l) -eq 1 ]; then # Only the last 2 years are taken for the PEM
            mv Xdiurnalave.nc diags/data2reshape${iPCM}.nc
        else
            cp Xdiurnalave.nc diags/data2reshape${iPCM}.nc
            mv Xdiurnalave.nc data2reshape_Y${k}.nc
        fi
        cp restartfi.nc starts/startfi${iPCM}.nc
        mv restartfi.nc startfi.nc
        if [ -f "restart.nc" ]; then
            cp restart.nc starts/restart${iPCM}.nc
            mv restart.nc start.nc
        elif [ -f "restart1D.txt" ]; then
            cp restart1D.txt starts/restart1D${iPCM}.txt
            mv restart1D.txt start1D.txt
        fi
        ((iPCM++))
        ((i_myear++))
        echo "Done!"
    done
    sed -i "1s/.*/$i_myear $n_myear $convert_years/" info_PEM.txt
done

# Restore the previous value of LC_NUMERIC
LC_NUMERIC=$OLD_LC_NUMERIC

date
echo "The launching script ended."