#!/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 if SLURM is the job scheduler
function is_slurm() {
    if [ ! -x $(command -v squeue) ]; then
        echo "Error: the job scheduler is not SLURM on $machine!"
        echo "You need to adapt the script to your case."
        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

    is_slurm
    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 "jobPCM.slurm" ]; then
        echo "Error: file \"jobPCM.slurm\" does not exist in $dir!"
        errlaunch
    fi
    if [ ! -f "jobPEM.slurm" ]; then
        echo "Error: file \"$jobPEM.slurm\" 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
}

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

# To submit the PCM runs
submitPCM() {
    find . -type f -name "jobPCM*.slurm" ! -name "jobPCM.slurm" -exec rm {} +
    if [ $i_myear -lt $n_myear ]; then
        echo "Run PCM $iPCM: call 1/$1..."
        cp jobPCM.slurm jobPCM${iPCM}.slurm
        sed -i "s/#SBATCH --job-name=jobPCM.*/#SBATCH --job-name=jobPCM${iPCM}/" jobPCM${iPCM}.slurm
        sed -i "s/^k=[0-9]\+$/k=$(echo "3 - $nPCM_ini" | bc -l)/" jobPCM${iPCM}.slurm
        jobID=$(sbatch --parsable jobPCM${iPCM}.slurm)
        echo "#!/bin/bash" > kill_launchPEM.sh
        chmod +x kill_launchPEM.sh
        echo "scancel" $jobID >> kill_launchPEM.sh
        ((iPCM++))
        ((i_myear++))
    else
        endlaunch
    fi
    for ((i = 2; i <= $1; i++)); do
        if [ $i_myear -lt $n_myear ]; then
            echo "Run PCM $iPCM: call $i/$1..."
            cp jobPCM.slurm jobPCM${iPCM}.slurm
            sed -i "s/#SBATCH --job-name=jobPCM.*/#SBATCH --job-name=jobPCM${iPCM}/" jobPCM${iPCM}.slurm
            sed -i "s/^k=[0-9]\+$/k=$(echo "$i + 2 - $nPCM_ini" | bc -l)/" jobPCM${iPCM}.slurm
            jobID=$(sbatch --parsable --dependency=afterok:${jobID} jobPCM${iPCM}.slurm)
            echo "scancel" $jobID >> kill_launchPEM.sh
            ((iPCM++))
            ((i_myear++))
        else
            endlaunch
        fi
    done
}

# To submit the PEM run
submitPEM() {
    if [ $i_myear -lt $n_myear ]; then
        sed -i "s/#SBATCH --job-name=jobPEM.*/#SBATCH --job-name=jobPEM${iPEM}/" jobPEM.slurm
        jobID=$(sbatch --parsable --dependency=afterok:${jobID} jobPEM.slurm)
        echo "scancel" $jobID >> kill_launchPEM.sh
    else
        endlaunch
    fi
}

# To make one cycle of PCM and PEM runs
cyclelaunch() {
    # PCM runs
    submitPCM $1

    # PEM run
    submitPEM
}
