#!/bin/bash

cmd=$*
echo ${cmd}

# stuff necessary for LAMMPS to run
# ================================== 
module load openmpi-x86_64
export OMP_NUM_THREADS=4
mdcmd='mpirun -n 4 lmp_g++'
# ================================== 

# directory where to store structures (.gen, .png, .dump)
strdir='./structures'
outdir='./output'
[[ -d ${strdir} ]] || mkdir ${strdir}
[[ -d ${outdir} ]] || mkdir ${outdir}

# This assumes the structure is heated and adds an oxygen every T_equilibr ps of MD

ATOMS=$1   # input structure .dump file containing velocities!!!
nextra=$2      # number of extra O atoms to add
if [ "$#" -eq 3 ] ; then
        bondtype=$3
else
        bondtype='subox'
fi
structtype='wire'

# nstart=`grep 'atoms' ${ATOMS} | awk '{print $1}'` # this will work for .data file
# nstart=`head -n 1 ${ATOMS} | awk '{print $1}'` # this is for .gen file
nstart=`grep -A 1 'NUMBER OF ATOMS' ${ATOMS} | tail -1` # this is for .dump file
nend=$(( ${nstart} + ${nextra} ))
echo "Attempting to add ${nextra} O atoms to the ${nstart} atoms of ${ATOMS}."

i=${nstart}
natoms=${i}
mderror=""
# exit 0

while [ ${i} -lt ${nend} ] || [ -n "${mderror}" ] ; do 

        if [ ${natoms} -lt ${nstart} ] ; then
                echo "CRITICAL ERROR: Less atoms then initially -- something whent terribly wrong!"
        fi

        # natoms=`grep 'atoms' ${ATOMS} | awk '{print $1}'` # this is for .data file
        # natoms=`head -n 1 ${ATOMS} | awk '{print $1}'` # this is for .gen file
        natoms=`grep -A 1 'NUMBER OF ATOMS' ${ATOMS} | tail -1` # this is for .dump file

        echo "Current structure: ${ATOMS}; i=${i}; natoms=${natoms}"

        # Check if lammps have lost an atom; this happens...
        if [ ${natoms} -eq ${i} ] ; then
                # Equilibrated system due to successful MD run: attempt to add oxygen
                echo "Attempting to add O $(($i+1))."
                poxy_oxidise -m ${bondtype} -s ${structtype} -n 1 ${ATOMS} add.data && ((i++))
                if [ ${natoms} -eq ${i} ] ; then
                        # we extpected i > natoms, if addition of O had went well
                        echo ""
                        echo "ERROR: O addition $(($i+1)) failed. Will do extra MD steps before attempting again"
                        echo ""
                else
                        echo "Successfully added 1 O: i=${i}"
                        # record the addition of O
                        /bin/mv poxy_oxidise.dbg.log output/poxy_oxidise.${i}.log
                fi
        else
                # We've tried to run MD with N+1 atoms but it failed, so we look at the 
                # last good .dump file, which has N atoms, while i = N+1
                echo ""
                echo ${mderror}
                echo "ERROR: MD failed after addition of O ${i}. Doing extra MD steps before attempting again."
                echo ""
                poxy_structure ${ATOMS} add.data
                (( i-- ))
                # exit 0
        fi

        echo ""
        echo "Running MD with ${i} atoms." 
        echo ""
        # The following is LAMMPS specific invocation of MD engine
        ${mdcmd} -in md.equilibrate.in -sc none
        mderror=`grep ERROR add.log`
        if [ -z "${mderror}" ] ; then
                echo "MD finished successfully"
                # view the structure
                poxy_structure add.dump add.gen
                poxy_view -r '90y' '-c {"Si":1.385, "O":0.85}' add.gen
                # Store tagged version of output and structure files
                # to structures directory:
                /bin/cp -f add.gen ${strdir}/add.${i}.gen   # structure after equilibration.
                /bin/mv    add.png ${strdir}/add.${i}.png   # picture of structure after MD equilibration
                # to output directory:
                /bin/mv    add.dump    ${outdir}/add.${i}.dump       # structure dump from lammps MD during equilibration.
                /bin/mv    add.log     ${outdir}/add.${i}.log       # thermo output from lammps MD during equilibration
                # make up a new current structure
                ATOMS=${outdir}/add.${i}.dump
        fi
done

# comment the line below if ffmpeg not available
ffmpeg -y -framerate 5 -start_number ${nstart} -i structures/add.%02d.png -c:v libx264 -vf scale=640:-2 -r 30 -pix_fmt yuv420p oxidation.mp4

exit 0