#!/bin/bash
########################################################################
######## Library of bash functions for the PEM launching script ########
########################################################################

# To end the launching script
endlaunch() {
    # Restore the previous value of LC_NUMERIC
    LC_NUMERIC=$OLD_LC_NUMERIC

    date
    echo "Successful end of the launching script for the PEM simulation."
    exit 0
}

# To end the launching script with error
errlaunch() {
    # Restore the previous value of LC_NUMERIC
    LC_NUMERIC=$OLD_LC_NUMERIC

    date
    echo "End with error of the launching script for the PEM."
    exit 1
}

# To check what is the job scheduler
function job_scheduler() {
    if command -v squeue &> /dev/null; then
        echo "SLURM is installed on $machine."
        name_job="#SBATCH --job-name="
        kill_job="scancel"
        submit_job="sbatch --parsable"
        submit_dependjob="sbatch --parsable --dependency"
        sed -i 's/\$PBS_JOBID/\$SLURM_JOB_ID/g' PEMrun.job
    elif command -v qstat &> /dev/null; then
        echo "PBS/TORQUE is installed on $machine."
        name_job="#PBS -N "
        kill_job="qdel"
        submit_job="qsub"
        submit_dependjob="qsub -W depend"
        sed -i 's/\$SLURM_JOB_ID/\$PBS_JOBID/g' PEMrun.job
    else
        echo "Error: neither SLURM nor TORQUE/PBS is installed on $machine!"
        echo "You need to adapt the script to your job scheduler."
        errlaunch
    fi
}

# To check if everything necessary for the launching script is ok
checklaunch() {
    # 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

    if [ -v n_mars_years ] && [ ! -z "$n_mars_years" ]; then
        if [ $n_mars_years -lt 1 ]; then
            echo "Error: the value of 'n_mars_years' must be >0!"
            errlaunch
        fi
    elif [ -v n_earth_years ] && [ ! -z "$n_earth_years" ]; then
        if [ $n_earth_years -lt 1 ]; then
            echo "Error: the value of 'n_earth_years' must be >0!"
            errlaunch
        fi
    else
        echo "Error: no number of years to be simulated has been set!"
        errlaunch
    fi
    if [ $nPCM_ini -lt 2 ] || [ -z "$nPCM_ini" ]; then
        echo "Error: the value of 'nPCM_ini' must be >1!"
        errlaunch
    fi
    if [ $nPCM -lt 2 ] || [ -z "$nPCM" ]; then
        echo "Error: the value of 'nPCM' must be >1!"
        errlaunch
    fi
    if [ ! -f "PCMrun.job" ]; then
        echo "Error: file \"PCMrun.job\" does not exist in $dir!"
        errlaunch
    fi
    if [ ! -f "PEMrun.job" ]; then
        echo "Error: file \"PEMrun.job\" does not exist in $dir!"
        errlaunch
    fi
    if [ ! -f "run_PCM.def" ]; then
        echo "Error: file \"run_PCM.def\" does not exist in $dir!"
        errlaunch
    fi
    if [ ! -f "run_PEM.def" ]; then
        echo "Error: file \"run_PEM.def\" does not exist in $dir!"
        errlaunch
    fi
    if [ ! -f "context_lmdz_physics.xml" ]; then
        echo "Error: file \"context_lmdz_physics.xml\" does not exist in $dir!"
        errlaunch
    fi
    if [ ! -f "field_def_physics_mars.xml" ]; then
        echo "Error: file \"field_def_physics_mars.xml\" does not exist in $dir!"
        errlaunch
    fi
    if [ ! -f "file_def_physics_mars.xml" ]; then
        echo "Error: file \"file_def_physics_mars.xml\" does not exist in $dir!"
        errlaunch
    fi
    if [ ! -f "iodef.xml" ]; then
        echo "Error: file \"iodef.xml\" does not exist in $dir!"
        errlaunch
    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
    if [ $dim -ne 1 ]; then
        job_scheduler
    fi
}

# To convert Earth years into Mars years
convertyears() {
    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 ]; then
        n_myear=$n_mars_years
        echo "Number of years to be simulated: $n_myear Martian years."
    elif [ -v 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."
    fi
}

# To initialize the launching script
initlaunch() {
    echo "This is a chained simulation for PEM and PCM runs in $dir on $machine by $user."
    convertyears
    i_myear=0
    iPEM=1
    iPCM=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
    echo $i_myear $n_myear $convert_years $iPCM $iPEM $nPCM $nPCM_ini > info_PEM.txt
}

# To submit the PCM runs
# arg1: model dimension
# arg2: number of PCM runs to launch
# arg3: local number of the PCM run from which to start (optional)
submitPCM() {
    find . -type f -name "PCMrun*.job" ! -name "PCMrun.job" -delete
    ii=1
    if [ ! -z $3 ]; then
        ii=$3
    fi
    if [ $i_myear -lt $n_myear ]; then
        echo "Run PCM $iPCM: call $ii/$2..."
        if [ $1 -eq 1 ]; then # 1D model
            sed -i "s/^k=[0-9]\+$/k=$(echo "3 - $nPCM_ini" | bc -l)/" PCMrun.job
            ./PCMrun.job
        else # 3D model
            cp PCMrun.job PCMrun${iPCM}.job
            sed -i -E "s/($name_job[^0-9]*[0-9]*[^0-9]*)[0-9]+$/\1${iPCM}/" PCMrun${iPCM}.job
            sed -i "s/^k=[0-9]\+$/k=$(echo "3 - $nPCM_ini" | bc -l)/" PCMrun${iPCM}.job
            jobID=$(eval "$submit_job PCMrun${iPCM}.job")
            # Create a file to cancel the dependent jobs of the cycle
            echo "#!/bin/bash" > kill_launchPEM.sh
            chmod +x kill_launchPEM.sh
            echo $kill_job $jobID >> kill_launchPEM.sh
        fi
        ((iPCM++))
        ((i_myear++))
        ((ii++))
    else
        endlaunch
    fi
    for ((i = $ii; i <= $2; i++)); do
        if [ $i_myear -lt $n_myear ]; then
            echo "Run PCM $iPCM: call $i/$2..."
            if [ $1 -eq 1 ]; then # 1D model
                sed -i "s/^k=[0-9]\+$/k=$(echo "$i + 2 - $nPCM_ini" | bc -l)/" PCMrun.job
                ./PCMrun.job
            else # 3D model
                cp PCMrun.job PCMrun${iPCM}.job
                sed -i -E "s/($name_job[^0-9]*[0-9]*[^0-9]*)[0-9]+$/\1${iPCM}/" PCMrun${iPCM}.job
                sed -i "s/^k=[0-9]\+$/k=$(echo "$i + 2 - $nPCM_ini" | bc -l)/" PCMrun${iPCM}.job
                jobID=$(eval "$submit_dependjob=afterok:${jobID} PCMrun${iPCM}.job")
                echo $kill_job $jobID >> kill_launchPEM.sh
            fi
            ((iPCM++))
            ((i_myear++))
        else
            endlaunch
        fi
    done
}

# To submit the PEM run
# arg1: model dimension
submitPEM() {
    if [ $i_myear -lt $n_myear ]; then
        echo "Run PEM $iPEM"
        if [ $1 -eq 1 ]; then # 1D model
            ./PEMrun.job
        else # 3D model
            sed -i -E "s/($name_job[^0-9]*[0-9]*[^0-9]*)[0-9]+$/\1${iPEM}/" PEMrun.job
            jobID=$(eval "$submit_job PEMrun.job")
            # Create a file to cancel the dependent jobs of the cycle
            echo "#!/bin/bash" > kill_launchPEM.sh
            chmod +x kill_launchPEM.sh
            echo $kill_job $jobID >> kill_launchPEM.sh
        fi
    else
        endlaunch
    fi
}

# To make one cycle of PCM and PEM runs
# arg1: model dimension
# arg2: number of PCM runs to launch
# arg3: local number of the PCM run from which to start (optional)
cyclelaunch() {
    # PCM runs
    submitPCM $1 $2 $3

    # PEM run
    if [ $i_myear -lt $n_myear ]; then
        echo "Run PEM $iPEM"
        if [ $1 -eq 1 ]; then # 1D model
            ./PEMrun.job
        else # 3D model
            sed -i -E "s/($name_job[^0-9]*[0-9]*[^0-9]*)[0-9]+$/\1${iPEM}/" PEMrun.job
            jobID=$(eval "$submit_dependjob=afterok:${jobID} PEMrun.job")
            echo $kill_job $jobID >> kill_launchPEM.sh
        fi
    else
        endlaunch
    fi
}

# To clean files after the starting run of the relaunch
# arg1: file name prefix to clean
# arg2: file name extension to clean
# arg3: file number from which to clean
cleanfiles() {
    prefix=$1
    extension=$2
    if [ -z "$extension" ]; then
        for file in ${prefix}*; do
            num=${file#$prefix}
            if [[ $num =~ ^[0-9]+$ ]] && [ $num -gt $3 ]; then
                rm $file
            fi
        done
    else
        for file in ${prefix}*${extension}; do
            num=${file#$prefix}
            num=${num%$extension}
            if [[ $num =~ ^[0-9]+$ ]] && [ $num -gt $3 ]; then
                rm $file
            fi
        done
    fi
}

# To relaunch from PCM run
# arg1: model dimension
relaunchPCM() {
    iPCM=$(($irelaunch + 1))
    cleanfiles diags/diagfi .nc $irelaunch
    cleanfiles diags/data2reshape .nc $irelaunch
    cleanfiles "out_PCM/run" "" $irelaunch
    cleanfiles starts/restart1D .txt $irelaunch
    cleanfiles starts/restart .nc $irelaunch
    cleanfiles starts/restartfi .nc $irelaunch
    cp starts/restartfi${irelaunch}.nc startfi.nc
    if [ -f "starts/restart${irelaunch}.nc" ]; then
        cp starts/restart${irelaunch}.nc start.nc
    elif [ -f "starts/restart1D${irelaunch}.txt" ]; then
        cp starts/restart1D${irelaunch}.txt start1D.txt
    fi
    if [ $irelaunch -le $nPCM_ini ]; then
        # PCM relaunch during the initialization cycle
        iPEM=1
        cleanfiles diags/diagpem .nc $iPEM
        cleanfiles "out_PEM/run" "" $iPEM
        cleanfiles starts/restart1D_postPEM .txt $iPEM
        cleanfiles starts/restart_postPEM .nc $iPEM
        cleanfiles starts/restartfi_postPEM .nc $iPEM
        cleanfiles starts/restartpem .nc $iPEM
        i_myear=$irelaunch
        sed -i "1s/.*/$i_myear $n_myear $convert_years $iPCM $iPEM $nPCM $nPCM_ini/" info_PEM.txt
        rm -f startpem.nc
        if [ $irelaunch -eq $(($nPCM_ini - 1)) ]; then
            cp diags/data2reshape${irelaunch}.nc data2reshape_Y1.nc
            cyclelaunch $1 $nPCM_ini $iPCM
        elif [ $irelaunch -eq $nPCM_ini ]; then
            cp diags/data2reshape$(($irelaunch - 1)).nc data2reshape_Y1.nc
            cp diags/data2reshape${irelaunch}.nc data2reshape_Y2.nc
            submitPEM $1 # The next job is a PEM run
        else
            cyclelaunch $1 $nPCM_ini $iPCM
        fi
    else
        # PCM relaunch during a cycle
        iPEM=$((($irelaunch - $nPCM_ini)/$nPCM + 1))
        il=$(echo "($irelaunch - $nPCM_ini)%$nPCM" | bc -l)
        cleanfiles diags/diagpem .nc $iPEM
        cleanfiles "out_PEM/run" "" $iPEM
        cleanfiles starts/restart1D_postPEM .txt $iPEM
        cleanfiles starts/restart_postPEM .nc $iPEM
        cleanfiles starts/restartfi_postPEM .nc $iPEM
        cleanfiles starts/restartpem .nc $iPEM
        cp starts/restartpem${iPEM}.nc startpem.nc
        if [ $il -eq $(($nPCM - 1)) ]; then
            i_myear=$(($(awk "NR==$iPEM {print \$1}" "info_PEM.txt") + $il))
            sed -i "1s/.*/$i_myear $n_myear $convert_years $iPCM $iPEM $nPCM $nPCM_ini/" info_PEM.txt
            cp diags/data2reshape${irelaunch}.nc data2reshape_Y1.nc
            cyclelaunch $1 $nPCM $il
        elif [ $il -eq 0 ]; then
            i_myear=$(($(awk "NR==$iPEM {print \$1}" "info_PEM.txt") + $nPCM))
            sed -i "1s/.*/$i_myear $n_myear $convert_years $iPCM $iPEM $nPCM $nPCM_ini/" info_PEM.txt
            cp diags/data2reshape$(($irelaunch - 1)).nc data2reshape_Y1.nc
            cp diags/data2reshape${irelaunch}.nc data2reshape_Y2.nc
            submitPEM $1 # The next job is a PEM run
        else
            i_myear=$(($(awk "NR==$iPEM {print \$1}" "info_PEM.txt") + $il))
            sed -i "1s/.*/$i_myear $n_myear $convert_years $iPCM $iPEM $nPCM $nPCM_ini/" info_PEM.txt
            cyclelaunch $1 $nPCM $il
        fi
    fi
}

# To relaunch from PEM run
# arg1: model dimension
relaunchPEM() {
    iPEM=$irelaunch
    iPCM=$(($nPCM_ini + ($nPCM - 1)*$irelaunch + 1))
    i_myear=$(awk "NR==$(($iPEM + 1)) {print \$1}" "info_PEM.txt")
    sed -i "1s/.*/$i_myear $n_myear $convert_years $iPCM $iPEM $nPCM $nPCM_ini/" info_PEM.txt
    cleanfiles diags/diagfi .nc $(($iPCM - 1))
    cleanfiles "out_PCM/run" "" $(($iPCM - 1))
    cleanfiles starts/restart1D .txt $(($iPCM - 1))
    cleanfiles starts/restart .nc $(($iPCM - 1))
    cleanfiles starts/restartfi .nc $(($iPCM - 1))
    cleanfiles diags/data2reshape .nc $(($iPCM - 1))
    cleanfiles diags/diagpem .nc $irelaunch
    cleanfiles "out_PEM/run" "" $irelaunch
    cleanfiles starts/restart1D_postPEM .txt $irelaunch
    cleanfiles starts/restart_postPEM .nc $irelaunch
    cleanfiles starts/restartfi_postPEM .nc $irelaunch
    cleanfiles starts/restartpem .nc $irelaunch
    cp starts/restartpem${irelaunch}.nc startpem.nc
    cp starts/restartfi_postPEM${irelaunch}.nc startfi.nc
    if [ -f "starts/restart_postPEM${irelaunch}.nc" ]; then
        cp starts/restart_postPEM${irelaunch}.nc start.nc
    elif [ -f "starts/restart1D_postPEM${irelaunch}.txt" ]; then
        cp starts/restart1D_postPEM${irelaunch}.txt start1D.txt
    fi
    cyclelaunch $1 $nPCM
}