#!/bin/bash
#SBATCH -J AQME
#SBATCH -t 36:00:00
#SBATCH -c 48
#SBATCH --mem=96GB

##################################################
################    VARIABLES    #################
##################################################

totalcpu=48
childcpu=12 # Number of cpus per parallel job during batching
childmem=24GB # Amount of mem per parallel job during batching
jbatch=$(expr $totalcpu / $childcpu) # this variable controls the number of simultaneous QM calcs running
excdir=$(pwd)

##################################################
################    MODULES     ##################
##################################################

# Load modules related to AQME
module load cesga/2020 miniconda3/4.11.0
module load intel/2021.3.0 crest/2.11.2
conda activate cheminf

# Load modules related to Gaussian
module load g16/c1

##################################################
################    FUNCTIONS    #################
##################################################

### Run QM calculation. In theory any program could be called here.
task () {

g16 < $1.com > $2.log

}

### Wrapper for batching  and running jobs
gexc () {
#Variables; Use local here or otherwise it becomes global and overwrites original
local jobname=$1 # The current job
local jbatch=$2 # Number or parallel jobs.

if [ -d $jobname ]; then
        cd $jobname # We don't exit this directory within the function. Might cause issues.
else
        echo "No $jobname folder found, exiting the script"
        exit
fi
start=$(date +%s)
for i in $(ls -1vd *.com); do
        (
        file=$(echo $i | sed 's/.com//')
        echo "$file"
        task $file $jobname/$file
        ) &
        # allow to execute up to $N jobs in parallel
        if [[ $(jobs -r -p | wc -l) -ge $jbatch ]]; then
        # now there are $N jobs already running, so wait here for any job
        # to be finished so there is a place to start next one.
                wait
        fi
done
wait
end=$(date +%s)
sec=$(expr $end - $start)
echo "Time $jobname: $sec seconds"

}

##################################################
###############    MAIN SCRIPT     ###############
##################################################

# Run CSEARCH conformational search with CREST
python -m aqme --csearch --program "crest" --input "CSEARCH_SMILES.csv" --nprocs "$totalcpu"

# Run QPREP to create Gaussian input files for OPT+FREQ (different for GS and TS calcs)
python -m aqme --qprep --program "gaussian" --files "CSEARCH/S0*.sdf" --qm_input "opt=(calcfc) freq=noraman wb97xd/6-31+g(d) scrf=(smd,solvent=dichloromethane)" --mem "$childmem" --nprocs "$childcpu"
python -m aqme --qprep --program "gaussian" --files "CSEARCH/T1*.sdf" --qm_input "opt=(calcfc) freq=noraman wb97xd/6-31+g(d) scrf=(smd,solvent=dichloromethane)" --mem "$childmem" --nprocs "$childcpu"
python -m aqme --qprep --program "gaussian" --files "CSEARCH/TS*.sdf" --qm_input "opt=(calcfc,ts,noeigen,maxstep=5) freq=noraman wb97xd/6-31+g(d) scrf=(smd,solvent=dichloromethane)" --mem "$childmem" --nprocs "$childcpu"

# Run Gaussian with the generated input files
echo "Starting opt-freq calculations"
gexc "$excdir/QCALC" $jbatch
wait
cd $excdir

# Run first QCORR analysis
python -m aqme --qcorr --files "$excdir/QCALC/*log" --mem "$childmem" --nprocs "$childcpu" --freq_conv 'opt=(calcfc,maxstep=5)' --isom_type 'com' --isom_inputs "$excdir/QCALC" --amplitude_ifreq 0.3

# Run second round of Gaussian/QCORR analysis
if [ -d "$excdir/QCALC/failed/run_1/fixed_QM_inputs" ]; then
    echo "Starting one rerun of failed calculations"
            gexc "$excdir/QCALC/failed/run_1/fixed_QM_inputs" $jbatch
            wait
            cd $excdir
            python -m aqme --qcorr --files "$excdir/QCALC/failed/run_1/fixed_QM_inputs/*.log" --mem "$childmem" --nprocs "$childcpu" --isom_type 'com' --isom_inputs "$excdir/QCALC/failed/run_1/fixed_QM_inputs"
    else
            echo "No fixable errors found"
    fi

echo "Done with opt-freq calculations"
cd $excdir

# Create SPC input files
python -m aqme --qprep --program gaussian --files "$excdir/QCALC/success/*.log" --destination "$excdir/QCALC/success" --qm_input 'wb97xd/def2qzvpp scrf=(smd,solvent=dichloromethane)' --mem "$childmem" --nprocs "$childcpu" --suffix "SPC"

# Run Gaussian with the generated SPC input files
echo "Starting SP calculations"
gexc "$excdir/QCALC/success" $jbatch
wait
cd $excdir

# Run GoodVibes to obtain thermochemistry data, Boltzmann population and XYZ coordinates
python -m goodvibes --xyz -c 1 "$excdir/QCALC/success/*.log" --spc SPC --imag --pes Grel.yaml