#!/usr/local/bin/bash
#CATEGORY: Visualization and post-processing
#SYNOPSIS: Create a VASP CHGCAR file from a JDFTx output file

function printUsageAndExit()
{
	echo '
	Create a VASP CHGCAR file from a JDFTx output file. Usage:
	
		createVASP <jdftxOutFile> <vaspFile> [<nFilename>] [<Ninterp>=1]
	
	Save the final structure and electron density from JDFTx output
	file <jdftxOutFile> to a VASP CHGCAR format file in <vaspFile>.
	The electron density is read from the appropriate filename
	determined by parsing the output; this can be overriden by
	explicitly specifying <nFilename>.
	
	Optional <Ninterp> specifies a grid multiplier for Fourier interpolation.
	
	Note: this tool requires awk and octave to be available in path.
	'
	exit $1
}

if [[ "$1" == "-h" || "$1" == "--help" ]]; then printUsageAndExit 0; fi
if [ -z "$2" ]; then printUsageAndExit 1; fi

jdftxOutFile="$1"
vaspFile="$2"
nFilename="$3"
Ninterp="$4"
if [ -z "$Ninterp" ]; then
	Ninterp=1
fi

command -v awk >/dev/null 2>&1 || { echo "createXSF requires awk, which was not found in path.  Aborting." >&2; exit 1; }
command -v octave >/dev/null 2>&1 || { echo "createXSF requires octave, which was not found in path.  Aborting." >&2; exit 1; }

if [ ! -f "$jdftxOutFile" ]; then
	echo "Could not open '$jdftxOutFile' for reading"
	exit 1
fi

tmpDir="createVASP.tmp"
mkdir -p "$tmpDir"
if [ ! -d "$tmpDir" ]; then
	echo "Could not create temporary working directory '$tmpDir'"
	exit 1
fi

#Extract relevant information from the output file (and save in separate files in tmpDir)
awk "
	#Header mode
	BEGIN { mode=1; nAtoms=0; nSteps=0; Sprinted=0; }
	mode==1 && \$1==\"coords-type\" { posCoords = (\$2==\"lattice\" || \$2==\"Lattice\" ? 0 : 1); }
	mode==1 && \$1==\"dump-name\" { fp=\"$tmpDir/dumpName\"; print \$2  > fp; close(fp) }
	mode==1 && ( \$1==\"ion\" || \$1==\"ion-vel\" ) {
		fpIonNames=\"$tmpDir/ionNames\"; print \$2 > fpIonNames;
		fpIonPos=\"$tmpDir/ionpos.0\"; print \$3, \$4, \$5 > fpIonPos;
		nAtoms++; }
	#Initial lattice vectors
	mode==1 && \$0~/Initializing the Grid/ { mode=2; close(fpIonNames); close(fpIonPos);
		refLine=NR; Sprinted=0; Ssuffix=\"grid\"; fp=\"$tmpDir/lattice.0\" }
	mode==2 && NR>refLine+1 && NR<=refLine+4 { print \$2, \$3, \$4 > fp; }
	mode==2 && NR==refLine+4 { close(fp) }
	mode==2 && !Sprinted && \$0~/Chosen fftbox/ { fp=(\"$tmpDir/S\" Ssuffix); print \$(NF-3), \$(NF-2), \$(NF-1) > fp; close(fp); Sprinted=1 }
	#Ionic positions
	/# Ionic positions in / { mode=3; refLine=NR; nSteps++; fp=(\"$tmpDir/ionpos.\" nSteps) }
	mode==3 && NR>refLine && NR<=refLine+nAtoms { print \$3, \$4, \$5 > fp }
	mode==3 && NR==refLine+nAtoms { close(fp) }
	#Lattice vectors
	mode==3 && \$0~/R =/ { mode=4; refLine=NR; fp=(\"$tmpDir/lattice.\" nSteps) }
	mode==4 && NR>refLine && NR<=refLine+3 { print \$2, \$3, \$4 > fp }
	mode==4 && NR==refLine+3 { close(fp) }
	END { print nSteps, nAtoms, posCoords > \"$tmpDir/dimensions\";	}
" $jdftxOutFile

#Put together into the target VASP CHGCAR file
cat > "$tmpDir/process.m" <<EOF
	#Read dimensions
	dims = load("-ascii", "$tmpDir/dimensions");
	nSteps = dims(1);
	nAtoms = dims(2);
	posCoords = dims(3);
	Sgrid = load("-ascii", "$tmpDir/Sgrid"); #charge-density grid
	fp = fopen("$tmpDir/dumpName", "r");
	dumpName = fscanf(fp, "%s");
	fclose(fp);
	
	#Read charge density:
	fname = "$nFilename";
	if length(fname)==0
		if index(dumpName,"\$STAMP")
			printf("Cannot use variable names when \$STAMP is used in dump-name; use full file names instead\n");
			exit(1);
		endif
		fname = strrep(dumpName, "\$VAR", "n");
	endif
	fp = fopen(fname, "r");
	n = fread(fp, "double");
	fclose(fp);
	#--- Handle endianness, if necessary:
	[comp,maxSize,endian] = computer();
	if strcmp(endian,"B")
		n = swapbytes(n); #files are always little-endian, so swap bytes in memory to make big-endian
	endif
	#--- Check length:
	if length(n) != prod(Sgrid)
		printf("Wrong length of '%s': should be exactly %d bytes\n", fname, 8*prod(Sgrid));
		exit(1);
	endif
	#--- C -> Fortran data order:
	n = reshape(permute(reshape(n,fliplr(Sgrid)), [3,2,1]), [],1);
	
	#Interpolate charge density if necessary:
	Ninterp = ceil($Ninterp);
	if Ninterp > 1
		Sshift = floor(Sgrid/2)
		SgridNew = Sgrid * Ninterp;
		nTildeNew = zeros(SgridNew);
		nTildeNew(1:Sgrid(1),1:Sgrid(2),1:Sgrid(3)) = circshift(fftn(reshape(n, Sgrid)), Sshift);
		n = reshape(real(ifftn(circshift(nTildeNew, -Sshift))), [],1) * (Ninterp^3);
		Sgrid = SgridNew;
	endif
	
	#Read atom names
	fp = fopen("$tmpDir/ionNames", "r");
	ionNamesAll = textscan(fp, "%s"){1};
	fclose(fp);
	#--- convert to list of unique names with counts:
	ionNames = {};
	ionCounts = [];
	iIonStart = 1;
	while(iIonStart <= length(ionNamesAll))
		ionNameCur = ionNamesAll{iIonStart};
		iIonEnd = iIonStart;
		while((iIonEnd<length(ionNamesAll)) && strcmp(ionNameCur,ionNamesAll{iIonEnd+1}))
			iIonEnd += 1;
		endwhile
		ionNames{length(ionNames)+1} = ionNameCur;
		ionCounts = [ ionCounts; (iIonEnd-iIonStart+1) ];
		iIonStart = iIonEnd + 1;
	endwhile
	
	#Read final ionpos and lattice:
	ionpos = load("-ascii", sprintf("$tmpDir/ionpos.%d",nSteps));
	for iStep = nSteps:-1:0
		latFile = sprintf("$tmpDir/lattice.%d",iStep);
		if exist(latFile, "file")
			lattice = load("-ascii", latFile);
			break;
		endif
	endfor
	#--- ionpos to lattice coordinates:
	if posCoords==1
		ionpos = ionpos * inv(lattice)'; #convert Cartesian to Lattice
	endif
	#--- ionpos to [0,1] range:
	ionpos -= floor(ionpos);
	#--- multiply electron density by volume:
	n *= abs(det(lattice));
	#--- lattice to Angstroms:
	Angstrom = 1/0.5291772;
	lattice = lattice/Angstrom;

	#Write VASP CHGCAR file:
	fp = fopen("$vaspFile", "w");
	#--- Atom names:
	for iIon=1:length(ionNames)
		fprintf(fp, "%s ", ionNames{iIon});
	endfor
	fprintf(fp, "\n");
	#--- Lattice vectors:
	fprintf(fp, "1.000\n"); #Scale factor
	fprintf(fp, "%12.6f%12.6f%12.6f\n", lattice);
	#--- Atom counts:
	fprintf(fp, "%4d", ionCounts);
	fprintf(fp, "\n");
	#--- Ion positions:
	fprintf(fp, "Direct\n");
	fprintf(fp, "%10.6f%10.6f%10.6f\n", ionpos');
	fprintf(fp, "\n");
	#--- Charge density:
	fprintf(fp, "%5d%5d%5d\n", Sgrid);
	fprintf(fp, "%18.10e%18.10e%18.10e%18.10e%18.10e\n", n);
	fprintf(fp, "\n");
	fclose(fp);
	exit(0);
EOF
octave -qf "$tmpDir/process.m"
rm -rf "$tmpDir"
