#!/usr/local/bin/bash
#CATEGORY: Output examination and debugging
#SYNOPSIS: Plot convergence in JDFTx output files using gnuplot

if [ "$1" == "-h" ] || [ "$1" == "--help" ]; then
	echo '
	Plot convergence in JDFTx output files using gnuplot. Usage:
	
		plotConvergence <file1> <file2> ...
	
	If multiple output files are specified, they will be plotted alongside.
	
	The behavior is controlled by several environment variables:
		nLines      Plot the final $nLines iterations (default: 1000)
		skipLines   Ignore the first $skipLines iterations (default: 0)
		iterType    Select iteration type: Elec (default), Ionic, Fluid etc.
		logTarget   Plot energy difference from $logTarget on a log axis.
		xAxis       "iter" (default) or "time" to plot versus iteration count or run time
		term        GNUPLOT terminal with options (default: GNUPLOT default with X, "dumb" without)
		out         GNUPLOT output file; match this with term (default None)
	'
	exit 0
fi

if [[ ( -z $xAxis ) || ( "$xAxis" == "iter" )  || ( "$xAxis" == "Iter" ) ]]; then
	xCol=0
	xLabel="Iteration number"
elif [[ ( "$xAxis" == "time" ) || ( "$xAxis" == "Time" ) ]]; then
	xCol=1
	xLabel="Time [s]"
else
	echo "Unknown option: \$xAxis='$xAxis', must be 'iter' (default) or 'time'"
	exit 1
fi

if [ -z $logTarget ]; then
	using="${xCol}:2"
	yformat='"%.6f"'
else
	logCommand="Final=$logTarget; set logscale y"
	using="${xCol}:(\$2-Final)"
	yformat='"10^{%L}"'
fi

if [ -z $skipLines ]; then
        skipLines=0
fi

if [ -z $nLines ]; then
	nLines=1000
fi

if [ -z $iterType ]; then
	iterType="Elec"
fi

if [ -z "$term" ]; then
	#Check if GNUPLOT can access graphics (works without DISPLAY on MacOS)
	if [[ ( -z $DISPLAY ) &&  ( "$(uname -s)" != "Darwin" ) ]]; then
		termCommand="set term dumb"
	else
		termCommand=""  #use default terminal
	fi
else
	termCommand="set term $term" #use specified terminal (with options)
fi

if [ -z "$out" ]; then
	outCommand=""
else
	outCommand="set output \"$out\""
fi

for outFile in $@; do
	tmpFile="/tmp/${outFile////_._}.dat"
	
	awk "BEGIN {t=0}
	NF>2 && \$(NF-1)==\"t[s]:\" {t=\$NF}
	\$1==\"${iterType}Minimize:\" && \$2==\"Iter:\" { print t,\$5 }
	\"${iterType}\"==\"Elec\" && \$1==\"SCF:\" && \$2==\"Cycle:\" { print t,\$5 }
	" $outFile | tail -n +$skipLines | tail -n -$nLines > $tmpFile
	
	if  [ -z "$plotCommand" ]; then
		plotCommand="plot"
	else
		plotCommand="$plotCommand,"
	fi
	plotCommand="$plotCommand \"$tmpFile\" u $using w l title \"$outFile\""
done

#Plot output
gnuplot -persist << EOF
$outCommand
$termCommand
set format y $yformat
set ylabel "Energy"
set xlabel "$xLabel"
$logCommand
$plotCommand
pause mouse
EOF

