#!/usr/local/bin/bash
#CATEGORY: Input creation
#SYNOPSIS: Generate pulay correction file for a pseudopotential

if [ "$1" == "-h" ] || [ "$1" == "--help" ]; then
	echo '
	Generate pulay correction file for a pseudopotential. Usage:
	
		calcPulay <pspFilename> [<EcutStart>:<dEcut>:<EcutStop>] [<outFile>]
	
	where
	
		<pspFilename> specifies the pseudopotential filename (cannot be a wildcard)
	
		<EcutStart>:<dEcut>:<EcutStop> specifies Pulay corrections should be
			computed from <EcutStart> to <EcutStop> in increments of <dEcut>.
			Defaults to "5:1:50".
	
		<outFile> is the path to the computed pulay correction file. Defaults
			to the pseudopotential filename with the extension replaced by .pulay.
	
	The details of the computation can be controlled by environment variables:
	
		boxSize: Periodic image separation of the atoms in bohrs (Default: 20)
	
		exCorr: Exchange-correlation to use (Default: gga-PBE)
	
		JDFTx: Path to executable jdftx (or jdftx_gpu) (Default: jdftx)
			This is required if jdftx is not in $PATH, to use jdftx_gpu, or
			to use a job management system eg. JDFTx="srun jdftx"
	
		scratchDir: Temporary directory to store wavefunctions and output files.
			Must be writeable. Defaults to current directory.
	'
	exit 0
fi

#---------- Initialize environment ----------

if [ -z "$boxSize" ]; then boxSize="20"; fi
if [ -z "$exCorr" ]; then exCorr="gga-PBE"; fi
if [ -z "$JDFTx" ]; then JDFTx="jdftx"; fi
if [ -z "$scratchDir" ]; then scratchDir="$PWD"; fi

workDir="$scratchDir/tmp-calcPulay.$$"
mkdir -p "$workDir"

#--------- Check inputs ----------------

pspFilename="$1"
EcutRange="$2"
outFile="$3"

if [ -z "$EcutRange" ]; then
	EcutRange="5:1:50"
fi

if [ -z "$outFile" ]; then
	outFile="${pspFilename%.*}.pulay"
fi

speciesName="$(basename $pspFilename)" #remove path if any
speciesNameLength=$(expr match "$speciesName" '[a-zA-Z0-9]*') #alphanumeric characters only
speciesName=${speciesName:0:$speciesNameLength}
echo $speciesName

Ecuts=( $(echo "$EcutRange" | awk -F ':' '
	{	EcutStart = $1<0 ? 5 : $1;
		dEcut = $2<=0 ? 1 : $2;
		EcutStop = $3<EcutStart ? EcutStart : $3;
		for(Ecut=EcutStart; Ecut<=EcutStop; Ecut+=dEcut) { printf("%g\n", Ecut); }
	}' ) )
echo "Computing pulay corrections for species '$speciesName' for Ecut in (${Ecuts[*]}) to be written to '$outFile'"
rm "$outFile" #Make sure pulay file doesn't exist while running calculations with it

#---------- JDFTx evaluator --------------

function runJDFTx
{
	runJDFTxCmdLine="$2" #Options for JDFTx
	runJDFTxOutFile="$1" #Output file
	
	aHlf="$(echo "scale=15; $boxSize/sqrt(2)" | bc)"
	if [ -z "$prevEcut" ]; then
		wfnsParams="lcao"
	else
		wfnsParams="read $workDir/state.wfns 0 $prevEcut"
	fi

	#Run JDFTx with input from a pipe
	$JDFTx $runJDFTxCmdLine <<EOF | tee $runJDFTxOutFile
		lattice 0 1 1  1 0 1  1 1 0
		latt-scale $aHlf $aHlf $aHlf # FCC lattice with image speration = $boxSize bohrs
		ion-species $pspFilename
		ion $speciesName  0 0 0  0
		elec-ex-corr $exCorr
		elec-cutoff $curEcut
		$fftboxCmd
		electronic-SCF energyDiffThreshold 1e-9
		elec-smearing Fermi 0.01
		coulomb-interaction Spherical
		wavefunction $wfnsParams
		dump End State
		dump-name $workDir/state.\$VAR
EOF
	if [ $? -ne 0 ]; then
		echo "Error running JDFTx (please see output above for details)"
		exit 1
	fi
}

#incrementEcut <Ecut> +1 echoes <Ecut> incremented by a suitable dEcut, and
#incrementEcut <Ecut> -1 echoes <Ecut> decremented by it.
function incrementEcut
{
	echo "$1 $2 $boxSize" | awk '{ print (($1)^(3./2) + ($2)*($1)*1./$3) ^ (2./3) }'
}

#Test inputs by performing a dry run:
runJDFTx "$workDir/dryRun.out" "--dry-run"

#Output header:
echo "#Pulay corrections for pseudopotential specification '$pspFilename'
#(Generated by calcPulay from the JDFTx scripts collection)
#Number of Ecuts:
${#Ecuts[@]}
#Ecut  dE/dnbasis (in atomic units)" > "$workDir/pulay"

for Ecut in "${Ecuts[@]}"; do
	#Calculation at slightly larger Ecut
	prevEcut="$curEcut"
	curEcut=$(incrementEcut $Ecut +1)
	fftboxCmd="" #Run with default fftbox
	runJDFTx "$workDir/plus.out"

	#Calculation at slightly smaller Ecut:
	prevEcut="$curEcut"
	curEcut=$(incrementEcut $Ecut -1)
	fftboxCmd="$(awk '/Chosen fftbox/ { print "fftbox",$7,$8,$9 }' $workDir/plus.out)" #Run with same fftbox as the plus calculation
	runJDFTx "$workDir/minus.out"
	
	#Compute pulay correction
	Eminus=$(awk '$1=="IonicMinimize:" && $2=="Iter:" { Energy=$5 } END {printf("%.15f",Energy)}' $workDir/minus.out)
	Eplus=$( awk '$1=="IonicMinimize:" && $2=="Iter:" { Energy=$5 } END {printf("%.15f",Energy)}' $workDir/plus.out)
	NbasisMinus=$(awk '/average nbasis/ { print $4 }' $workDir/minus.out)
	NbasisPlus=$( awk '/average nbasis/ { print $4 }' $workDir/plus.out)
	Vol=$(echo "scale=16; sqrt(0.5)*($boxSize^3)" | bc) #FCC unit cell volume
	dEdnbasis=$(echo "scale=16; $Vol * (($Eplus) - ($Eminus)) / ($NbasisPlus - $NbasisMinus)" | bc)
	echo "$Ecut     $dEdnbasis" >> "$workDir/pulay"
done
mv "$workDir/pulay" "$outFile"
rm -r "$workDir"

echo
echo "Generated '$outFile' with contents:"
echo
cat $outFile
