#!/bin/bash
#
# workscript.sh - GlueX simulation samples script for generation
#                 of Moller samples using gxtwist.
#
# author: richard.t.jones at uconn.edu
# version: february 10, 2024
#
# usage: $ ./workscript.sh <configNo> <seqNo> [<offset> [<nevents>]]

runNo=120000
nthreads=1
nevents=10000000
jmax=500

xrootdURL="root://cn440.storrs.hpc.uconn.edu"
remotepath="/Gluex/resilient/simulation/targetpol-7-2024"
inputURL="http://zeus.phys.uconn.edu/halld/targetpol-7-2024"
cvmfspath="/cvmfs/gluex.osgstorage.org/gluex/uconn1/resilient/simulation/targetpol-7-2024/xxx"
wget="wget --ca-directory=/etc/grid-security/certificates"

function usage() {
    echo "Usage: workscript.sh <configNo> <seqNo> [<offset> [<nevents>]]"
    echo "  where <configNo> = configuration number, 1..4"
    echo "        <seqNo> = job sequence number, 1..Jmax"
    echo "        <offset> : job sequence number += <offset>"
    echo "        <nevents> = number of beam electrons to simulate"
}

function clean_exit() {
    ls -l 
    if [ "$1" = "" -o "$1" = "0" ]; then
        rm -f setup randoms postcol.hddm control*.in beam.hddm bhgen.hddm *.astate *.py run.mac
        echo "Successful exit from workscript."
        exit 0
    else
        echo "Error $1 in workscript, $2"
        rm -f setup randoms postcol.hddm control*.in beam.hddm bhgen.hddm *.astate *.py run.mac
        echo "Failed exit from workscript."
        exit $1
    fi
}

function fetch_input() {
    cp $cvmfspath/$1 $2 || \
    $wget -O $2 $inputURL/$1
    return $?
}

function save_output() {
    maxretry=5
    retry=0
    while [[ $retry -le $maxretry ]]; do
        xrdcp -f $1 $xrootdURL/$remotepath/$2 2>xrdcp.err
        retcode=$?
        if [[ $retcode != 0 ]]; then
            cat xrdcp.err
        fi
        rm xrdcp.err
        if [[ $retcode = 0 ]]; then
            rm $1
            break
        elif [[ $retry -lt $maxretry ]]; then
            retry=$(expr $retry + 1)
            echo "xrdcp returned error code $retcode, waiting $retry minutes before retrying"
            sleep $(expr $retry \* 60)
        else
            retry=$(expr $retry + 1)
            echo "xrdcp returned error code $retcode, giving up"
        fi
    done
    # fall through to allow job file transfer return results, failure not fatal
    if [[ "$1" != "$2" ]]; then
        mv $1 $(basename $2)
    fi
    return 0
}

if [ $# -lt 2 -o $# -gt 4 ]; then
    usage
    exit 
fi

configNo=$1
if [[ $# -gt 3 ]]; then
    nevents=$3
fi
if [[ $# -gt 2 ]]; then
    seqNo=$(expr $2 + $3 + 1)
    seqNo3=$seqNo
else
    seqNo=$(expr $2 + 1)
    seqNo3=$(echo $seqNo | awk '{printf("%03d",$1)}')
fi
echo "job $seqNo running on" $(hostname)

fetch_input scripts.d/setup_alma9.sh setup || clean_exit $? "cannot fetch setup_alma9.sh from web server"
fetch_input config.d/randoms.in randoms || clean_exit $? "cannot fetch randoms.in from web server"
fetch_input config.d/control_postcol.in control_postcol.in  || clean_exit $? "cannot fetch control_postcol.in from web server"
fetch_input scripts.d/filter_postcol.python filter_postcol.py || clean_exit $? "cannot fetch filter_postcol.py from web server"
fetch_input config.d/control_bhgen.in control_bhgen.in  || clean_exit $? "cannot fetch control_bhgen.in from web server"
fetch_input scripts.d/filter_bhgen.python filter_bhgen.py || clean_exit $? "cannot fetch filter_bhgen.py from web server"

source ./setup
which hdgeant4

seedseq=$(expr \( $configNo \* $jmax \) + $seqNo)
rseeds=$(head -n $seedseq randoms | tail -n1)
[ $seqNo -ge 0 -o $seqNo -ge $jmax ] || clean_exit $? "simulation seqNo $seqNo is out of range"
rm -f control.in && ln -s control_postcol.in control.in || clean_exit $? "cannot create control.in"
sed -i "s/^RNDM.*/RNDM $rseeds/" control.in || clean_exit $? "RNDM card not found in control.in"
sed -i "s/^RUN[NG].*/RUNNO $runNo/" control.in || clean_exit $? "RUN card not found in control.in"
sed -i "s/^OUTF.*/OUTFILE 'postcol.hddm'/" control.in || clean_exit $? "OUTF card not found in control.in"
sed -i "s/^TRIG.*/TRIG $nevents/" control.in || clean_exit $? "TRIG card not found in control.in"
hdgeant4 || clean_exit $? "hdgeant4 crashed in photon beam simulation"

python3 filter_postcol.py $configNo postcol.hddm beam.hddm

rm -f control.in && ln -s control_postcol.in control.in || clean_exit $? "cannot create control.in"
sed -i "s/^RNDM.*/RNDM $rseeds/" control.in || clean_exit $? "RNDM card not found in control.in"
sed -i "s/^RUN[NG].*/RUNNO $runNo/" control.in || clean_exit $? "RUN card not found in control.in"
sed -i "s/^INFI.*/INFILE 'beam.hddm'/" control.in || clean_exit $? "INFI card not found in control.in"
sed -i "s/^OUTF.*/OUTFILE 'bhgen.hddm'/" control.in || clean_exit $? "OUTF card not found in control.in"
sed -i "s/^TRIG.*/TRIG $nevents/" control.in || clean_exit $? "TRIG card not found in control.in"
if [ $configNo = 1 ]; then
    cat <<EOI >run.mac
/hdgeant4/targetNuclearPolarization 0 0 1
/hdgeant4/targetElectronPolarization 0 0 1
/run/beamOn $nevents
EOI
    outfile="targetpol_++_$seqNo.hddm"
elif [ $configNo = 2 ]; then
    cat <<EOI >run.mac
/hdgeant4/targetNuclearPolarization 0 0 -1
/hdgeant4/targetElectronPolarization 0 0 1
/run/beamOn $nevents
EOI
    outfile="targetpol_+-_$seqNo.hddm"
elif [ $configNo = 3 ]; then
    cat <<EOI >run.mac
/hdgeant4/targetNuclearPolarization 0 0 1
/hdgeant4/targetElectronPolarization 0 0 1
/run/beamOn $nevents
EOI
    outfile="targetpol_-+_$seqNo.hddm"
elif [ $configNo = 4 ]; then
    cat <<EOI >run.mac
/hdgeant4/targetNuclearPolarization 0 0 -1
/hdgeant4/targetElectronPolarization 0 0 1
/run/beamOn $nevents
EOI
    outfile="targetpol_--_$seqNo.hddm"
fi
hdgeant4 run.mac || clean_exit $? "hdgeant4 crashed in pair production simulation"
rm -f control.in

python3 filter_bhgen.py $configNo bhgen.hddm $outfile

save_output $outfile $outfile || clean_exit $? "save of $outfile failed"
clean_exit