#!/bin/bash
dohelp=0
cutoff="0"
cmdsuffix=""
volrelaxcommand="runstruct_vasp -w vaspvol.wrap"
relaxcommand="runstruct_vasp -w vasp.wrap"
chaincommand="runstruct_vasp -w vaspneb.wrap"
staticcommand="runstruct_vasp -w vaspstatic.wrap"
infdetcommand="runstruct_vasp -p -w vaspid.wrap"
doid=0
idopt="-d"
idf="0.5"
idja="0.0"
idjc="0.0"
doneb=0
doextract=0
nimage=6
dovib=0
makevaspfiles=0
contif=0

if [[ $1 == "" || $1 == "-h" ]]
then
    cat - <<EOF
This command implements the "inflection-detection" method introduced in
A. van de Walle, Q. Hong, S. Kadkhodaei and R. Sun, "The free energy of mechanically unstable phases"
Nat. Com. 6, 7559 (2015) doi:10.1038/ncomms8559 .
and
A. van de Walle and S. Kadkhodaei and R. Sun and Q.-J. Hong, "Epicycle method for elasticity limit calculations",
Phys. Rev. B 95 144113 (2017) doi:10.1103/PhysRevB.95.144113 .

It calculates the energy of a structure that is potentially mechanically unstable.
In a nutshell: the method finds either a local minimum (for a mechanically stable phase)
or the inflection point on a path joining the unrelaxed and fully relaxed structures
(for a mechanically unstable phase).

Syntax: robustrelax_vasp [options] command_prefix
where options are
  -id             Inflection Detection method
  -idop "options" Options passed to the infdet command.
  -neb            Use the Nudged Elastic Band method to find minimum energy path
                  Otherwise: simply interpolate linearly.
  -ex             Only extract data form VASP files. Do not run anything.
  -c  [real]      Relaxation magnitude cutoff needed to activate robust relax algorithm (default: 0).
  -ni [integer]   Number of images in the minimum energy path (default: 6).
  -rc [string]    Command for full relaxation          (Default: runstruct_vasp -w vasp.wrap).
  -vc [string]    Command for volume only relaxation   (Default: runstruct_vasp -w vaspvol.wrap).
  -cc [string]    Command for chain (neb) calculations (Default: runstruct_vasp -w vaspneb.wrap).
  -sc [string]    Command for static runs              (Default: runstruct_vasp -w vaspstatic.wrap).
  -ic [string]    Command for inflection-detection runs (-id option) (Default: runstruct_vasp -p -w vaspid.wrap).
  -idf [real]     Starting point fraction for inflection detection (default: 0.5)
  -ja
  -jc             By how much to jitter atoms (-ja) and cell parameters (-jt) for the initial conditions
                  for inflection-detection method.
  -vib            Update vibrational data in svib_ht file from vol_0/svib_ht file.
  -mk             Make auxiliary vasp input files (vaspvol.wrap, vaspneb.wrap, vaspstatic.wrap, vaspf.wrap) from vasp.wrap.
                  Do not run anything.
  -cln            clean files for a full restart.
  -cip            continue (an prior run) if possible (only implemented for -id algorithm).
  -d              Use all defaults values.
  -h              Display more help.
  command_prefix  Prefix needed for vasp to run on a remote machine or in parallel (e.g. mpirun).

Note: the commands specified by -rc,-vc,-cc,-sc can be set to empty "" to skip the corresponding
  calculation step (useful to restart an aborted run).
EOF
    if [[ $1 == "" ]]
    then
	exit
    fi
fi
while [[ $1 != "" ]]
do
    if [[ $1 == "-d" ]]
    then
	echo -n
    elif [[ $1 == "-h" ]]
    then
	dohelp=1
    elif [[ $1 == "-mk" ]]
    then
	makevaspfiles=1
    elif [[ $1 == "-cln" ]]
    then
	clnrestart=1
    elif [[ $1 == "-vib" ]]
    then
	dovib=1
    elif [[ $1 == "-cc" ]]
    then
	shift
	chaincommand="$1";
    elif [[ $1 == "-rc" ]]
    then
	shift
	relaxcommand="$1";
    elif [[ $1 == "-vc" ]]
    then
	shift
	volrelaxcommand="$1";
    elif [[ $1 == "-sc" ]]
    then
	shift
	staticcommand="$1";
    elif [[ $1 == "-ic" ]]
    then
	shift
	infdetcommand="$1";
    elif [[ $1 == "-cip" ]]
    then
	contif=1;
    elif [[ $1 == "-c" ]]
    then
	shift
	cutoff=$1
    elif [[ $1 == "-ni" ]]
    then
	shift
	nimage=$1
    elif [[ $1 == "-id" ]]
    then
	doid=1
    elif [[ $1 == "-idop" ]]
    then
	shift
	idopt="$1"
    elif [[ $1 == "-idf" ]]
    then
	shift
	idf="$1"
    elif [[ $1 == "-ja" ]]
    then
	shift
	idja="$1"
    elif [[ $1 == "-jc" ]]
    then
	shift
	idjc="$1"
    elif [[ $1 == "-neb" ]]
    then
	doneb=1
    elif [[ $1 == "-ex" ]]
    then
	doextract=1
    else
	cmdsuffix="$*"
	break
    fi
    shift
done

if [[ $dohelp == 1 ]]
then
    cat - <<EOF

This command does the following:
1) Fully relax the structure (in str_hint.out or str.out) to create str_relax.out,
   just like runstruct_vasp.
2) If the relaxation cutoff (-c=...) no exceeded, stop. Typically, should be set to -c=0.05.
3) Relax the volume only of structure in str_sup.out or str.out. Results are in directory 00/
4) Do a neb (-neb) or infection detection (-id) calculation between the two structures obtained in step 1 and 3
   or just a linear interpolation between them (if no -neb or -id). Results are in directories ??/
   Note that you need VTST Tools from http://theory.cm.utexas.edu/vtsttools/ for the neb method
   with variable cell shape.
5) Do static calculations in each directory ??/
6) Find the inflection point on the path optimized in step 4. If it does not exist, find the
   minimum energy along that path.
7) Report the energy found in 6 in the file energy and the corresponding geometry in str_relax.out

Notes:
1) str_beg.out and str_end.out contain the extremities of the path.
   If you want to restart from scratch, delete these files.
2) For high symmetry phases (e.g. bcc,fcc,hcp) you may need to break the symmetry of the phase for it to relax.
   For that purpose, you can use, e.g.,
      symbrklib W fcc bcc
   or
      cellcvrt -ja=0.01 -jc=0.01 < str.out > str_hint.out .
   Sometimes a supercell is needed (a phonon analysis will tell you that - or symbrklib for bcc,fcc,hcp).
EOF
    exit
fi

if [[ $makevaspfiles == 1 ]]
then
    if [[ ! -e vasp.wrap ]]
    then
	echo Please provide a vasp.wrap file.
	exit
    fi
    grep -v -i -e ISIF -e DOSTATIC vasp.wrap > vaspvol.wrap
    echo ISIF = 7 >> vaspvol.wrap
    grep -v -i -e  ISMEAR -e SIGMA -e NSW -e IBRION -e ISIF -e DOSTATIC vasp.wrap > vaspstatic.wrap
    cat - >> vaspstatic.wrap <<EOF
ISMEAR = -5
IBRION = -1
ISIF = 2
EOF
    grep -v -i -e IBRION -e DOSTATIC vasp.wrap > vaspneb.wrap
    echo IMAGES = $(( nimage -2 )) >> vaspneb.wrap
	cat - >>vaspneb.wrap <<EOF
SPRING = -5
LCLIMB = .FALSE.
LNEBCELL = .TRUE.
IBRION = 3
POTIM = 0
EDIFFG = -5e-2
IOPT = 3
EOF
    grep -v -i -e NSW -e IBRION -e ISIF -e DOSTATIC vasp.wrap > vaspf.wrap
    cat - >> vaspf.wrap <<EOF
LREAL = AUTO
ICHARG = 1
ISTART = 2
IBRION = -1
ISIF = 2
EOF
    grep -v -i -e NSW -e IBRION -e ISIF -e DOSTATIC -e KPPRA -e EDIFF -e EDIFFG vasp.wrap > vaspid.wrap
    cat - >> vaspid.wrap <<EOF
IBRION = -1
ISIF = 2
ICHARG = 1
KPPRA = UPDIR
ISYM = 0
EDIFF = 1e-5
ISTART = 2
EOF
    exit
fi

if [[ $clnrestart == 1 ]]
then
    rm str_beg.out str_end.out
    rm -r [0-9][0-9]
    exit
fi

if [[ $contif == 1 ]]
then
    if [[ -e OSZICAR && ! -e str_relax.out ]]
    then
	runstruct_vasp -ex
	rm -f error
	if [[ -e str_relax.out ]]
	then
	    if [[ -e str_hint.out ]]
	    then
		cp str_hint.out str_hint.out.org
	    fi
	    echo Re-relaxing pre-relaxed geometry...
	    mv -f str_relax.out str_hint.out
	fi
    fi
    if [[ -e str_end.out ]]
    then
	relaxcommand="";
	echo Skipping relax step...
	if [[ -e str_beg.out ]]
	then
	    volrelaxcommand="";
	    echo Skipping volume relax step...
	    if [[ -e 01/cstr_relax.out ]]
	    then
		infdetcommand="";
		echo Skipping inflection detection step...
		diff 01/cstr_relax.out str_relax.out > /dev/null
		if [[ $? == 0 ]]
		then
		    staticcommand="";
		    echo Skipping static run step...
		fi
	    fi
	fi
    fi
fi

if [[ $dovib == 1 ]]
then
    for subdir in vol_0 smsqs/vol_0
    do
	if [[ -e ${subdir}/svib_ht ]]
	then
	    if [[ ! -e confmult.in ]]
	    then
		echo using scaled ${subdir}/svib_ht
		echo `cat ${subdir}/svib_ht` `cellcvrt -pn < str.out` `cellcvrt -pn < ${subdir}/str_relax.out` | awk '{print $1*$2/$3}' > svib_ht
	    else
		echo using scaled ${subdir}/svib_ht and confmult.in
		echo `cat ${subdir}/svib_ht` `cellcvrt -pn < str.out` `cellcvrt -pn < ${subdir}/str_relax.out` `cat confmult.in` | awk '{print $1*$2/$3+$2*log($4)}' > svib_ht
	    fi
	    exit
	fi
    done
    echo No svib_ht file in `pwd`/vol_0 or `pwd`/smsqs/vol_0
    exit
fi


if [[ $doextract == 0 ]]
then
    if [[  $relaxcommand != "" ]]
    then
	echo $relaxcommand $cmdsuffix
	$relaxcommand $cmdsuffix
    fi
    if [[ -e error ]]
    then
	echo Error during relaxation run.
	exit 1
    fi

    relmag=`checkrelax -1`

#    if [[ ! (( $relmag > $cutoff || -e str_sup.out )) ]]
    if [[ ! (( $relmag > $cutoff )) ]]
    then
	if [[ $relaxcommand != "" ]]
	then
	    mv energy energy_sup
	    echo `cat energy_sup` `cellcvrt -pn < str.out` `cellcvrt -pn < str_relax.out` | awk '{print $1*$2/$3}' > energy
	fi
    else
	if [[ $relaxcommand != "" ]]
	then
	    mv energy energy_end
	    if [[ ! -e str_end.out ]]
	    then
		cp str_relax.out str_end.out
	    fi
	fi
	mkdir 00
	if [[ $volrelaxcommand != "" ]]
	then
	    if [[ -e str_sup.out ]]
	    then
		cp str_sup.out 00/str.out
	    else
		cp str.out 00/
	    fi
	    cd 00
	    echo "Now working in" `pwd`
	    echo $volrelaxcommand $cmdsuffix
	    $volrelaxcommand $cmdsuffix
	    if [[ -e error ]]
	    then
		echo Error during static run in 00/
		touch ../error
		exit 1
	    fi
	    cd ..
	fi
	if [[ ! -e str_beg.out ]]
	then
	    cp 00/str_relax.out str_beg.out
	fi
	if [[ $doid == 1 ]]
	then
	    if [[ ! -e KPOINTS ]]
	    then
		cp KPOINTS.static KPOINTS
	    fi

	    mkdir 01
	    if [[ -e 01/str_hint.in ]]
	    then
		cp 01/str_hint.in 01/str.in
	    else
		strpath -s1=str_beg.out -s2=str_end.out -f=${idf} | cellcvrt -ja=$idja -jc=$idjc > 01/str.in
	    fi
	    if [[ $infdetcommand != "" ]]
	    then
		pushd 01 > /dev/null
		echo "Now working in" `pwd`
		if [[ -e str_hint.in ]]
		then
		    cp str_hint.in str.in
		fi
		if [[ -e busy ]]
		then
		    rm busy
		fi
		echo Running infdet $idopt
		rm -f error stop busy
                if [[ -e infdet.log ]]
                then
                    for (( baknum=1 ; ; baknum++ ))
                    do
                        if [[ ! -e infdet.save${baknum} ]]
                        then
                            cp infdet.log infdet.save${baknum}
                            break;
                        fi
                    done
                fi
		( infdet $idopt > infdet.log ; if [[ $? == 1 ]] ; then touch error ; fi ; touch stop )  &
		while [[ 1 ]]
		do
		    while [[ ! -e busy && ! -e stop && ! -e error ]]
		    do
			sleep 2
		    done
		    if [[ -e stop || -e error ]]
		    then
			break
		    fi
		    echo vasp
		    rm stress.out force.out > /dev/null
		    $infdetcommand $cmdsuffix
		    cat stress.out force.out
		    rm busy
		    echo waiting
		done
		rm -f stop busy
		if [[ -e error ]]
		then
                    echo Error during inflection detection run in 01/
		    touch ../error
                    exit 1
		fi
		if [[ ! -e cstr_relax.out ]]
		then
		    cp str_current.out cstr_relax.out
		fi
		popd
	    fi
	    if [[ $staticcommand != "" ]]
	    then
		pushd 01
		cp cstr_relax.out str.out
		echo "$staticcommand $cmdsuffix"
		$staticcommand $cmdsuffix
		popd
	    fi
            if [[ -e error ]]
            then
                echo Error during static run in 01/
                exit 1
            fi
	    echo `cat 01/energy` `cellcvrt -pn < str.out` `cellcvrt -pn < 01/str_relax.out` | awk '{print $1*$2/$3}' > energy
	    ( echo `cellcvrt -pn < str.out` `cellcvrt -pn < 01/str_relax.out` ; cat 01/dos.out )| awk 'BEGIN {getline; sc=$1/$2; getline; print  $1*sc,$2,$3} {print $1*sc;}' > dos.out
	    cp 01/cstr_relax.out str_relax.out
	    exit 0
	elif [[ $doneb == 1 ]]
	then
	    strpath -s1=str_beg.out -s2=str_end.out -ni=$nimage
	    if [[ $chaincommand != "" ]]
	    then
		lastdir=`ls -d -1 ?? | sort -n | tail -n -1`
		cp ${lastdir}/str.out ${lastdir}/str_relax.out

		echo "$chaincommand -nr"
		foreachfile str.out "$chaincommand -nr"
		for file in POTCAR KPOINTS INCAR vasp.in
		do
		    rm ??/${file}
		done
		echo  $chaincommand "$cmdsuffix"
		$chaincommand "$cmdsuffix"
		foreachfile str.out "$chaincommand -ex"
		foreachfile str_relax.out cp str_relax.out str_hint.out
		foreachfile str_relax.out cp OSZICAR OSZICAR.neb
		cp OUTCAR OUTCAR.neb \; gzip OUTCAR.neb
		if [[ $staticcommand != "" ]]
		then
		    echo "$staticcommand $cmdsuffix"
		    foreachfile str_relax.out "$staticcommand $cmdsuffix"
		fi
	    fi
	else
	    strpath -s1=str_beg.out -s2=str_end.out -ni=$nimage
	    if [[ $volrelaxcommand != "" ]]
	    then
		echo "$volrelaxcommand $cmdsuffix"
		foreachfile str.out "$volrelaxcommand $cmdsuffix"
	    fi
	    if [[ $staticcommand != "" ]]
	    then
		echo "$staticcommand $cmdsuffix"
		foreachfile str_relax.out "$staticcommand $cmdsuffix"
	    fi
	fi
    fi
fi
if [[ $doextract == 1 || (( $relmag > $cutoff || -e str_sup.out )) ]]
then
    if [[ ! -e 00 ]]
    then
	runstruct_vasp -ex
	cp energy energy_end
	echo `cat energy_end` `cellcvrt -pn < str.out` `cellcvrt -pn < str_relax.out` | awk '{print $1*$2/$3}' > energy
    else
	if [[ ! -e energy_end ]]
	then
	    cp energy energy_end
	fi
	lastenergy=`ls -d -1 ?? | sort -n | tail -n -1`
#echo $lastenergy
	if [[ ! -e ${lastenergy}/energy ]]
	then
	    cp energy_end ${lastenergy}/energy
	fi
	((curnimage=lastenergy+0))
#echo $curnimage
	cat `ls -1 ??/energy | sort -n` > epath.out
#    cat epath.out | strpath -ci > energy
	echo `cat epath.out | strpath -ci` `cellcvrt -pn < str.out` `cellcvrt -pn < str_beg.out` | awk '{print $1*$2/$3}' > energy
	iploc=$(cat epath.out | strpath -cil | awk '{print $1*('$curnimage')}')
#echo $iploc
	dirb=$(echo $iploc | awk '{n=int($1); printf("%02i",n)}')
#echo $dirb
	if [[ $iploc == $curnimage ]]
	then
#echo cp only
	    cp ${dirb}/str_relax.out .
	else
	    dire=$(echo $iploc | awk '{n=int($1)+1; printf("%02i",n)}')
#echo $dire
	    frac=$(echo $iploc | awk '{print $1-int($1);}')
#echo $frac
	    strpath -f=${frac} -s1=${dirb}/str_relax.out -s2=${dire}/str_relax.out > str_relax.out
	    echo	strpath -f=${frac} -s1=${dirb}/str_relax.out -s2=${dire}/str_relax.out 
	fi
    fi

fi